feat(otelgrpc): support adding custom metric attributes to gRPCContext

open-telemetry · Oct 29, 2024 · 370d0b2 · 370d0b2
1 parent c0f6f09
commit 370d0b2
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 16 deletions.
diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/internal/test/test_utils.go b/instrumentation/google.golang.org/grpc/otelgrpc/internal/test/test_utils.go
@@ -36,11 +36,14 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/grpclog"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 
-	testpb "google.golang.org/grpc/interop/grpc_testing"
+	"go.opentelemetry.io/otel/attribute"
+
+	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 )
 
 var (
@@ -241,6 +244,7 @@ func NewTestServer() testpb.TestServiceServer {
 }
 
 func (s *testServer) EmptyCall(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
+	otelgrpc.AddMetricAttributes(ctx, attribute.String("custom_test_metric", "OK"))
 	return new(testpb.Empty), nil
 }
 
@@ -261,6 +265,7 @@ func serverNewPayload(t testpb.PayloadType, size int32) (*testpb.Payload, error)
 }
 
 func (s *testServer) UnaryCall(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+	otelgrpc.AddMetricAttributes(ctx, attribute.String("custom_test_metric", "OK"))
 	st := in.GetResponseStatus()
 	if md, ok := metadata.FromIncomingContext(ctx); ok {
 		if initialMetadata, ok := md[initialMetadataKey]; ok {

diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go b/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go
@@ -30,6 +30,17 @@ type gRPCContext struct {
 	record           bool
 }
 
+// AddMetricAttributes adds extra metric attributes which be added to recorded metrics.
+// This can be called from within your gRPC handlers.
+func AddMetricAttributes(ctx context.Context, attributes ...attribute.KeyValue) {
+	gctx, ok := ctx.Value(gRPCContextKey{}).(*gRPCContext)
+	if !ok {
+		return
+	}
+
+	gctx.metricAttrs = append(gctx.metricAttrs, attributes...)
+}
+
 type serverHandler struct {
 	*config
 }

diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_stats_handler_test.go b/instrumentation/google.golang.org/grpc/otelgrpc/test/grpc_stats_handler_test.go
@@ -16,6 +16,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
 	"google.golang.org/grpc/status"
 
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
@@ -26,17 +27,15 @@ import (
 	"go.opentelemetry.io/otel/sdk/metric"
 	"go.opentelemetry.io/otel/sdk/metric/metricdata"
 	"go.opentelemetry.io/otel/sdk/metric/metricdata/metricdatatest"
-
 	"go.opentelemetry.io/otel/sdk/trace"
 	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	semconv "go.opentelemetry.io/otel/semconv/v1.17.0"
-
-	testpb "google.golang.org/grpc/interop/grpc_testing"
 )
 
 var (
-	testSpanAttr   = attribute.String("test_span", "OK")
-	testMetricAttr = attribute.String("test_metric", "OK")
+	testSpanAttr         = attribute.String("test_span", "OK")
+	testMetricAttr       = attribute.String("test_metric", "OK")
+	customTestMetricAttr = attribute.String("custom_test_metric", "OK")
 )
 
 func TestStatsHandler(t *testing.T) {
@@ -1115,15 +1114,19 @@ func checkServerMetrics(t *testing.T, reader metric.Reader) {
 								semconv.RPCMethod("EmptyCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+								customTestMetricAttr,
+							),
 						},
 						{
 							Attributes: attribute.NewSet(
 								semconv.RPCGRPCStatusCodeOk,
 								semconv.RPCMethod("UnaryCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+								customTestMetricAttr,
+							),
 						},
 						{
 							Attributes: attribute.NewSet(
@@ -1164,7 +1167,8 @@ func checkServerMetrics(t *testing.T, reader metric.Reader) {
 								semconv.RPCMethod("EmptyCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+							),
 							Bounds:       []float64{0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000},
 							BucketCounts: []uint64{1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
 							Max:          metricdata.NewExtrema(int64(0)),
@@ -1177,7 +1181,8 @@ func checkServerMetrics(t *testing.T, reader metric.Reader) {
 								semconv.RPCMethod("UnaryCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+							),
 							Bounds:       []float64{0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000},
 							BucketCounts: []uint64{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1},
 							Max:          metricdata.NewExtrema(int64(271840)),
@@ -1239,7 +1244,9 @@ func checkServerMetrics(t *testing.T, reader metric.Reader) {
 								semconv.RPCMethod("EmptyCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+								customTestMetricAttr,
+							),
 							Bounds:       []float64{0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000},
 							BucketCounts: []uint64{1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
 							Max:          metricdata.NewExtrema(int64(0)),
@@ -1252,7 +1259,9 @@ func checkServerMetrics(t *testing.T, reader metric.Reader) {
 								semconv.RPCMethod("UnaryCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+								customTestMetricAttr,
+							),
 							Bounds:       []float64{0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000},
 							BucketCounts: []uint64{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1},
 							Max:          metricdata.NewExtrema(int64(314167)),
@@ -1315,7 +1324,9 @@ func checkServerMetrics(t *testing.T, reader metric.Reader) {
 								semconv.RPCMethod("EmptyCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+								customTestMetricAttr,
+							),
 							Bounds:       []float64{0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000},
 							BucketCounts: []uint64{0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
 							Max:          metricdata.NewExtrema(int64(1)),
@@ -1329,7 +1340,9 @@ func checkServerMetrics(t *testing.T, reader metric.Reader) {
 								semconv.RPCMethod("UnaryCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+								customTestMetricAttr,
+							),
 							Bounds:       []float64{0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000},
 							BucketCounts: []uint64{0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
 							Max:          metricdata.NewExtrema(int64(1)),
@@ -1395,7 +1408,9 @@ func checkServerMetrics(t *testing.T, reader metric.Reader) {
 								semconv.RPCMethod("EmptyCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+								customTestMetricAttr,
+							),
 							Bounds:       []float64{0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000},
 							BucketCounts: []uint64{0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
 							Max:          metricdata.NewExtrema(int64(1)),
@@ -1409,7 +1424,9 @@ func checkServerMetrics(t *testing.T, reader metric.Reader) {
 								semconv.RPCMethod("UnaryCall"),
 								semconv.RPCService("grpc.testing.TestService"),
 								semconv.RPCSystemGRPC,
-								testMetricAttr),
+								testMetricAttr,
+								customTestMetricAttr,
+							),
 							Bounds:       []float64{0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000},
 							BucketCounts: []uint64{0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
 							Max:          metricdata.NewExtrema(int64(1)),

diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/test/stats_handler_test.go b/instrumentation/google.golang.org/grpc/otelgrpc/test/stats_handler_test.go
@@ -42,11 +42,15 @@ func TestStatsHandlerHandleRPCServerErrors(t *testing.T) {
 			serviceName := "TestGrpcService"
 			methodName := serviceName + "/" + name
 			fullMethodName := "/" + methodName
+
 			// call the server handler
 			ctx := serverHandler.TagRPC(context.Background(), &stats.RPCTagInfo{
 				FullMethodName: fullMethodName,
 			})
 
+			// add custom metric attribute
+			otelgrpc.AddMetricAttributes(ctx, customTestMetricAttr)
+
 			grpcErr := status.Error(check.grpcCode, check.grpcCode.String())
 			serverHandler.HandleRPC(ctx, &stats.End{
 				Error: grpcErr,
@@ -81,6 +85,7 @@ func assertStatsHandlerServerMetrics(t *testing.T, reader metric.Reader, service
 								otelgrpc.RPCSystemGRPC,
 								otelgrpc.GRPCStatusCodeKey.Int64(int64(code)),
 								testMetricAttr,
+								customTestMetricAttr,
 							),
 						},
 					},
@@ -100,6 +105,7 @@ func assertStatsHandlerServerMetrics(t *testing.T, reader metric.Reader, service
 								otelgrpc.RPCSystemGRPC,
 								otelgrpc.GRPCStatusCodeKey.Int64(int64(code)),
 								testMetricAttr,
+								customTestMetricAttr,
 							),
 						},
 					},
@@ -119,6 +125,7 @@ func assertStatsHandlerServerMetrics(t *testing.T, reader metric.Reader, service
 								otelgrpc.RPCSystemGRPC,
 								otelgrpc.GRPCStatusCodeKey.Int64(int64(code)),
 								testMetricAttr,
+								customTestMetricAttr,
 							),
 						},
 					},