Add Otel Tracing API to grpc-go

stats/handlers.go

@@ -38,6 +38,14 @@ type RPCTagInfo struct {
 	// FailFast indicates if this RPC is failfast.
 	// This field is only valid on client side, it's always false on server side.
 	FailFast bool
+	// NameResolutionDelay indicates whether there was a delay in name
+	// resolution.
+	//
+	// This field is only valid on client side, it's always false on server side.
+	NameResolutionDelay bool
+	// IsTransparentRetry indicates whether the stream is undergoing a
+	// transparent retry.
+	IsTransparentRetry bool
 }
 
 // Handler defines the interface for the related stats handling (e.g., RPCs, connections).

stats/opentelemetry/client_metrics.go

@@ -18,6 +18,7 @@ package opentelemetry
 
 import (
 	"context"
+	"go.opentelemetry.io/otel"
 	"sync/atomic"
 	"time"
 
@@ -68,6 +69,15 @@ func (h *clientStatsHandler) initializeMetrics() {
 	rm.registerMetrics(metrics, meter)
 }
 
+func (h *clientStatsHandler) initializeTracing() {
+	if h.options.TraceOptions.TracerProvider == nil || h.options.TraceOptions.TextMapPropagator == nil {
+		return
+	}
+
+	otel.SetTextMapPropagator(h.options.TraceOptions.TextMapPropagator)
+	otel.SetTracerProvider(h.options.TraceOptions.TracerProvider)
+}
+
 func (h *clientStatsHandler) unaryInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 	ci := &callInfo{
 		target: cc.CanonicalTarget(),

stats/opentelemetry/internal/tracing/custom_map_carrier.go

@@ -0,0 +1,109 @@
+/*
+ *
+ * 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 tracing implements the OpenTelemetry carrier for context propagation
+// in gRPC tracing.
+package tracing
+
+import (
+	"context"
+
+	otelpropagation "go.opentelemetry.io/otel/propagation"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
+)
+
+// GRPCTraceBinHeaderKey is the gRPC metadata header key `grpc-trace-bin` used
+// to propagate trace context in binary format.
+const GRPCTraceBinHeaderKey = "grpc-trace-bin"
+
+// CustomCarrier is a TextMapCarrier that uses gRPC context to store and
+// retrieve any propagated key-value pairs in text format along with binary
+// format for `grpc-trace-bin` header
+type CustomCarrier struct {
+	otelpropagation.TextMapCarrier
+
+	ctx context.Context
+}
+
+// NewCustomCarrier creates a new CustomMapCarrier with
+// the given context.
+func NewCustomCarrier(ctx context.Context) *CustomCarrier {
+	return &CustomCarrier{
+		ctx: ctx,
+	}
+}
+
+// Get returns the string value associated with the passed key from the gRPC
+// context. It returns an empty string if the key is not present in the
+// context.
+func (c *CustomCarrier) Get(key string) string {
+	md, ok := metadata.FromIncomingContext(c.ctx)
+	if !ok {
+		return ""
+	}
+	values := md.Get(key)
+	if len(values) == 0 {
+		return ""
+	}
+	return values[0]
+}
+
+// Set stores the key-value pair in string format in the gRPC context.
+// If the key already exists, its value will be overwritten.
+func (c *CustomCarrier) Set(key, value string) {
+	md, ok := metadata.FromOutgoingContext(c.ctx)
+	if !ok {
+		md = metadata.MD{}
+	}
+	md.Set(key, value)
+	c.ctx = metadata.NewOutgoingContext(c.ctx, md)
+}
+
+// GetBinary returns the binary value from the gRPC context in the incoming RPC,
+// associated with the header `grpc-trace-bin`.
+func (c CustomCarrier) GetBinary() []byte {
+	values := stats.Trace(c.ctx)
+	if len(values) == 0 {
+		return nil
+	}
+
+	return values
+}
+
+// SetBinary sets the binary value to the gRPC context, which will be sent in
+// the outgoing RPC with the header `grpc-trace-bin`.
+func (c *CustomCarrier) SetBinary(value []byte) {
+	c.ctx = stats.SetTrace(c.ctx, value)
+}
+
+// Keys returns the keys stored in the gRPC context for the outgoing RPC.
+func (c *CustomCarrier) Keys() []string {
+	md, _ := metadata.FromOutgoingContext(c.ctx)
+	keys := make([]string, 0, len(md))
+	for k := range md {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
+// Context returns the underlying *context.Context associated with the
+// CustomCarrier.
+func (c *CustomCarrier) Context() context.Context {
+	return c.ctx
+}

stats/opentelemetry/internal/tracing/custom_map_carrier_test.go

@@ -0,0 +1,183 @@
+/*
+ *
+ * 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
+ *
+ *     htestp://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 tracing
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestGet(t *testing.T) {
+	tests := []struct {
+		name string
+		md   metadata.MD
+		key  string
+		want string
+	}{
+		{
+			name: "existing key",
+			md:   metadata.Pairs("key1", "value1"),
+			key:  "key1",
+			want: "value1",
+		},
+		{
+			name: "non-existing key",
+			md:   metadata.Pairs("key1", "value1"),
+			key:  "key2",
+			want: "",
+		},
+		{
+			name: "empty key",
+			md:   metadata.MD{},
+			key:  "key1",
+			want: "",
+		},
+	}
+
+	for _, test := range tests {
+		ctx, cancel := context.WithCancel(context.Background())
+		t.Run(test.name, func(t *testing.T) {
+			c := NewCustomCarrier(metadata.NewIncomingContext(ctx, test.md))
+			got := c.Get(test.key)
+			if got != test.want {
+				t.Fatalf("got %s, want %s", got, test.want)
+			}
+			cancel()
+		})
+	}
+}
+
+func (s) TestSet(t *testing.T) {
+	tests := []struct {
+		name      string
+		initialMD metadata.MD // Metadata to initialize the context with
+		setKey    string      // Key to set using c.Set()
+		setValue  string      // Value to set using c.Set()
+		wantKeys  []string    // Expected keys returned by c.Keys()
+	}{
+		{
+			name:      "set new key",
+			initialMD: metadata.MD{},
+			setKey:    "key1",
+			setValue:  "value1",
+			wantKeys:  []string{"key1"},
+		},
+		{
+			name:      "override existing key",
+			initialMD: metadata.MD{"key1": []string{"oldvalue"}},
+			setKey:    "key1",
+			setValue:  "newvalue",
+			wantKeys:  []string{"key1"},
+		},
+		{
+			name:      "set key with existing unrelated key",
+			initialMD: metadata.MD{"key2": []string{"value2"}},
+			setKey:    "key1",
+			setValue:  "value1",
+			wantKeys:  []string{"key2", "key1"}, // Order matesters here!
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			c := NewCustomCarrier(metadata.NewOutgoingContext(ctx, test.initialMD))
+
+			c.Set(test.setKey, test.setValue)
+
+			gotKeys := c.Keys()
+			equalKeys := cmp.Equal(test.wantKeys, gotKeys, cmpopts.SortSlices(func(a, b string) bool {
+				return a < b
+			}))
+			if !equalKeys {
+				t.Fatalf("got keys %v, want %v", gotKeys, test.wantKeys)
+			}
+			gotMD, _ := metadata.FromOutgoingContext(c.Context())
+			if gotMD.Get(test.setKey)[0] != test.setValue {
+				t.Fatalf("got value %s, want %s, for key %s", gotMD.Get(test.setKey)[0], test.setValue, test.setKey)
+			}
+			cancel()
+		})
+	}
+}
+
+func (s) TestGetBinary(t *testing.T) {
+	t.Run("get grpc-trace-bin header", func(t *testing.T) {
+		want := []byte{0x01, 0x02, 0x03}
+		ctx, cancel := context.WithCancel(context.Background())
+		c := NewCustomCarrier(stats.SetIncomingTrace(ctx, want))
+		got := c.GetBinary()
+		if got == nil {
+			t.Fatalf("got nil, want %v", got)
+		}
+		if string(got) != string(want) {
+			t.Fatalf("got %s, want %s", got, want)
+		}
+		cancel()
+	})
+
+	t.Run("get non grpc-trace-bin header", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		c := NewCustomCarrier(metadata.NewIncomingContext(ctx, metadata.Pairs("non-trace-bin", "\x01\x02\x03")))
+		got := c.GetBinary()
+		if got != nil {
+			t.Fatalf("got %v, want nil", got)
+		}
+		cancel()
+	})
+}
+
+func (s) TestSetBinary(t *testing.T) {
+	t.Run("set grpc-trace-bin header", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		want := []byte{0x01, 0x02, 0x03}
+		c := NewCustomCarrier(stats.SetIncomingTrace(ctx, want))
+		c.SetBinary(want)
+		got := stats.OutgoingTrace(c.Context())
+		if !reflect.DeepEqual(got, want) {
+			t.Fatalf("got %v, want %v", got, want)
+		}
+		cancel()
+	})
+
+	t.Run("set non grpc-trace-bin header", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		c := NewCustomCarrier(metadata.NewOutgoingContext(ctx, metadata.MD{"non-trace-bin": []string{"value"}}))
+		got := stats.OutgoingTrace(c.Context())
+		if got != nil {
+			t.Fatalf("got %v, want nil", got)
+		}
+		cancel()
+	})
+}