diff --git a/v2/callctx/callctx.go b/v2/callctx/callctx.go index f5af5c990..b28df61d0 100644 --- 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 index af1eebe3e..594ac168c 100644 --- 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 index 2a2dd4b2b..98dc909f5 100644 --- 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) } }