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

Author: sinhasubham

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,
                 )

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

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

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

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.

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

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

tests/unit/test_metrics_concurrency.py

@@ -0,0 +1,80 @@
+import threading
+import time
+import unittest
+from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import (
+    SpannerMetricsTracerFactory,
+)
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+
+
+class TestMetricsConcurrency(unittest.TestCase):
+    def setUp(self):
+        # Reset factory singleton
+        SpannerMetricsTracerFactory._metrics_tracer_factory = None
+
+    def test_concurrent_tracers(self):
+        """Verify that concurrent threads have isolated tracers."""
+        factory = SpannerMetricsTracerFactory(enabled=True)
+        # Ensure enabled
+        factory.enabled = True
+
+        errors = []
+
+        def worker(idx):
+            try:
+                # Simulate a request workflow
+                with MetricsCapture():
+                    # Capture should have set a tracer
+                    tracer = SpannerMetricsTracerFactory.get_current_tracer()
+                    if tracer is None:
+                        errors.append(f"Thread {idx}: Tracer is None inside Capture")
+                        return
+
+                    # Set a unique attribute for this thread
+                    project_name = f"project-{idx}"
+                    tracer.set_project(project_name)
+
+                    # Simulate some work
+                    time.sleep(0.01)
+
+                    # Verify verify we still have OUR tracer
+                    current_tracer = SpannerMetricsTracerFactory.get_current_tracer()
+                    if current_tracer.client_attributes["project_id"] != project_name:
+                        errors.append(
+                            f"Thread {idx}: Tracer project mismatch. Expected {project_name}, got {current_tracer.client_attributes.get('project_id')}"
+                        )
+
+                    # Check interceptor logic (simulated)
+                    # Interceptor reads from factory.current_metrics_tracer
+                    interceptor_tracer = (
+                        SpannerMetricsTracerFactory.get_current_tracer()
+                    )
+                    if interceptor_tracer is not tracer:
+                        errors.append(f"Thread {idx}: Interceptor tracer mismatch")
+
+            except Exception as e:
+                errors.append(f"Thread {idx}: Exception {e}")
+
+        threads = []
+        for i in range(10):
+            t = threading.Thread(target=worker, args=(i,))
+            threads.append(t)
+            t.start()
+
+        for t in threads:
+            t.join()
+
+        self.assertEqual(errors, [], f"Concurrency errors found: {errors}")
+
+    def test_context_var_cleanup(self):
+        """Verify tracer is cleaned up after ContextVar reset."""
+        SpannerMetricsTracerFactory(enabled=True)
+
+        with MetricsCapture():
+            self.assertIsNotNone(SpannerMetricsTracerFactory.get_current_tracer())
+
+        self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer())
+
+
+if __name__ == "__main__":
+    unittest.main()