-
Notifications
You must be signed in to change notification settings - Fork 104
fix(metrics): prevent thread leak by ensuring singleton initialization #1492
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -23,10 +23,12 @@ | |
| * a :class:`~google.cloud.spanner_v1.instance.Instance` owns a | ||
| :class:`~google.cloud.spanner_v1.database.Database` | ||
| """ | ||
|
|
||
| import grpc | ||
| import os | ||
| import logging | ||
| import warnings | ||
| import threading | ||
|
|
||
| from google.api_core.gapic_v1 import client_info | ||
| from google.auth.credentials import AnonymousCredentials | ||
|
|
@@ -99,11 +101,50 @@ def _get_spanner_optimizer_statistics_package(): | |
|
|
||
| log = logging.getLogger(__name__) | ||
|
|
||
| _metrics_monitor_initialized = False | ||
sinhasubham marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| _metrics_monitor_lock = threading.Lock() | ||
|
|
||
|
|
||
| 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. | ||
|
|
||
|
|
@@ -252,30 +293,13 @@ def __init__( | |
| ): | ||
| warnings.warn(_EMULATOR_HOST_HTTP_SCHEME) | ||
| # Check flag to enable Spanner builtin metrics | ||
| global _metrics_monitor_initialized | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is this still needed here? |
||
| if ( | ||
| _get_spanner_enable_builtin_metrics_env() | ||
| and not disable_builtin_metrics | ||
| and HAS_GOOGLE_CLOUD_MONITORING_INSTALLED | ||
| ): | ||
| 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() | ||
| except Exception as e: | ||
| log.warning( | ||
| "Failed to initialize Spanner built-in metrics. Error: %s", e | ||
| ) | ||
| _initialize_metrics(project, credentials) | ||
| else: | ||
| SpannerMetricsTracerFactory(enabled=False) | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -23,13 +23,18 @@ | |
| from .spanner_metrics_tracer_factory import SpannerMetricsTracerFactory | ||
|
|
||
|
|
||
| from contextvars import Token | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: group this with the other import |
||
|
|
||
|
|
||
| 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 | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What is this? Could we add a small comment for it? |
||
|
|
||
| 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 | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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: | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: maybe add set_current_tracer and reset_current_tracer methods? |
||
| return SpannerMetricsTracerFactory._current_metrics_tracer_ctx.get() | ||
|
|
||
| @property | ||
| def current_metrics_tracer(self) -> MetricsTracer: | ||
| return SpannerMetricsTracerFactory._current_metrics_tracer_ctx.get() | ||
|
|
||
|
Comment on lines
+94
to
+97
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You've introduced both a static method The property Since all new code in this PR uses the clear and unambiguous static method |
||
| @staticmethod | ||
| def _generate_client_uid() -> str: | ||
| """Generate a client UID in the form of uuidv4@pid@hostname. | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,13 @@ | ||
| import pytest | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: add copyright header |
||
| 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 | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think that the existing code was correct. There are two different things that can be retried in Spanner:
attemptto be increased, even in this case, where the entire transaction is just a singleCommitcall.attemptshould be increased.