feat(storage): add OpenTelemetry tracing support for gRPC bidi APIs

packages/google-cloud-storage/google/cloud/storage/_bucket_metadata_cache.py

@@ -14,6 +14,8 @@
 
 """In-memory LRU cache for bucket metadata supporting App-centric Observability (ACO)."""
 
+import asyncio
+import contextvars
 import logging
 import threading
 
@@ -37,6 +39,20 @@ def __init__(self, client, max_size=10000):
         self._lock = threading.Lock()
         self._inflight_fetches = set()
         self._inflight_checks = set()
+        self._background_tasks = set()
+
+    def _spawn_async_task(self, coro, inflight_set, bucket_name):
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            coro.close()
+            with self._lock:
+                inflight_set.discard(bucket_name)
+            return
+        # Detach from the caller's tracing context.
+        task = contextvars.Context().run(loop.create_task, coro)
+        self._background_tasks.add(task)
+        task.add_done_callback(self._background_tasks.discard)
 
     def get(self, bucket_name):
         """Thread-safely retrieve cached metadata without queueing fetch."""
@@ -49,21 +65,29 @@ def get_or_queue_fetch(self, bucket_name):
         Returns None immediately on cache miss so caller does not block.
         """
         with self._lock:
-            if bucket_name in self._cache:
-                return self._cache.get(bucket_name)
-            elif bucket_name in self._inflight_fetches:
+            cached = self._cache.get(bucket_name)
+            if cached is not None:
+                return cached
+            if bucket_name in self._inflight_fetches:
                 # This handles a thundering herd where 'n' threads
                 # simultaneously experience a cache miss while 1 is already
                 # fetching metadata. The remaining n - 1 threads should
                 # bypass starting duplicate fetches.
                 return None
-            else:
-                # fire a background thread and get bucket metadata.
-                self._inflight_fetches.add(bucket_name)
-                threading.Thread(
-                    target=self._fetch_background, args=(bucket_name,), daemon=True
-                ).start()
-                return None
+            # fire a background fetch and get bucket metadata.
+            self._inflight_fetches.add(bucket_name)
+
+        if getattr(self._client, "_is_async_grpc_client", False) is True:
+            self._spawn_async_task(
+                self._fetch_background_async(bucket_name),
+                self._inflight_fetches,
+                bucket_name,
+            )
+        else:
+            threading.Thread(
+                target=self._fetch_background, args=(bucket_name,), daemon=True
+            ).start()
+        return None
 
     def check_and_evict(self, bucket_name):
         """Asynchronously verify if a bucket exists on 404 and evict if deleted."""
@@ -73,6 +97,14 @@ def check_and_evict(self, bucket_name):
             if bucket_name in self._inflight_checks:
                 return
             self._inflight_checks.add(bucket_name)
+
+        if getattr(self._client, "_is_async_grpc_client", False) is True:
+            self._spawn_async_task(
+                self._verify_existence_background_async(bucket_name),
+                self._inflight_checks,
+                bucket_name,
+            )
+        else:
             threading.Thread(
                 target=self._verify_existence_background,
                 args=(bucket_name,),
@@ -92,6 +124,27 @@ def _verify_existence_background(self, bucket_name):
             with self._lock:
                 self._inflight_checks.discard(bucket_name)
 
+    async def _verify_existence_background_async(self, bucket_name):
+        try:
+            from google.cloud import _storage_v2 as storage_v2
+
+            request = storage_v2.GetBucketRequest(
+                name=f"projects/_/buckets/{bucket_name}"
+            )
+            bucket = await self._client.grpc_client.get_bucket(
+                request=request, timeout=10.0
+            )
+            self.update_from_bucket(bucket, bucket_name=bucket_name)
+        except (NotFound, api_exceptions.NotFound):
+            self.evict(bucket_name)
+        except Exception as e:
+            logger.debug(
+                f"Async background verification for bucket existence failed for {bucket_name}: {e}"
+            )
+        finally:
+            with self._lock:
+                self._inflight_checks.discard(bucket_name)
+
     def _fetch_background(self, bucket_name):
         """Asynchronously fetch bucket metadata and update the cache."""
         try:
@@ -114,28 +167,76 @@ def _fetch_background(self, bucket_name):
             with self._lock:
                 self._inflight_fetches.discard(bucket_name)
 
-    def update_from_bucket(self, bucket):
-        """Update cache from a Bucket instance."""
-        if not bucket or not bucket.name:
+    async def _fetch_background_async(self, bucket_name):
+        """Asynchronously fetch bucket metadata via gRPC and update the cache."""
+        try:
+            from google.cloud import _storage_v2 as storage_v2
+
+            request = storage_v2.GetBucketRequest(
+                name=f"projects/_/buckets/{bucket_name}"
+            )
+            bucket = await self._client.grpc_client.get_bucket(
+                request=request, timeout=10.0
+            )
+            self.update_from_bucket(bucket, bucket_name=bucket_name)
+        except (NotFound, api_exceptions.NotFound):
+            self.evict(bucket_name)
+        except api_exceptions.Forbidden:
+            self.update_cache(
+                bucket_name,
+                f"//storage.googleapis.com/projects/_/buckets/{bucket_name}",
+                "global",
+            )
+        except Exception as e:
+            logger.debug(
+                f"Async background fetch for bucket metadata failed for {bucket_name}: {e}"
+            )
+        finally:
+            with self._lock:
+                self._inflight_fetches.discard(bucket_name)
+
+    def update_from_bucket(self, bucket, bucket_name=None):
+        """Update cache from a Bucket instance or storage_v2.Bucket proto."""
+        if not bucket:
+            return
+        name = bucket_name or getattr(bucket, "name", None)
+        if not name or not isinstance(name, str):
             return
+        if name.startswith("projects/") and "/buckets/" in name:
+            name = name.split("/buckets/", 1)[1]
 
         project_number = getattr(bucket, "project_number", None)
-        location = getattr(bucket, "location", None) or "global"
-        location = location.lower()
-        location_type = getattr(bucket, "location_type", None) or "region"
-        location_type = location_type.lower()
+        if not project_number:
+            proj_attr = getattr(bucket, "project", None)
+            if isinstance(proj_attr, (str, int)):
+                proj_str = str(proj_attr)
+                if proj_str.startswith("projects/"):
+                    project_number = proj_str.split("projects/", 1)[1]
+                elif proj_str:
+                    project_number = proj_str
+
+        loc_attr = getattr(bucket, "location", None)
+        location = (
+            loc_attr.lower() if isinstance(loc_attr, str) and loc_attr else "global"
+        )
+        loc_type_attr = getattr(bucket, "location_type", None)
+        location_type = (
+            loc_type_attr.lower()
+            if isinstance(loc_type_attr, str) and loc_type_attr
+            else "region"
+        )
 
         if location_type in ("multi-region", "dual-region"):
             location = "global"
 
-        if project_number:
-            destination_id = f"//storage.googleapis.com/projects/{project_number}/buckets/{bucket.name}"
-        else:
+        if project_number and str(project_number) != "_":
             destination_id = (
-                f"//storage.googleapis.com/projects/_/buckets/{bucket.name}"
+                f"//storage.googleapis.com/projects/{project_number}/buckets/{name}"
             )
+        else:
+            destination_id = f"//storage.googleapis.com/projects/_/buckets/{name}"
 
-        self.update_cache(bucket.name, destination_id, location)
+        self.update_cache(name, destination_id, location)
 
     def update_cache(self, bucket_name, destination_id, location):
         """Thread-safely update or insert a cache entry with bounded size."""

packages/google-cloud-storage/google/cloud/storage/_helpers.py

@@ -32,18 +32,16 @@
 from google.auth import environment_vars
 from google.cloud.exceptions import NotFound
 
-from google.cloud.storage._opentelemetry_tracing import (
-    _is_bucket_metadata_disabled,
-)
-from google.cloud.storage._opentelemetry_tracing import (
-    create_trace_span as _base_create_trace_span,
-)
+from google.cloud.storage import _opentelemetry_tracing
 from google.cloud.storage.constants import _DEFAULT_TIMEOUT
 from google.cloud.storage.retry import (
     DEFAULT_RETRY,
     DEFAULT_RETRY_IF_METAGENERATION_SPECIFIED,
 )
 
+# Kept as a patch point for unit tests.
+_base_create_trace_span = None
+
 _logger = logging.getLogger(__name__)
 
 STORAGE_EMULATOR_ENV_VAR = "STORAGE_EMULATOR_HOST"  # Despite name, includes scheme.
@@ -149,61 +147,127 @@ def _validate_name(name):
     return name
 
 
-@contextmanager
-def create_trace_span_helper(client, bucket_name, name, attributes=None, **kwargs):
-    span_attrs = dict(attributes) if attributes else {}
+_BUCKET_METADATA_FETCH_SPAN_NAMES = frozenset(
+    (
+        "Storage.Client.getBucket",
+        "Storage.Client.lookupBucket",
+        "Storage.Bucket.reload",
+        "Storage.Bucket.exists",
+    )
+)
 
-    if (
-        bucket_name
-        and isinstance(bucket_name, str)
-        and client
-        and hasattr(client, "_bucket_metadata_cache")
-        and client._bucket_metadata_cache
-        and not _is_bucket_metadata_disabled()
-    ):
-        try:
-            if name in (
-                "Storage.Client.getBucket",
-                "Storage.Client.lookupBucket",
-                "Storage.Bucket.reload",
-                "Storage.Bucket.exists",
-            ):
-                cached = client._bucket_metadata_cache.get(bucket_name)
-            else:
-                cached = client._bucket_metadata_cache.get_or_queue_fetch(bucket_name)
-
-            if cached and isinstance(cached, tuple) and len(cached) == 2:
-                dest_id, loc = cached
-                span_attrs.update(
-                    {
-                        "gcp.resource.destination.id": dest_id,
-                        "gcp.resource.destination.location": loc,
-                    }
-                )
-        except Exception as e:
-            _logger.debug(f"Failed cache lookup in create_trace_span_helper: {e}")
 
-    if "client" not in kwargs and client:
-        kwargs["client"] = client
+class _TraceSpanHelperContext:
+    """Context manager supporting both sync and async tracing span creation with bucket metadata."""
 
-    with _base_create_trace_span(name, attributes=span_attrs, **kwargs) as span:
-        try:
-            yield span
-        except (NotFound, api_exceptions.NotFound):
-            if (
-                bucket_name
-                and isinstance(bucket_name, str)
-                and client
-                and hasattr(client, "_bucket_metadata_cache")
-                and client._bucket_metadata_cache
-            ):
+    __slots__ = (
+        "client",
+        "bucket_name",
+        "name",
+        "attributes",
+        "kwargs",
+        "_base_cm",
+    )
+
+    def __init__(self, client, bucket_name, name, attributes=None, **kwargs):
+        self.client = client
+        self.bucket_name = bucket_name
+        self.name = name
+        self.attributes = attributes
+        self.kwargs = kwargs
+        self._base_cm = None
+
+    def _prepare_base_cm(self):
+        span_attrs = dict(self.attributes) if self.attributes else {}
+        client = self.client
+        bucket_name = self.bucket_name
+        name = self.name
+
+        if (
+            _opentelemetry_tracing._is_otel_traces_enabled()
+            and not _opentelemetry_tracing._is_bucket_metadata_disabled()
+            and bucket_name
+            and isinstance(bucket_name, str)
+            and client is not None
+        ):
+            cache = getattr(client, "_bucket_metadata_cache", None)
+            if cache is not None:
                 try:
-                    client._bucket_metadata_cache.check_and_evict(bucket_name)
+                    if name in _BUCKET_METADATA_FETCH_SPAN_NAMES:
+                        cached = cache.get(bucket_name)
+                    else:
+                        cached = cache.get_or_queue_fetch(bucket_name)
+
+                    if cached and isinstance(cached, tuple) and len(cached) == 2:
+                        dest_id, loc = cached
+                        span_attrs["gcp.resource.destination.id"] = dest_id
+                        span_attrs["gcp.resource.destination.location"] = loc
+                except Exception as e:
+                    _logger.debug(
+                        f"Failed cache lookup in create_trace_span_helper: {e}"
+                    )
+
+        kwargs = self.kwargs
+        if "client" not in kwargs and client is not None:
+            kwargs = dict(kwargs)
+            kwargs["client"] = client
+
+        create_span_fn = (
+            _base_create_trace_span
+            if _base_create_trace_span is not None
+            else _opentelemetry_tracing.create_trace_span
+        )
+        self._base_cm = create_span_fn(name, attributes=span_attrs, **kwargs)
+        return self._base_cm
+
+    def _handle_not_found(self):
+        if (
+            _opentelemetry_tracing._is_otel_traces_enabled()
+            and self.bucket_name
+            and isinstance(self.bucket_name, str)
+            and self.client is not None
+        ):
+            cache = getattr(self.client, "_bucket_metadata_cache", None)
+            if cache is not None:
+                try:
+                    cache.check_and_evict(self.bucket_name)
                 except Exception as e:
                     _logger.debug(
                         f"Failed cache eviction on 404 in create_trace_span_helper: {e}"
                     )
-            raise
+
+    def __enter__(self):
+        self._prepare_base_cm()
+        return self._base_cm.__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        suppressed = False
+        if self._base_cm is not None:
+            suppressed = bool(self._base_cm.__exit__(exc_type, exc_val, exc_tb))
+        if (
+            not suppressed
+            and exc_val is not None
+            and isinstance(exc_val, (NotFound, api_exceptions.NotFound))
+        ):
+            self._handle_not_found()
+        return suppressed
+
+    async def __aenter__(self):
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return self.__exit__(exc_type, exc_val, exc_tb)
+
+
+def create_trace_span_helper(client, bucket_name, name, attributes=None, **kwargs):
+    if (
+        _base_create_trace_span is None
+        and not _opentelemetry_tracing._is_otel_traces_enabled()
+    ):
+        return _opentelemetry_tracing._NOOP_TRACE_SPAN_CONTEXT
+    return _TraceSpanHelperContext(
+        client, bucket_name, name, attributes=attributes, **kwargs
+    )
 
 
 class _PropertyMixin(object):