moved contextvar into ActiveOperationMetric

Author: daniel-sanche

google/cloud/bigtable/data/_async/metrics_interceptor.py

@@ -96,11 +96,6 @@ class AsyncBigtableMetricsInterceptor(
     An async gRPC interceptor to add client metadata and print server metadata.
     """
 
-    _routing_cookie_context: ClassVar[
-        contextvars.ContextVar[dict[str, str|bytes]]
-    ] = contextvars.ContextVar("routing_cooke_context")
-
-
     @CrossSync.convert
     @_with_active_operation
     async def intercept_unary_unary(
@@ -136,7 +131,7 @@ async def intercept_unary_stream(
           - SampleRowKeys
         """
         try:
-            for k,v in self._routing_cookie_context.get({}).items():
+            for k,v in operation.routing_cookie.items():
                 client_call_details.metadata.add(k, v)
             return self._streaming_generator_wrapper(
                 operation, await continuation(client_call_details, request)

google/cloud/bigtable/data/_metrics/data_model.py

@@ -172,6 +172,8 @@ class ActiveOperationMetric:
     first_response_latency_ns: int | None = None
     # time waiting on flow control, in nanoseconds
     flow_throttling_time_ns: int = 0
+    # routing cookie, passed in between retry attempts
+    routing_cookie: dict[str, str|bytes] = field(default_factory=dict)
 
     _active_operation_context: ClassVar[
         contextvars.ContextVar[ActiveOperationMetric]

google/cloud/bigtable/data/_metrics/tracked_retry.py

@@ -64,7 +64,7 @@ def wrapper(exc: Exception) -> None:
                 # check for routing cookie:
                 cookie_headers = {k:v for k,v in metadata_dict.items() if k.startswith("x-goog-cbt-cookie")}
                 if cookie_headers:
-                    AsyncBigtableMetricsInterceptor._routing_cookie_context.set(cookie_headers)
+                    operation.routing_cookie = cookie_headers
         except Exception:
             # ignore errors in metadata collection
             pass