-
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 |
|---|---|---|
|
|
@@ -20,6 +20,8 @@ | |
| performance monitoring. | ||
| """ | ||
|
|
||
| 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 |
||
|
|
||
| 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 | ||
|
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? |
||
| """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 | ||
| 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,26 @@ 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() | ||
|
|
||
| @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) | ||
|
|
||
| @property | ||
| def current_metrics_tracer(self) -> MetricsTracer: | ||
| return SpannerMetricsTracerFactory._current_metrics_tracer_ctx.get() | ||
|
|
||
|
Comment on lines
+102
to
+105
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,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 | ||
|
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.