stats/internal: OpenTelemetry tracing GRPCTraceBinPropagator

go.mod

@@ -15,6 +15,7 @@ require (
 	go.opentelemetry.io/otel/metric v1.31.0
 	go.opentelemetry.io/otel/sdk v1.31.0
 	go.opentelemetry.io/otel/sdk/metric v1.31.0
+	go.opentelemetry.io/otel/trace v1.31.0
 	golang.org/x/net v0.30.0
 	golang.org/x/oauth2 v0.23.0
 	golang.org/x/sync v0.8.0
@@ -32,7 +33,6 @@ require (
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
-	go.opentelemetry.io/otel/trace v1.31.0 // indirect
 	golang.org/x/text v0.19.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20241015192408-796eee8c2d53 // indirect
 )

stats/opentelemetry/grpc_trace_bin_propagator.go

@@ -0,0 +1,119 @@
+/*
+ *
+ * Copyright 2024 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package opentelemetry
+
+import (
+	"context"
+
+	otelpropagation "go.opentelemetry.io/otel/propagation"
+	oteltrace "go.opentelemetry.io/otel/trace"
+)
+
+// GRPCTraceBinHeaderKey is the gRPC metadata header key `grpc-trace-bin` used
+// to propagate trace context in binary format.
+const GRPCTraceBinHeaderKey = "grpc-trace-bin"
+
+// GRPCTraceBinPropagator is an OpenTelemetry TextMapPropagator which is used
+// to extract and inject trace context data from and into headers exchanged by
+// gRPC applications. It propagates trace data in binary format using the
+// `grpc-trace-bin` header.
+type GRPCTraceBinPropagator struct{}
+
+// Inject sets OpenTelemetry span context from the Context into the carrier as
+// a `grpc-trace-bin` header if span context is valid.
+//
+// If span context is not valid, it ruturns without setting `grpc-trace-bin`
+// header.
+func (GRPCTraceBinPropagator) Inject(ctx context.Context, carrier otelpropagation.TextMapCarrier) {
+	sc := oteltrace.SpanFromContext(ctx)
+	if !sc.SpanContext().IsValid() {
+		return
+	}
+
+	bd := binary(sc.SpanContext())
+	carrier.Set(GRPCTraceBinHeaderKey, string(bd))
+}
+
+// Extract reads OpenTelemetry span context from the `grpc-trace-bin` header of
+// carrier into the provided context, if present.
+//
+// If a valid span context is retrieved from `grpc-trace-bin`, it returns a new
+// context containing the extracted OpenTelemetry span context marked as
+// remote.
+//
+// If `grpc-trace-bin` header is not present, it returns the context as is.
+func (GRPCTraceBinPropagator) Extract(ctx context.Context, carrier otelpropagation.TextMapCarrier) context.Context {
+	h := carrier.Get(GRPCTraceBinHeaderKey)
+	if h == "" {
+		return ctx
+	}
+
+	sc, ok := fromBinary([]byte(h))
+	if !ok {
+		return ctx
+	}
+	return oteltrace.ContextWithRemoteSpanContext(ctx, sc)
+}
+
+// Fields returns the keys whose values are set with Inject.
+//
+// GRPCTraceBinPropagator always returns a slice containing only
+// `grpc-trace-bin` key because it only sets the `grpc-trace-bin` header for
+// propagating trace context.
+func (GRPCTraceBinPropagator) Fields() []string {
+	return []string{GRPCTraceBinHeaderKey}
+}
+
+// Binary returns the binary format representation of a SpanContext.
+//
+// If sc is the zero value, returns nil.
+func binary(sc oteltrace.SpanContext) []byte {
+	if sc.Equal(oteltrace.SpanContext{}) {
+		return nil
+	}
+	var b [29]byte
+	traceID := oteltrace.TraceID(sc.TraceID())
+	copy(b[2:18], traceID[:])
+	b[18] = 1
+	spanID := oteltrace.SpanID(sc.SpanID())
+	copy(b[19:27], spanID[:])
+	b[27] = 2
+	b[28] = byte(oteltrace.TraceFlags(sc.TraceFlags()))
+	return b[:]
+}
+
+// FromBinary returns the SpanContext represented by b with Remote set to true.
+//
+// It returns with zero value SpanContext and false, if any of the
+// below condition is not satisfied:
+// - Valid header: len(b) = 29
+// - Valid version: b[0] = 0
+// - Valid traceID prefixed with 0: b[1] = 0
+// - Valid spanID prefixed with 1: b[18] = 1
+// - Valid traceFlags prefixed with 2: b[27] = 2
+func fromBinary(b []byte) (oteltrace.SpanContext, bool) {
+	if len(b) != 29 || b[0] != 0 || b[1] != 0 || b[18] != 1 || b[27] != 2 {
+		return oteltrace.SpanContext{}, false
+	}
+
+	return oteltrace.SpanContext{}.WithTraceID(
+		oteltrace.TraceID(b[2:18])).WithSpanID(
+		oteltrace.SpanID(b[19:27])).WithTraceFlags(
+		oteltrace.TraceFlags(b[28])).WithRemote(true), true
+}

stats/opentelemetry/grpc_trace_bin_propagator_test.go

@@ -0,0 +1,226 @@
+/*
+ *
+ * Copyright 2024 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package opentelemetry
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	oteltrace "go.opentelemetry.io/otel/trace"
+	"google.golang.org/grpc/metadata"
+	itracing "google.golang.org/grpc/stats/opentelemetry/internal/tracing"
+)
+
+var validSpanContext = oteltrace.SpanContext{}.WithTraceID(
+	oteltrace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}).WithSpanID(
+	oteltrace.SpanID{17, 18, 19, 20, 21, 22, 23, 24}).WithTraceFlags(
+	oteltrace.TraceFlags(1))
+
+// TestInject verifies that the GRPCTraceBinPropagator correctly injects
+// existing binary trace data or OpenTelemetry span context as `grpc-trace-bin`
+// header in the provided carrier's metadata.
+//
+// It verifies that if a valid span context is injected, same span context can
+// can be retreived from the carrier's metadata.
+//
+// If an invalid span context is injected, it verifies that `grpc-trace-bin`
+// header is not set in the carrier's metadata.
+func (s) TestInject(t *testing.T) {
+	tests := []struct {
+		name     string
+		injectSC oteltrace.SpanContext
+		wantSC   oteltrace.SpanContext
+	}{
+		{
+			name:     "valid OpenTelemetry span context",
+			injectSC: validSpanContext,
+			wantSC:   validSpanContext,
+		},
+		{
+			name:     "invalid OpenTelemetry span context",
+			injectSC: oteltrace.SpanContext{},
+			wantSC:   oteltrace.SpanContext{},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			p := GRPCTraceBinPropagator{}
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			ctx = oteltrace.ContextWithSpanContext(ctx, test.injectSC)
+
+			c := itracing.NewCustomCarrier(&metadata.MD{})
+			p.Inject(ctx, c)
+			gotH := c.Get(GRPCTraceBinHeaderKey)
+			if !test.wantSC.IsValid() {
+				if gotH != "" {
+					t.Fatalf("got non-empty value from CustomCarrier's metadata grpc-trace-bin header, want empty")
+				}
+				return
+			}
+			if gotH == "" {
+				t.Fatalf("got empty value from CustomCarrier's metadata grpc-trace-bin header, want valid span context: %v", test.wantSC)
+			}
+			gotSC, ok := fromBinary([]byte(gotH))
+			if !ok {
+				t.Fatalf("got invalid span context from CustomCarrier's metadata grpc-trace-bin header, want valid span context: %v", test.wantSC)
+			}
+			if test.wantSC.TraceID() != gotSC.TraceID() && test.wantSC.SpanID() != gotSC.SpanID() && test.wantSC.TraceFlags() != gotSC.TraceFlags() {
+				t.Fatalf("got span context = %v, want span contexts %v", gotSC, test.wantSC)
+			}
+		})
+	}
+}
+
+// TestExtract verifies that the GRPCTraceBinPropagator correctly extracts
+// OpenTelemetry span context data from the provided context using carrier.
+//
+// If a valid span context was injected, it verifies same trace span context
+// is extracted from carrier's metadata for `grpc-trace-bin` header key.
+//
+// If invalid span context was injected, it verifies that valid trace span
+// context is not extracted.
+func (s) TestExtract(t *testing.T) {
+	tests := []struct {
+		name   string
+		wantSC oteltrace.SpanContext // expected span context from carrier
+	}{
+		{
+			name:   "valid OpenTelemetry span context",
+			wantSC: validSpanContext.WithRemote(true),
+		},
+		{
+			name:   "invalid OpenTelemetry span context",
+			wantSC: oteltrace.SpanContext{},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			p := GRPCTraceBinPropagator{}
+			bd := binary(test.wantSC)
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			c := itracing.NewCustomCarrier(&metadata.MD{GRPCTraceBinHeaderKey: []string{string(bd)}})
+
+			tCtx := p.Extract(ctx, c)
+			got := oteltrace.SpanContextFromContext(tCtx)
+			if !got.Equal(test.wantSC) {
+				t.Fatalf("got span context: %v, want span context: %v", got, test.wantSC)
+			}
+		})
+	}
+}
+
+// TestBinary verifies that the binary() function correctly serializes a valid
+// OpenTelemetry span context into its binary format representation. If span
+// context is invalid, it verifies that serialization is nil.
+func (s) TestBinary(t *testing.T) {
+	tests := []struct {
+		name string
+		sc   oteltrace.SpanContext
+		want []byte
+	}{
+		{
+			name: "valid context",
+			sc:   validSpanContext,
+			want: binary(validSpanContext),
+		},
+		{
+			name: "zero value context",
+			sc:   oteltrace.SpanContext{},
+			want: nil,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if got := binary(test.sc); !cmp.Equal(got, test.want) {
+				t.Fatalf("binary() = %v, want %v", got, test.want)
+			}
+		})
+	}
+}
+
+// TestFromBinary verifies that the fromBinary() function correctly
+// deserializes a binary format representation of a valid OpenTelemetry span
+// context into its corresponding span context format. If span context's binary
+// representation is invalid, it verifies that deserialization is zero value
+// span context.
+func (s) TestFromBinary(t *testing.T) {
+	tests := []struct {
+		name string
+		b    []byte
+		want oteltrace.SpanContext
+		ok   bool
+	}{
+		{
+			name: "valid",
+			b:    []byte{0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 1, 17, 18, 19, 20, 21, 22, 23, 24, 2, 1},
+			want: validSpanContext.WithRemote(true),
+			ok:   true,
+		},
+		{
+			name: "invalid length",
+			b:    []byte{0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 1, 17, 18, 19, 20, 21, 22, 23, 24, 2},
+			want: oteltrace.SpanContext{},
+			ok:   false,
+		},
+		{
+			name: "invalid version",
+			b:    []byte{1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 1, 17, 18, 19, 20, 21, 22, 23, 24, 2, 1},
+			want: oteltrace.SpanContext{},
+			ok:   false,
+		},
+		{
+			name: "invalid traceID field ID",
+			b:    []byte{0, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 1, 17, 18, 19, 20, 21, 22, 23, 24, 2, 1},
+			want: oteltrace.SpanContext{},
+			ok:   false,
+		},
+		{
+			name: "invalid spanID field ID",
+			b:    []byte{0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 0, 17, 18, 19, 20, 21, 22, 23, 24, 2, 1},
+			want: oteltrace.SpanContext{},
+			ok:   false,
+		},
+		{
+			name: "invalid traceFlags field ID",
+			b:    []byte{0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 1, 17, 18, 19, 20, 21, 22, 23, 24, 1, 1},
+			want: oteltrace.SpanContext{},
+			ok:   false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got, ok := fromBinary(test.b)
+			if ok != test.ok {
+				t.Fatalf("fromBinary() ok = %v, want %v", ok, test.ok)
+				return
+			}
+			if !got.Equal(test.want) {
+				t.Fatalf("fromBinary() got = %v, want %v", got, test.want)
+			}
+		})
+	}
+}