fix: Metrics thread-safety refactor and Batch.commit idempotency fix

sinhasubham · sinhasubham · commit ffb60f6143e9 · 2026-02-11T15:03:19.000+05:30
diff --git a/.kokoro/presubmit/presubmit.cfg b/.kokoro/presubmit/presubmit.cfg
@@ -3,5 +3,5 @@
 # Only run a subset of all nox sessions
 env_vars: {
     key: "NOX_SESSION"
-    value: "unit-3.9 unit-3.12 cover docs docfx"
+    value: "unit-3.10 unit-3.12 cover docs docfx"
 }
diff --git a/google/cloud/spanner_v1/batch.py b/google/cloud/spanner_v1/batch.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Context manager for Cloud Spanner batched writes."""
+
 import functools
 from typing import List, Optional
 
@@ -27,7 +28,6 @@
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
     _merge_Transaction_Options,
-    AtomicCounter,
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
 from google.cloud.spanner_v1 import RequestOptions
@@ -252,9 +252,6 @@ def wrapped_method():
                     max_commit_delay=max_commit_delay,
                     request_options=request_options,
                 )
-                # This code is retried due to ABORTED, hence nth_request
-                # should be increased. attempt can only be increased if
-                # we encounter UNAVAILABLE or INTERNAL.
                 call_metadata, error_augmenter = database.with_error_augmentation(
                     getattr(database, "_next_nth_request", 0),
                     1,
@@ -376,8 +373,6 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
             observability_options=getattr(database, "observability_options", None),
             metadata=metadata,
         ) as span, MetricsCapture():
-            attempt = AtomicCounter(0)
-            nth_request = getattr(database, "_next_nth_request", 0)
 
             def wrapped_method():
                 batch_write_request = BatchWriteRequest(
@@ -390,8 +385,8 @@ def wrapped_method():
                     api.batch_write,
                     request=batch_write_request,
                     metadata=database.metadata_with_request_id(
-                        nth_request,
-                        attempt.increment(),
+                        getattr(database, "_next_nth_request", 0),
+                        1,
                         metadata,
                         span,
                     ),
diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py
@@ -23,6 +23,7 @@
 * a :class:`~google.cloud.spanner_v1.instance.Instance` owns a
   :class:`~google.cloud.spanner_v1.database.Database`
 """
+
 import grpc
 import os
 import logging
@@ -108,6 +109,42 @@ def _get_spanner_enable_builtin_metrics_env():
     return os.getenv(SPANNER_DISABLE_BUILTIN_METRICS_ENV_VAR) != "true"
 
 
+def _initialize_metrics(project, credentials):
+    """
+    Initializes the Spanner built-in metrics.
+
+    This function sets up the OpenTelemetry MeterProvider and the SpannerMetricsTracerFactory.
+    It uses a lock to ensure that initialization happens only once.
+    """
+    global _metrics_monitor_initialized
+    if not _metrics_monitor_initialized:
+        with _metrics_monitor_lock:
+            if not _metrics_monitor_initialized:
+                meter_provider = metrics.NoOpMeterProvider()
+                try:
+                    if not _get_spanner_emulator_host():
+                        meter_provider = MeterProvider(
+                            metric_readers=[
+                                PeriodicExportingMetricReader(
+                                    CloudMonitoringMetricsExporter(
+                                        project_id=project,
+                                        credentials=credentials,
+                                    ),
+                                    export_interval_millis=METRIC_EXPORT_INTERVAL_MS,
+                                ),
+                            ]
+                        )
+                    metrics.set_meter_provider(meter_provider)
+                    SpannerMetricsTracerFactory()
+                    _metrics_monitor_initialized = True
+                except Exception as e:
+                    # log is already defined at module level
+                    log.warning(
+                        "Failed to initialize Spanner built-in metrics. Error: %s",
+                        e,
+                    )
+
+
 class Client(ClientWithProject):
     """Client for interacting with Cloud Spanner API.
 
@@ -255,38 +292,12 @@ def __init__(
             "http://" in self._emulator_host or "https://" in self._emulator_host
         ):
             warnings.warn(_EMULATOR_HOST_HTTP_SCHEME)
-        # Check flag to enable Spanner builtin metrics
-        global _metrics_monitor_initialized
         if (
             _get_spanner_enable_builtin_metrics_env()
             and not disable_builtin_metrics
             and HAS_GOOGLE_CLOUD_MONITORING_INSTALLED
         ):
-            if not _metrics_monitor_initialized:
-                with _metrics_monitor_lock:
-                    if not _metrics_monitor_initialized:
-                        meter_provider = metrics.NoOpMeterProvider()
-                        try:
-                            if not _get_spanner_emulator_host():
-                                meter_provider = MeterProvider(
-                                    metric_readers=[
-                                        PeriodicExportingMetricReader(
-                                            CloudMonitoringMetricsExporter(
-                                                project_id=project,
-                                                credentials=credentials,
-                                            ),
-                                            export_interval_millis=METRIC_EXPORT_INTERVAL_MS,
-                                        ),
-                                    ]
-                                )
-                            metrics.set_meter_provider(meter_provider)
-                            SpannerMetricsTracerFactory()
-                            _metrics_monitor_initialized = True
-                        except Exception as e:
-                            log.warning(
-                                "Failed to initialize Spanner built-in metrics. Error: %s",
-                                e,
-                            )
+            _initialize_metrics(project, credentials)
         else:
             SpannerMetricsTracerFactory(enabled=False)
 
diff --git a/google/cloud/spanner_v1/metrics/metrics_capture.py b/google/cloud/spanner_v1/metrics/metrics_capture.py
@@ -20,6 +20,8 @@
 performance monitoring.
 """
 
+from contextvars import Token
+
 from .spanner_metrics_tracer_factory import SpannerMetricsTracerFactory
 
 
@@ -30,6 +32,9 @@ class MetricsCapture:
     the start and completion of metrics tracing for a given operation.
     """
 
+    _token: Token
+    """Token to reset the context variable after the operation completes."""
+
     def __enter__(self):
         """Enter the runtime context related to this object.
 
@@ -45,11 +50,11 @@ def __enter__(self):
             return self
 
         # Define a new metrics tracer for the new operation
-        SpannerMetricsTracerFactory.current_metrics_tracer = (
-            factory.create_metrics_tracer()
-        )
-        if SpannerMetricsTracerFactory.current_metrics_tracer:
-            SpannerMetricsTracerFactory.current_metrics_tracer.record_operation_start()
+        # Set the context var and keep the token for reset
+        tracer = factory.create_metrics_tracer()
+        self._token = SpannerMetricsTracerFactory.set_current_tracer(tracer)
+        if tracer:
+            tracer.record_operation_start()
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -70,6 +75,11 @@ def __exit__(self, exc_type, exc_value, traceback):
         if not SpannerMetricsTracerFactory().enabled:
             return False
 
-        if SpannerMetricsTracerFactory.current_metrics_tracer:
-            SpannerMetricsTracerFactory.current_metrics_tracer.record_operation_completion()
+        tracer = SpannerMetricsTracerFactory.get_current_tracer()
+        if tracer:
+            tracer.record_operation_completion()
+
+        # Reset the context var using the token
+        if getattr(self, "_token", None):
+            SpannerMetricsTracerFactory.reset_current_tracer(self._token)
         return False  # Propagate the exception if any
diff --git a/google/cloud/spanner_v1/metrics/metrics_interceptor.py b/google/cloud/spanner_v1/metrics/metrics_interceptor.py
@@ -97,22 +97,17 @@ def _set_metrics_tracer_attributes(self, resources: Dict[str, str]) -> None:
         Args:
             resources (Dict[str, str]): A dictionary containing project, instance, and database information.
         """
-        if SpannerMetricsTracerFactory.current_metrics_tracer is None:
+        tracer = SpannerMetricsTracerFactory.get_current_tracer()
+        if tracer is None:
             return
 
         if resources:
             if "project" in resources:
-                SpannerMetricsTracerFactory.current_metrics_tracer.set_project(
-                    resources["project"]
-                )
+                tracer.set_project(resources["project"])
             if "instance" in resources:
-                SpannerMetricsTracerFactory.current_metrics_tracer.set_instance(
-                    resources["instance"]
-                )
+                tracer.set_instance(resources["instance"])
             if "database" in resources:
-                SpannerMetricsTracerFactory.current_metrics_tracer.set_database(
-                    resources["database"]
-                )
+                tracer.set_database(resources["database"])
 
     def intercept(self, invoked_method, request_or_iterator, call_details):
         """Intercept gRPC calls to collect metrics.
@@ -126,31 +121,32 @@ def intercept(self, invoked_method, request_or_iterator, call_details):
             The RPC response
         """
         factory = SpannerMetricsTracerFactory()
-        if (
-            SpannerMetricsTracerFactory.current_metrics_tracer is None
-            or not factory.enabled
-        ):
+        tracer = SpannerMetricsTracerFactory.get_current_tracer()
+        if tracer is None or not factory.enabled:
             return invoked_method(request_or_iterator, call_details)
 
         # Setup Metric Tracer attributes from call details
-        ## Extract  Project / Instance / Databse from header information
-        resources = self._extract_resource_from_path(call_details.metadata)
-        self._set_metrics_tracer_attributes(resources)
+        ## Extract Project / Instance / Database from header information if not already set
+        if not (
+            tracer.client_attributes.get("project_id")
+            and tracer.client_attributes.get("instance_id")
+            and tracer.client_attributes.get("database")
+        ):
+            resources = self._extract_resource_from_path(call_details.metadata)
+            self._set_metrics_tracer_attributes(resources)
 
         ## Format method to be be spanner.<method name>
         method_name = self._remove_prefix(
             call_details.method, SPANNER_METHOD_PREFIX
         ).replace("/", ".")
 
-        SpannerMetricsTracerFactory.current_metrics_tracer.set_method(method_name)
-        SpannerMetricsTracerFactory.current_metrics_tracer.record_attempt_start()
+        tracer.set_method(method_name)
+        tracer.record_attempt_start()
         response = invoked_method(request_or_iterator, call_details)
-        SpannerMetricsTracerFactory.current_metrics_tracer.record_attempt_completion()
+        tracer.record_attempt_completion()
 
         # Process and send GFE metrics if enabled
-        if SpannerMetricsTracerFactory.current_metrics_tracer.gfe_enabled:
+        if tracer.gfe_enabled:
             metadata = response.initial_metadata()
-            SpannerMetricsTracerFactory.current_metrics_trace.record_gfe_metrics(
-                metadata
-            )
+            tracer.record_gfe_metrics(metadata)
         return response
diff --git a/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py b/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py
@@ -19,6 +19,7 @@
 import os
 import logging
 from .constants import SPANNER_SERVICE_NAME
+import contextvars
 
 try:
     import mmh3
@@ -43,7 +44,9 @@ class SpannerMetricsTracerFactory(MetricsTracerFactory):
     """A factory for creating SpannerMetricsTracer instances."""
 
     _metrics_tracer_factory: "SpannerMetricsTracerFactory" = None
-    current_metrics_tracer: MetricsTracer = None
+    _current_metrics_tracer_ctx = contextvars.ContextVar(
+        "current_metrics_tracer", default=None
+    )
 
     def __new__(
         cls, enabled: bool = True, gfe_enabled: bool = False
@@ -80,10 +83,22 @@ def __new__(
             cls._metrics_tracer_factory.gfe_enabled = gfe_enabled
 
             if cls._metrics_tracer_factory.enabled != enabled:
-                cls._metrics_tracer_factory.enabeld = enabled
+                cls._metrics_tracer_factory.enabled = enabled
 
         return cls._metrics_tracer_factory
 
+    @staticmethod
+    def get_current_tracer() -> MetricsTracer:
+        return SpannerMetricsTracerFactory._current_metrics_tracer_ctx.get()
+
+    @staticmethod
+    def set_current_tracer(tracer: MetricsTracer) -> contextvars.Token:
+        return SpannerMetricsTracerFactory._current_metrics_tracer_ctx.set(tracer)
+
+    @staticmethod
+    def reset_current_tracer(token: contextvars.Token):
+        SpannerMetricsTracerFactory._current_metrics_tracer_ctx.reset(token)
+
     @staticmethod
     def _generate_client_uid() -> str:
         """Generate a client UID in the form of uuidv4@pid@hostname.
diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
@@ -0,0 +1,27 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from unittest.mock import patch
+
+
+@pytest.fixture(autouse=True)
+def mock_periodic_exporting_metric_reader():
+    """Globally mock PeriodicExportingMetricReader to prevent real network calls."""
+    with patch(
+        "google.cloud.spanner_v1.client.PeriodicExportingMetricReader"
+    ) as mock_client_reader, patch(
+        "opentelemetry.sdk.metrics.export.PeriodicExportingMetricReader"
+    ):
+        yield mock_client_reader
diff --git a/tests/unit/test_metrics.py b/tests/unit/test_metrics.py
diff --git a/tests/unit/test_metrics_concurrency.py b/tests/unit/test_metrics_concurrency.py
diff --git a/tests/unit/test_metrics_interceptor.py b/tests/unit/test_metrics_interceptor.py

Original file line number	Diff line number	Diff line change
`@@ -3,5 +3,5 @@`
`3`	`3`	`# Only run a subset of all nox sessions`
`4`	`4`	`env_vars: {`
`5`	`5`	`key: "NOX_SESSION"`
`6`		`- value: "unit-3.9 unit-3.12 cover docs docfx"`
	`6`	`+ value: "unit-3.10 unit-3.12 cover docs docfx"`
`7`	`7`	`}`