Server Side CSM Observability Changes

Author: zasweq

stats/opentelemetry/opentelemetry.go

@@ -157,6 +157,8 @@ func DialOption(o Options) grpc.DialOption {
 	return joinDialOptions(grpc.WithChainUnaryInterceptor(csh.unaryInterceptor), grpc.WithChainStreamInterceptor(csh.streamInterceptor), grpc.WithStatsHandler(csh))
 }
 
+var joinServerOptions = internal.JoinServerOptions.(func(...grpc.ServerOption) grpc.ServerOption)
+
 // ServerOption returns a server option which enables OpenTelemetry
 // instrumentation code for a grpc.Server.
 //
@@ -172,7 +174,7 @@ func DialOption(o Options) grpc.DialOption {
 func ServerOption(o Options) grpc.ServerOption {
 	ssh := &serverStatsHandler{o: o}
 	ssh.initializeMetrics()
-	return grpc.StatsHandler(ssh)
+	return joinServerOptions(grpc.ChainUnaryInterceptor(ssh.unaryInterceptor), grpc.ChainStreamInterceptor(ssh.streamInterceptor), grpc.StatsHandler(ssh))
 }
 
 // callInfo is information pertaining to the lifespan of the RPC client side.

stats/opentelemetry/server_metrics.go

@@ -23,6 +23,7 @@ import (
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
 	"google.golang.org/grpc/status"
 
@@ -55,6 +56,118 @@ func (ssh *serverStatsHandler) initializeMetrics() {
 	ssh.serverMetrics.callDuration = createFloat64Histogram(setOfMetrics, "grpc.server.call.duration", meter, metric.WithUnit("s"), metric.WithDescription("End-to-end time taken to complete a call from server transport's perspective."), metric.WithExplicitBucketBoundaries(DefaultLatencyBounds...))
 }
 
+// attachLabelsTransport stream intercepts SetHeader and SendHeader calls of the
+// underlying ServerTransportStream to attach metadataExchangeLabels.
+type attachLabelsTransportStream struct {
+	grpc.ServerTransportStream
+
+	attachedLabels         atomic.Bool
+	metadataExchangeLabels metadata.MD
+}
+
+func (alts *attachLabelsTransportStream) SetHeader(md metadata.MD) error {
+	if !alts.attachedLabels.Swap(true) {
+		val := alts.metadataExchangeLabels.Get("x-envoy-peer-metadata")
+		md.Append("x-envoy-peer-metadata", val...)
+	}
+	return alts.ServerTransportStream.SetHeader(md)
+}
+
+func (alts *attachLabelsTransportStream) SendHeader(md metadata.MD) error {
+	if !alts.attachedLabels.Swap(true) {
+		val := alts.metadataExchangeLabels.Get("x-envoy-peer-metadata")
+		md.Append("x-envoy-peer-metadata", val...)
+	}
+
+	return alts.ServerTransportStream.SendHeader(md)
+}
+
+func (ssh *serverStatsHandler) unaryInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+	var metadataExchangeLabels metadata.MD
+	if ssh.o.MetricsOptions.pluginOption != nil {
+		metadataExchangeLabels = ssh.o.MetricsOptions.pluginOption.GetMetadata()
+	}
+
+	sts := grpc.ServerTransportStreamFromContext(ctx)
+
+	alts := &attachLabelsTransportStream{
+		ServerTransportStream:  sts,
+		metadataExchangeLabels: metadataExchangeLabels,
+	}
+	ctx = grpc.NewContextWithServerTransportStream(ctx, alts)
+
+	any, err := handler(ctx, req)
+	if err != nil { // error returned, so trailers only
+		if !alts.attachedLabels.Swap(true) {
+			alts.SetTrailer(alts.metadataExchangeLabels)
+		}
+		logger.Infof("RPC failed with error: %v", err)
+	} else { // headers will be written; a message was sent
+		if !alts.attachedLabels.Swap(true) {
+			alts.SetHeader(alts.metadataExchangeLabels)
+		}
+	}
+
+	return any, err
+}
+
+// attachLabelsStream embeds a grpc.ServerStream, and intercepts the
+// SetHeader/SendHeader/SendMsg/SendTrailer call to attach metadata exchange
+// labels.
+type attachLabelsStream struct {
+	grpc.ServerStream
+
+	attachedLabels         atomic.Bool
+	metadataExchangeLabels metadata.MD
+}
+
+func (als *attachLabelsStream) SetHeader(md metadata.MD) error {
+	if !als.attachedLabels.Swap(true) {
+		val := als.metadataExchangeLabels.Get("x-envoy-peer-metadata")
+		md.Append("x-envoy-peer-metadata", val...)
+	}
+
+	return als.ServerStream.SetHeader(md)
+}
+
+func (als *attachLabelsStream) SendHeader(md metadata.MD) error {
+	if !als.attachedLabels.Swap(true) {
+		val := als.metadataExchangeLabels.Get("x-envoy-peer-metadata")
+		md.Append("x-envoy-peer-metadata", val...)
+	}
+
+	return als.ServerStream.SendHeader(md)
+}
+
+func (als *attachLabelsStream) SendMsg(m any) error {
+	if !als.attachedLabels.Swap(true) {
+		als.ServerStream.SetHeader(als.metadataExchangeLabels)
+	}
+	return als.ServerStream.SendMsg(m)
+}
+
+func (ssh *serverStatsHandler) streamInterceptor(srv any, ss grpc.ServerStream, ssi *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+	var metadataExchangeLabels metadata.MD
+	if ssh.o.MetricsOptions.pluginOption != nil {
+		metadataExchangeLabels = ssh.o.MetricsOptions.pluginOption.GetMetadata()
+	}
+	als := &attachLabelsStream{
+		ServerStream:           ss,
+		metadataExchangeLabels: metadataExchangeLabels,
+	}
+	err := handler(srv, als)
+	if err != nil {
+		logger.Infof("RPC failed with error: %v", err)
+	}
+
+	// Add metadata exchange labels to trailers if never sent in headers,
+	// irrespective of whether or not RPC failed.
+	if !als.attachedLabels.Load() {
+		als.SetTrailer(als.metadataExchangeLabels)
+	}
+	return err
+}
+
 // TagConn exists to satisfy stats.Handler.
 func (ssh *serverStatsHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
 	return ctx
@@ -105,6 +218,10 @@ func (ssh *serverStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats)
 func (ssh *serverStatsHandler) processRPCData(ctx context.Context, s stats.RPCStats, mi *metricsInfo) {
 	switch st := s.(type) {
 	case *stats.InHeader:
+		if !mi.labelsReceived && ssh.o.MetricsOptions.pluginOption != nil {
+			mi.labels = ssh.o.MetricsOptions.pluginOption.GetLabels(st.Header)
+			mi.labelsReceived = true
+		}
 		ssh.serverMetrics.callStarted.Add(ctx, 1, metric.WithAttributes(attribute.String("grpc.method", mi.method)))
 	case *stats.OutPayload:
 		atomic.AddInt64(&mi.sentCompressedBytes, int64(st.CompressedLength))
@@ -123,8 +240,15 @@ func (ssh *serverStatsHandler) processRPCEnd(ctx context.Context, mi *metricsInf
 		s, _ := status.FromError(e.Error)
 		st = canonicalString(s.Code())
 	}
-	serverAttributeOption := metric.WithAttributes(attribute.String("grpc.method", mi.method), attribute.String("grpc.status", st))
+	attributes := []attribute.KeyValue{
+		attribute.String("grpc.method", mi.method),
+		attribute.String("grpc.status", st),
+	}
+	for k, v := range mi.labels {
+		attributes = append(attributes, attribute.String(k, v))
+	}
 
+	serverAttributeOption := metric.WithAttributes(attributes...)
 	ssh.serverMetrics.callDuration.Record(ctx, latency, serverAttributeOption)
 	ssh.serverMetrics.callSentTotalCompressedMessageSize.Record(ctx, atomic.LoadInt64(&mi.sentCompressedBytes), serverAttributeOption)
 	ssh.serverMetrics.callRcvdTotalCompressedMessageSize.Record(ctx, atomic.LoadInt64(&mi.recvCompressedBytes), serverAttributeOption)