feat(v2): add callctx telemetry helpers (#472)

quartzmo · web-flow · commit fa319ffc3093 · 2026-03-05T14:49:55.000-07:00
Replace gRPC metadata Context management with new telemetry Context helpers.
diff --git a/v2/callctx/callctx.go b/v2/callctx/callctx.go
@@ -98,3 +98,30 @@ func cloneHeaders(h map[string][]string) map[string][]string {
 	}
 	return c
 }
+
+// telemetryKey is a private type used to store/retrieve telemetry context values.
+type telemetryKey string
+
+// WithTelemetryContext injects telemetry attribute values (like resource name
+// or client version) into the context. In accordance with standard Go context
+// guidelines, this should only be used for data that transits processes and APIs,
+// and not for passing optional parameters to functions. keyvals should have a
+// corresponding value for every key provided. If there is an odd number of keyvals
+// this method will panic.
+func WithTelemetryContext(ctx context.Context, keyvals ...string) context.Context {
+	if len(keyvals)%2 != 0 {
+		panic(fmt.Sprintf("callctx: an even number of key value pairs must be provided, got %d", len(keyvals)))
+	}
+
+	for i := 0; i < len(keyvals); i = i + 2 {
+		ctx = context.WithValue(ctx, telemetryKey(keyvals[i]), keyvals[i+1])
+	}
+	return ctx
+}
+
+// TelemetryFromContext extracts a telemetry attribute value from the context.
+// The returned bool indicates a successful typecast of the value to a string.
+func TelemetryFromContext(ctx context.Context, key string) (string, bool) {
+	val, ok := ctx.Value(telemetryKey(key)).(string)
+	return val, ok
+}
diff --git a/v2/invoke.go b/v2/invoke.go
@@ -36,19 +36,19 @@ import (
 	"time"
 
 	"github.com/googleapis/gax-go/v2/apierror"
-	"google.golang.org/grpc/metadata"
+	"github.com/googleapis/gax-go/v2/callctx"
 )
 
 // APICall is a user defined call stub.
 type APICall func(context.Context, CallSettings) error
 
 // withRetryCount returns a new context with the retry count appended to
-// gRPC metadata. The retry count is the number of retries that have been
+// the telemetry context. The retry count is the number of retries that have been
 // attempted. On the initial request, retry count is 0.
 // On a second request (the first retry), retry count is 1.
 func withRetryCount(ctx context.Context, retryCount int) context.Context {
-	// Add to gRPC metadata so it's visible to StatsHandlers
-	return metadata.AppendToOutgoingContext(ctx, "gcp.grpc.resend_count", strconv.Itoa(retryCount))
+	// Add to telemetry context so it's visible to observability wrappers
+	return callctx.WithTelemetryContext(ctx, "resend_count", strconv.Itoa(retryCount))
 }
 
 // Invoke calls the given APICall, performing retries as specified by opts, if
diff --git a/v2/invoke_test.go b/v2/invoke_test.go
@@ -40,9 +40,9 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/googleapis/gax-go/v2/apierror"
+	"github.com/googleapis/gax-go/v2/callctx"
 	"google.golang.org/genproto/googleapis/rpc/errdetails"
 	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
@@ -285,9 +285,8 @@ func TestInvokeRetryCount(t *testing.T) {
 			calls := 0
 			apiCall := func(ctx context.Context, _ CallSettings) error {
 				calls++
-				md, _ := metadata.FromOutgoingContext(ctx)
-				if vals := md["gcp.grpc.resend_count"]; len(vals) > 0 {
-					if count, err := strconv.Atoi(vals[0]); err == nil {
+				if val, ok := callctx.TelemetryFromContext(ctx, "resend_count"); ok {
+					if count, err := strconv.Atoi(val); err == nil {
 						retryCounts = append(retryCounts, count)
 					}
 				}
