feat(client): Enable routing cookie and attempt headers for enhanced retries (#12964)

bhshkh · web-flow · commit 96cfd47ad729 · 2025-10-16T10:19:25.000-04:00
Implements client-side logic to support the Cloud Bigtable routing cookie protocol (see [go/cbt-support-routing-cookie](http://go/cbt-support-routing-cookie)). This feature enhances server-initiated retries and improves resilience. Specific modifications: * **Feature Flag:** The client now signals support for this feature by setting `RoutingCookie: true` in the `bigtable-features` header. * **Cookie Propagation:** Any headers prefixed with `x-goog-cbt-cookie-` received in response metadata are collected and sent back on subsequent requests for the same operation. * **Attempt Header:** * The `x-goog-cbt-attempt` header is added to each outgoing RPC attempt. * This integer value starts at 0 and is incremented before each call. * Crucially, the counter is reset to 0 if the client receives *any* metadata (headers or trailers) from the server, indicating the request reached the backend. This helps differentiate between network issues and server-side retries. * **Internal Tracking:** The `opTracer` in `metrics.go` now stores the current cookies and the `routingAttempt` number. These changes enable more sophisticated load balancing and retry strategies, by giving the service and routing layers visibility into the client's retry state and a mechanism to influence client behavior. Fixes: #12658
diff --git a/bigtable/bigtable.go b/bigtable/bigtable.go
@@ -55,6 +55,9 @@ const (
 	queryExpiredViolationType        = "PREPARED_QUERY_EXPIRED"
 	preparedQueryExpireEarlyDuration = time.Second
 	methodNameReadRows               = "ReadRows"
+
+	// For routing cookie
+	cookiePrefix = "x-goog-cbt-cookie-"
 )
 
 var errNegativeRowLimit = errors.New("bigtable: row limit cannot be negative")
@@ -398,6 +401,7 @@ type Table struct {
 // and enabled on the client
 func (c *Client) newFeatureFlags() metadata.MD {
 	ff := btpb.FeatureFlags{
+		RoutingCookie:            true,
 		ReverseScans:             true,
 		LastScannedRowResponses:  true,
 		ClientSideMetricsEnabled: c.metricsTracerFactory.enabled,
@@ -2281,15 +2285,37 @@ func gaxInvokeWithRecorder(ctx context.Context, mt *builtinMetricsTracer, method
 	mt.setMethod(method)
 
 	callWrapper := func(ctx context.Context, callSettings gax.CallSettings) error {
+		op := &mt.currOp
+		// Inject cookie and attempt information
+		md := metadata.New(nil)
+		for k, v := range op.cookies {
+			md.Append(k, v)
+		}
+
+		existingMD, _ := metadata.FromOutgoingContext(ctx)
+		finalMD := metadata.Join(existingMD, md)
+		newCtx := metadata.NewOutgoingContext(ctx, finalMD)
+
 		mt.recordAttemptStart()
 
 		// f makes calls to CBT service
-		err := f(ctx, &attemptHeaderMD, &attempTrailerMD, callSettings)
+		err := f(newCtx, &attemptHeaderMD, &attempTrailerMD, callSettings)
 
 		// Record attempt specific metrics
 		mt.recordAttemptCompletion(attemptHeaderMD, attempTrailerMD, err)
+
+		extractCookies(attemptHeaderMD, op)
+		extractCookies(attempTrailerMD, op)
 		return err
 	}
 
 	return gax.Invoke(ctx, callWrapper, opts...)
 }
+
+func extractCookies(md metadata.MD, op *opTracer) {
+	for k, v := range md {
+		if strings.HasPrefix(k, cookiePrefix) {
+			op.cookies[k] = v[len(v)-1]
+		}
+	}
+}
diff --git a/bigtable/internal/testproxy/known_failures.txt b/bigtable/internal/testproxy/known_failures.txt
@@ -1,6 +1 @@
-TestMutateRows_Retry_WithRoutingCookie\|
-TestReadRow_Retry_WithRoutingCookie\|
-TestReadRows_Retry_WithRoutingCookie\|
-TestReadRows_Retry_WithRoutingCookie_MultipleErrorResponses\|
-TestReadRows_Retry_WithRetryInfo_MultipleErrorResponse\|
-TestSampleRowKeys_Retry_WithRoutingCookie
+TestReadRows_Retry_WithRetryInfo_MultipleErrorResponse
diff --git a/bigtable/metrics.go b/bigtable/metrics.go
@@ -457,6 +457,9 @@ type opTracer struct {
 	currAttempt attemptTracer
 
 	appBlockingLatency float64
+
+	// For routing cookie and gRPC attempt number
+	cookies map[string]string
 }
 
 func (o *opTracer) setStartTime(t time.Time) {
@@ -524,14 +527,20 @@ func (a *attemptTracer) setServerLatencyErr(err error) {
 }
 
 func (tf *builtinMetricsTracerFactory) createBuiltinMetricsTracer(ctx context.Context, tableName string, isStreaming bool) builtinMetricsTracer {
-	if !tf.enabled {
-		return builtinMetricsTracer{builtInEnabled: false}
-	}
 	// Operation has started but not the attempt.
 	// So, create only operation tracer and not attempt tracer
-	currOpTracer := opTracer{}
+	currOpTracer := opTracer{
+		cookies: make(map[string]string),
+	}
 	currOpTracer.setStartTime(time.Now())
 
+	if !tf.enabled {
+		return builtinMetricsTracer{
+			builtInEnabled: false,
+			currOp:         currOpTracer,
+		}
+	}
+
 	return builtinMetricsTracer{
 		ctx:            ctx,
 		builtInEnabled: tf.enabled,
