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

sinhasubham · sinhasubham · commit 1341f21f2cf2 · 2026-02-10T12:37:02.000+05:30
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
 
@@ -242,6 +243,8 @@ def commit(
             observability_options=getattr(database, "observability_options", None),
             metadata=metadata,
         ) as span, MetricsCapture():
+            nth_request = getattr(database, "_next_nth_request", 0)
+            attempt = AtomicCounter(0)
 
             def wrapped_method():
                 commit_request = CommitRequest(
@@ -256,8 +259,8 @@ def wrapped_method():
                 # 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,
+                    nth_request,
+                    attempt.increment(),
                     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.
 
@@ -262,31 +299,7 @@ def __init__(
             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
@@ -23,13 +23,18 @@
 from .spanner_metrics_tracer_factory import SpannerMetricsTracerFactory
 
 
+from contextvars import Token
+
+
 class MetricsCapture:
     """Context manager for capturing metrics in Cloud Spanner operations.
 
     This class provides a context manager interface to automatically handle
     the start and completion of metrics tracing for a given operation.
     """
 
+    _token: Token
+
     def __enter__(self):
         """Enter the runtime context related to this object.
 
@@ -45,11 +50,13 @@ def __enter__(self):
             return self
 
         # Define a new metrics tracer for the new operation
-        SpannerMetricsTracerFactory.current_metrics_tracer = (
-            factory.create_metrics_tracer()
+        # Set the context var and keep the token for reset
+        tracer = factory.create_metrics_tracer()
+        self._token = SpannerMetricsTracerFactory._current_metrics_tracer_ctx.set(
+            tracer
         )
-        if SpannerMetricsTracerFactory.current_metrics_tracer:
-            SpannerMetricsTracerFactory.current_metrics_tracer.record_operation_start()
+        if tracer:
+            tracer.record_operation_start()
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -70,6 +77,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._current_metrics_tracer_ctx.get()
+        if tracer:
+            tracer.record_operation_completion()
+
+        # Reset the context var using the token
+        if getattr(self, "_token", None):
+            SpannerMetricsTracerFactory._current_metrics_tracer_ctx.reset(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,10 +121,8 @@ 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
@@ -142,15 +135,13 @@ def intercept(self, invoked_method, request_or_iterator, call_details):
             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,18 @@ 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()
+
+    @property
+    def current_metrics_tracer(self) -> MetricsTracer:
+        return SpannerMetricsTracerFactory._current_metrics_tracer_ctx.get()
+
     @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,13 @@
+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
@@ -65,19 +65,25 @@ def patched_client(monkeypatch):
 
     client_module._metrics_monitor_initialized = False
 
-    with patch("google.cloud.spanner_v1.client.CloudMonitoringMetricsExporter"):
+    with patch(
+        "google.cloud.spanner_v1.metrics.metrics_exporter.MetricServiceClient"
+    ), patch(
+        "google.cloud.spanner_v1.metrics.metrics_exporter.CloudMonitoringMetricsExporter"
+    ), patch(
+        "opentelemetry.sdk.metrics.export.PeriodicExportingMetricReader"
+    ):
         client = Client(
             project="test",
             credentials=TestCredentials(),
-            # client_options={"api_endpoint": "none"}
         )
         yield client
 
     # Resetting
     metrics.set_meter_provider(metrics.NoOpMeterProvider())
     SpannerMetricsTracerFactory._metrics_tracer_factory = None
-    SpannerMetricsTracerFactory.current_metrics_tracer = None
-    client_module._metrics_monitor_initialized = False
+    # Reset context var
+    ctx = SpannerMetricsTracerFactory._current_metrics_tracer_ctx
+    ctx.set(None)
 
 
 def test_metrics_emission_with_failure_attempt(patched_client):
@@ -92,10 +98,14 @@ def test_metrics_emission_with_failure_attempt(patched_client):
     original_intercept = metrics_interceptor.intercept
     first_attempt = True
 
+    captured_tracer_list = []
+
     def mocked_raise(*args, **kwargs):
         raise ServiceUnavailable("Service Unavailable")
 
     def mocked_call(*args, **kwargs):
+        # Capture the tracer while it is active
+        captured_tracer_list.append(SpannerMetricsTracerFactory.get_current_tracer())
         return _UnaryOutcome(MagicMock(), MagicMock())
 
     def intercept_wrapper(invoked_method, request_or_iterator, call_details):
@@ -113,11 +123,14 @@ def intercept_wrapper(invoked_method, request_or_iterator, call_details):
 
     metrics_interceptor.intercept = intercept_wrapper
     patch_path = "google.cloud.spanner_v1.metrics.metrics_exporter.CloudMonitoringMetricsExporter.export"
+
     with patch(patch_path):
         with database.snapshot():
             pass
 
     # Verify that the attempt count increased from the failed initial attempt
-    assert (
-        SpannerMetricsTracerFactory.current_metrics_tracer.current_op.attempt_count
-    ) == 2
+    # We use the captured tracer from the SUCCESSFUL attempt (the second one)
+    assert len(captured_tracer_list) > 0
+    tracer = captured_tracer_list[0]
+    assert tracer is not None
+    assert tracer.current_op.attempt_count == 2
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
