From c3ab7d339c858c3b7238349324893cda0dce4cd8 Mon Sep 17 00:00:00 2001 From: Chris Smith Date: Tue, 30 Dec 2025 16:09:22 -0700 Subject: [PATCH 1/5] feat(v2): update Invoke to add retry attempts to context Expose the retry attempt count to the transport layer for tracing --- v2/invoke.go | 26 +++++++++++++++++++++++++- v2/invoke_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/v2/invoke.go b/v2/invoke.go index 721d1af55..55e47a996 100644 --- a/v2/invoke.go +++ b/v2/invoke.go @@ -31,15 +31,32 @@ package gax import ( "context" + "strconv" "strings" "time" "github.com/googleapis/gax-go/v2/apierror" + "google.golang.org/grpc/metadata" ) // APICall is a user defined call stub. type APICall func(context.Context, CallSettings) error +type attemptKey struct{} + +// WithAttemptCount returns a new context with the attempt count attached. +func WithAttemptCount(ctx context.Context, count int) context.Context { + // For gRPC, we also append to metadata so it's visible to StatsHandlers + ctx = metadata.AppendToOutgoingContext(ctx, "gcp.grpc.resend_count", strconv.Itoa(count)) + return context.WithValue(ctx, attemptKey{}, count) +} + +// AttemptCountFromContext returns the attempt count from the context, if present. +func AttemptCountFromContext(ctx context.Context) (int, bool) { + v, ok := ctx.Value(attemptKey{}).(int) + return v, ok +} + // Invoke calls the given APICall, performing retries as specified by opts, if // any. func Invoke(ctx context.Context, call APICall, opts ...CallOption) error { @@ -78,8 +95,14 @@ func invoke(ctx context.Context, call APICall, settings CallSettings, sp sleeper ctx = c } + attempt := 0 + tracingEnabled := IsFeatureEnabled("TRACING") for { - err := call(ctx, settings) + ctxToUse := ctx + if tracingEnabled { + ctxToUse = WithAttemptCount(ctx, attempt) + } + err := call(ctxToUse, settings) if err == nil { return nil } @@ -110,5 +133,6 @@ func invoke(ctx context.Context, call APICall, settings CallSettings, sp sleeper } else if err = sp(ctx, d); err != nil { return err } + attempt++ } } diff --git a/v2/invoke_test.go b/v2/invoke_test.go index 8ae6d2803..2741d8bd9 100644 --- a/v2/invoke_test.go +++ b/v2/invoke_test.go @@ -32,6 +32,7 @@ package gax import ( "context" "errors" + "fmt" "testing" "time" @@ -264,3 +265,44 @@ func TestInvokeWithTimeout(t *testing.T) { }) } } + +func TestInvokeAttemptCount(t *testing.T) { + for _, enabled := range []bool{true, false} { + t.Run(fmt.Sprintf("enabled=%v", enabled), func(t *testing.T) { + TestOnlyResetIsFeatureEnabled() + defer TestOnlyResetIsFeatureEnabled() + + if enabled { + t.Setenv("GOOGLE_SDK_GO_EXPERIMENTAL_TRACING", "true") + } else { + t.Setenv("GOOGLE_SDK_GO_EXPERIMENTAL_TRACING", "false") + } + + const target = 3 + var attempts []int + calls := 0 + apiCall := func(ctx context.Context, _ CallSettings) error { + calls++ + if count, ok := AttemptCountFromContext(ctx); ok { + attempts = append(attempts, count) + } + if calls < target { + return errors.New("retry") + } + return nil + } + var settings CallSettings + WithRetry(func() Retryer { return boolRetryer(true) }).Resolve(&settings) + var sp recordSleeper + invoke(context.Background(), apiCall, settings, sp.sleep) + + var want []int + if enabled { + want = []int{0, 1, 2} + } + if diff := cmp.Diff(want, attempts); diff != "" { + t.Errorf("attempt count mismatch (-want +got):\n%s", diff) + } + }) + } +} From 38aae162aa1214e27a108c2b0252c67687f5abfc Mon Sep 17 00:00:00 2001 From: Chris Smith Date: Fri, 23 Jan 2026 13:11:26 -0700 Subject: [PATCH 2/5] rename attempt to retryCount --- v2/invoke.go | 24 +++++++++++++----------- v2/invoke_test.go | 12 ++++++------ 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/v2/invoke.go b/v2/invoke.go index 55e47a996..b24869766 100644 --- a/v2/invoke.go +++ b/v2/invoke.go @@ -42,18 +42,20 @@ import ( // APICall is a user defined call stub. type APICall func(context.Context, CallSettings) error -type attemptKey struct{} +type retryCountKey struct{} -// WithAttemptCount returns a new context with the attempt count attached. -func WithAttemptCount(ctx context.Context, count int) context.Context { +// withRetryCount returns a new context with the retry count attached. +// retryCount is the number of retries that have been attempted. +// The initial request is retryCount = 0. +func withRetryCount(ctx context.Context, retryCount int) context.Context { // For gRPC, we also append to metadata so it's visible to StatsHandlers - ctx = metadata.AppendToOutgoingContext(ctx, "gcp.grpc.resend_count", strconv.Itoa(count)) - return context.WithValue(ctx, attemptKey{}, count) + ctx = metadata.AppendToOutgoingContext(ctx, "gcp.grpc.resend_count", strconv.Itoa(retryCount)) + return context.WithValue(ctx, retryCountKey{}, retryCount) } -// AttemptCountFromContext returns the attempt count from the context, if present. -func AttemptCountFromContext(ctx context.Context) (int, bool) { - v, ok := ctx.Value(attemptKey{}).(int) +// retryCountFromContext returns the retry count from the context, if present. +func retryCountFromContext(ctx context.Context) (int, bool) { + v, ok := ctx.Value(retryCountKey{}).(int) return v, ok } @@ -95,12 +97,12 @@ func invoke(ctx context.Context, call APICall, settings CallSettings, sp sleeper ctx = c } - attempt := 0 + retryCount := 0 tracingEnabled := IsFeatureEnabled("TRACING") for { ctxToUse := ctx if tracingEnabled { - ctxToUse = WithAttemptCount(ctx, attempt) + ctxToUse = withRetryCount(ctx, retryCount) } err := call(ctxToUse, settings) if err == nil { @@ -133,6 +135,6 @@ func invoke(ctx context.Context, call APICall, settings CallSettings, sp sleeper } else if err = sp(ctx, d); err != nil { return err } - attempt++ + retryCount++ } } diff --git a/v2/invoke_test.go b/v2/invoke_test.go index 2741d8bd9..a88798e30 100644 --- a/v2/invoke_test.go +++ b/v2/invoke_test.go @@ -266,7 +266,7 @@ func TestInvokeWithTimeout(t *testing.T) { } } -func TestInvokeAttemptCount(t *testing.T) { +func TestInvokeRetryCount(t *testing.T) { for _, enabled := range []bool{true, false} { t.Run(fmt.Sprintf("enabled=%v", enabled), func(t *testing.T) { TestOnlyResetIsFeatureEnabled() @@ -279,12 +279,12 @@ func TestInvokeAttemptCount(t *testing.T) { } const target = 3 - var attempts []int + var retryCounts []int calls := 0 apiCall := func(ctx context.Context, _ CallSettings) error { calls++ - if count, ok := AttemptCountFromContext(ctx); ok { - attempts = append(attempts, count) + if count, ok := retryCountFromContext(ctx); ok { + retryCounts = append(retryCounts, count) } if calls < target { return errors.New("retry") @@ -300,8 +300,8 @@ func TestInvokeAttemptCount(t *testing.T) { if enabled { want = []int{0, 1, 2} } - if diff := cmp.Diff(want, attempts); diff != "" { - t.Errorf("attempt count mismatch (-want +got):\n%s", diff) + if diff := cmp.Diff(want, retryCounts); diff != "" { + t.Errorf("retry count mismatch (-want +got):\n%s", diff) } }) } From 6165b00c3164a4e2a14fa73778c661c590a06150 Mon Sep 17 00:00:00 2001 From: Chris Smith Date: Fri, 23 Jan 2026 13:38:04 -0700 Subject: [PATCH 3/5] remove unneeded context value and helper function --- v2/invoke.go | 11 +---------- v2/invoke_test.go | 9 +++++++-- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/v2/invoke.go b/v2/invoke.go index b24869766..eb5217718 100644 --- a/v2/invoke.go +++ b/v2/invoke.go @@ -42,21 +42,12 @@ import ( // APICall is a user defined call stub. type APICall func(context.Context, CallSettings) error -type retryCountKey struct{} - // withRetryCount returns a new context with the retry count attached. // retryCount is the number of retries that have been attempted. // The initial request is retryCount = 0. func withRetryCount(ctx context.Context, retryCount int) context.Context { // For gRPC, we also append to metadata so it's visible to StatsHandlers - ctx = metadata.AppendToOutgoingContext(ctx, "gcp.grpc.resend_count", strconv.Itoa(retryCount)) - return context.WithValue(ctx, retryCountKey{}, retryCount) -} - -// retryCountFromContext returns the retry count from the context, if present. -func retryCountFromContext(ctx context.Context) (int, bool) { - v, ok := ctx.Value(retryCountKey{}).(int) - return v, ok + return metadata.AppendToOutgoingContext(ctx, "gcp.grpc.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 a88798e30..a31b4ad73 100644 --- a/v2/invoke_test.go +++ b/v2/invoke_test.go @@ -33,6 +33,7 @@ import ( "context" "errors" "fmt" + "strconv" "testing" "time" @@ -41,6 +42,7 @@ import ( "github.com/googleapis/gax-go/v2/apierror" "google.golang.org/genproto/googleapis/rpc/errdetails" "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" ) @@ -283,8 +285,11 @@ func TestInvokeRetryCount(t *testing.T) { calls := 0 apiCall := func(ctx context.Context, _ CallSettings) error { calls++ - if count, ok := retryCountFromContext(ctx); ok { - retryCounts = append(retryCounts, count) + md, _ := metadata.FromOutgoingContext(ctx) + if vals := md["gcp.grpc.resend_count"]; len(vals) > 0 { + if count, err := strconv.Atoi(vals[0]); err == nil { + retryCounts = append(retryCounts, count) + } } if calls < target { return errors.New("retry") From d8b0b8ffa3318a8c42eea78da83dbe4bc9ed2ef2 Mon Sep 17 00:00:00 2001 From: Chris Smith Date: Fri, 23 Jan 2026 13:54:45 -0700 Subject: [PATCH 4/5] update withRetryCount docs --- v2/invoke.go | 10 ++++++---- v2/invoke_test.go | 8 ++++---- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/v2/invoke.go b/v2/invoke.go index eb5217718..57abb424d 100644 --- a/v2/invoke.go +++ b/v2/invoke.go @@ -42,11 +42,12 @@ import ( // APICall is a user defined call stub. type APICall func(context.Context, CallSettings) error -// withRetryCount returns a new context with the retry count attached. -// retryCount is the number of retries that have been attempted. -// The initial request is retryCount = 0. +// withRetryCount returns a new context with the retry count appended to +// gRPC metadata. The retry count is the number of retries that have been +// attempted. Immediately after the initial request, retry count is 0. +// Immediately after a second request (the first retry), retry count is 1. func withRetryCount(ctx context.Context, retryCount int) context.Context { - // For gRPC, we also append to metadata so it's visible to StatsHandlers + // Add to gRPC metadata so it's visible to StatsHandlers return metadata.AppendToOutgoingContext(ctx, "gcp.grpc.resend_count", strconv.Itoa(retryCount)) } @@ -89,6 +90,7 @@ func invoke(ctx context.Context, call APICall, settings CallSettings, sp sleeper } retryCount := 0 + // Feature gate: GOOGLE_SDK_GO_EXPERIMENTAL_TRACING=true tracingEnabled := IsFeatureEnabled("TRACING") for { ctxToUse := ctx diff --git a/v2/invoke_test.go b/v2/invoke_test.go index a31b4ad73..2a2dd4b2b 100644 --- a/v2/invoke_test.go +++ b/v2/invoke_test.go @@ -269,12 +269,12 @@ func TestInvokeWithTimeout(t *testing.T) { } func TestInvokeRetryCount(t *testing.T) { - for _, enabled := range []bool{true, false} { - t.Run(fmt.Sprintf("enabled=%v", enabled), func(t *testing.T) { + for _, tracingEnabled := range []bool{true, false} { + t.Run(fmt.Sprintf("tracingEnabled=%v", tracingEnabled), func(t *testing.T) { TestOnlyResetIsFeatureEnabled() defer TestOnlyResetIsFeatureEnabled() - if enabled { + if tracingEnabled { t.Setenv("GOOGLE_SDK_GO_EXPERIMENTAL_TRACING", "true") } else { t.Setenv("GOOGLE_SDK_GO_EXPERIMENTAL_TRACING", "false") @@ -302,7 +302,7 @@ func TestInvokeRetryCount(t *testing.T) { invoke(context.Background(), apiCall, settings, sp.sleep) var want []int - if enabled { + if tracingEnabled { want = []int{0, 1, 2} } if diff := cmp.Diff(want, retryCounts); diff != "" { From 3eba94ec76ae090cf3b87dd9caa05d3c49b3fe5c Mon Sep 17 00:00:00 2001 From: Chris Smith Date: Fri, 23 Jan 2026 14:01:42 -0700 Subject: [PATCH 5/5] Update docs Co-authored-by: Wesley Tarle --- v2/invoke.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/v2/invoke.go b/v2/invoke.go index 57abb424d..af1eebe3e 100644 --- a/v2/invoke.go +++ b/v2/invoke.go @@ -44,8 +44,8 @@ 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 -// attempted. Immediately after the initial request, retry count is 0. -// Immediately after a second request (the first retry), retry count is 1. +// 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))