diff --git a/packages/google-cloud-storage/google/cloud/storage/_bucket_metadata_cache.py b/packages/google-cloud-storage/google/cloud/storage/_bucket_metadata_cache.py index e431d00144fb..e375e4396142 100644 --- a/packages/google-cloud-storage/google/cloud/storage/_bucket_metadata_cache.py +++ b/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.""" diff --git a/packages/google-cloud-storage/google/cloud/storage/_helpers.py b/packages/google-cloud-storage/google/cloud/storage/_helpers.py index 9bea6c40845e..100d2def51d5 100644 --- a/packages/google-cloud-storage/google/cloud/storage/_helpers.py +++ b/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): diff --git a/packages/google-cloud-storage/google/cloud/storage/_opentelemetry_tracing.py b/packages/google-cloud-storage/google/cloud/storage/_opentelemetry_tracing.py index 1d9e4b88270b..75fa21b83ab2 100644 --- a/packages/google-cloud-storage/google/cloud/storage/_opentelemetry_tracing.py +++ b/packages/google-cloud-storage/google/cloud/storage/_opentelemetry_tracing.py @@ -16,10 +16,8 @@ import logging import os -from contextlib import contextmanager from urllib.parse import urlparse -from google.api_core import exceptions as api_exceptions from google.api_core import retry as api_retry from google.cloud.storage import __version__ @@ -44,6 +42,12 @@ def _is_bucket_metadata_disabled() -> bool: enable_otel_traces = _parse_bool_env( ENABLE_OTEL_TRACES_ENV_VAR, _DEFAULT_ENABLE_OTEL_TRACES_VALUE ) + + +def _is_otel_traces_enabled() -> bool: + return bool(HAS_OPENTELEMETRY and enable_otel_traces) + + logger = logging.getLogger(__name__) @@ -75,30 +79,116 @@ def _is_bucket_metadata_disabled() -> bool: } -@contextmanager -def create_trace_span(name, attributes=None, client=None, api_request=None, retry=None): +class _NoOpTraceSpanContext: + """Singleton no-op context manager supporting both sync and async protocols.""" + + __slots__ = () + + def __enter__(self): + return None + + def __exit__(self, exc_type, exc_val, exc_tb): + return False + + async def __aenter__(self): + return None + + async def __aexit__(self, exc_type, exc_val, exc_tb): + return False + + +_NOOP_TRACE_SPAN_CONTEXT = _NoOpTraceSpanContext() + + +class _TraceSpanContext: + """Context manager supporting both sync and async tracing spans.""" + + __slots__ = ( + "name", + "attributes", + "client", + "api_request", + "retry", + "rpc_system", + "_span_cm", + ) + + def __init__( + self, + name, + attributes=None, + client=None, + api_request=None, + retry=None, + rpc_system="http", + ): + self.name = name + self.attributes = attributes + self.client = client + self.api_request = api_request + self.retry = retry + self.rpc_system = rpc_system + self._span_cm = None + + def __enter__(self): + if not _is_otel_traces_enabled(): + return None + + tracer = trace.get_tracer(__name__) + final_attributes = _get_final_attributes( + self.attributes, + self.client, + self.api_request, + self.retry, + rpc_system=self.rpc_system, + ) + self._span_cm = tracer.start_as_current_span( + name=self.name, + kind=trace.SpanKind.CLIENT, + attributes=final_attributes, + ) + return self._span_cm.__enter__() + + def __exit__(self, exc_type, exc_val, exc_tb): + if self._span_cm is not None: + return self._span_cm.__exit__(exc_type, exc_val, exc_tb) + return False + + 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( + name, + attributes=None, + client=None, + api_request=None, + retry=None, + rpc_system="http", +): """Creates a context manager for a new span and set it as the current span - in the configured tracer. If no configuration exists yields None.""" - if not HAS_OPENTELEMETRY or not enable_otel_traces: - yield None - return - - tracer = trace.get_tracer(__name__) - final_attributes = _get_final_attributes(attributes, client, api_request, retry) - # Yield new span. - with tracer.start_as_current_span( - name=name, kind=trace.SpanKind.CLIENT, attributes=final_attributes - ) as span: - try: - yield span - except api_exceptions.GoogleAPICallError as error: - span.set_status(trace.Status(trace.StatusCode.ERROR)) - span.record_exception(error) - raise - - -def _get_final_attributes(attributes=None, client=None, api_request=None, retry=None): + in the configured tracer. Supports both sync and async context managers. + If no configuration exists yields None.""" + if not _is_otel_traces_enabled(): + return _NOOP_TRACE_SPAN_CONTEXT + return _TraceSpanContext( + name=name, + attributes=attributes, + client=client, + api_request=api_request, + retry=retry, + rpc_system=rpc_system, + ) + + +def _get_final_attributes( + attributes=None, client=None, api_request=None, retry=None, rpc_system="http" +): collected_attr = _default_attributes.copy() + collected_attr["rpc.system"] = rpc_system collected_attr.update(_cloud_trace_adoption_attrs) if api_request: collected_attr.update(_set_api_request_attr(api_request, client)) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/_utils.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/_utils.py index 53e8503c7f8b..b917c55be60c 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/_utils.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/_utils.py @@ -12,9 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextvars +from typing import Dict, FrozenSet, Optional, Union + import google_crc32c from google.api_core import exceptions +from google.cloud.storage import _opentelemetry_tracing +from google.cloud.storage._helpers import create_trace_span_helper + def raise_if_no_fast_crc32c(): """Check if the C-accelerated version of google-crc32c is available. @@ -38,3 +44,97 @@ def update_write_handle_if_exists(obj, response): """Update the write_handle attribute of an object if it exists in the response.""" if hasattr(response, "write_handle") and response.write_handle is not None: obj.write_handle = response.write_handle + + +try: + from opentelemetry.trace.propagation.tracecontext import ( + TraceContextTextMapPropagator, + ) + + _TRACE_CONTEXT_PROPAGATOR = TraceContextTextMapPropagator() +except ImportError: # pragma: NO COVER + _TRACE_CONTEXT_PROPAGATOR = None + + +def inject_traceparent_to_metadata(metadata=None): + """Inject W3C traceparent (and tracestate) into gRPC metadata tuple or list.""" + if ( + not _opentelemetry_tracing._is_otel_traces_enabled() + or _TRACE_CONTEXT_PROPAGATOR is None + ): + return () if metadata is None else metadata + + carrier = {} + try: + _TRACE_CONTEXT_PROPAGATOR.inject(carrier) + except Exception: + return () if metadata is None else metadata + + if not carrier: + return () if metadata is None else metadata + + meta_list = list(metadata) if metadata else [] + existing_keys = {k.lower() for k, _ in meta_list} + for key, val in carrier.items(): + if key.lower() not in existing_keys: + meta_list.append((key, val)) + + return ( + tuple(meta_list) + if isinstance(metadata, tuple) or metadata is None + else meta_list + ) + + +_ACTIVE_T2_SPAN_OWNERS: contextvars.ContextVar[FrozenSet[int]] = contextvars.ContextVar( + "_active_t2_span_owners", default=frozenset() +) + + +class _TopLevelSpanContext: + """Async context manager that tracks active top-level T2 spans per task.""" + + __slots__ = ("_owner_id", "_cm", "_token") + + def __init__(self, owner_id: int, cm): + self._owner_id = owner_id + self._cm = cm + self._token: Optional[contextvars.Token] = None + + async def __aenter__(self): + owners = _ACTIVE_T2_SPAN_OWNERS.get() + self._token = _ACTIVE_T2_SPAN_OWNERS.set(owners | {self._owner_id}) + try: + return await self._cm.__aenter__() + except BaseException: + _ACTIVE_T2_SPAN_OWNERS.reset(self._token) + raise + + async def __aexit__(self, exc_type, exc_val, exc_tb): + try: + return await self._cm.__aexit__(exc_type, exc_val, exc_tb) + finally: + _ACTIVE_T2_SPAN_OWNERS.reset(self._token) + + +def create_top_level_trace_span( + owner, + client, + bucket_name: str, + name: str, + attributes: Optional[Dict[str, Union[str, int]]] = None, +): + """Create a T2 OpenTelemetry span only for top-level public API calls per task.""" + owner_id = id(owner) + if owner_id in _ACTIVE_T2_SPAN_OWNERS.get(): + return _opentelemetry_tracing._NOOP_TRACE_SPAN_CONTEXT + cm = create_trace_span_helper( + client, + bucket_name, + name, + attributes=attributes, + rpc_system="grpc", + ) + if cm is _opentelemetry_tracing._NOOP_TRACE_SPAN_CONTEXT: + return cm + return _TopLevelSpanContext(owner_id, cm) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index da399659a343..b0081278cbc5 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -316,6 +316,20 @@ def __init__( self._flush_count = 0 self.blob: Optional[Blob] = None + def _trace_span( + self, + name: str, + attributes: Optional[Dict[str, Union[str, int]]] = None, + ): + """Creates a T2 OpenTelemetry span only for top-level public API calls.""" + return _utils.create_top_level_trace_span( + self, + self.client, + self.bucket_name, + name, + attributes=attributes, + ) + @classmethod def from_blob( cls, @@ -403,14 +417,15 @@ async def state_lookup(self) -> int: if not self._is_stream_open: raise ValueError("Stream is not open. Call open() before state_lookup().") - await self.write_obj_stream.send( - _storage_v2.BidiWriteObjectRequest( - state_lookup=True, + async with self._trace_span("Storage.AsyncAppendableObjectWriter.stateLookup"): + await self.write_obj_stream.send( + _storage_v2.BidiWriteObjectRequest( + state_lookup=True, + ) ) - ) - response = await self.write_obj_stream.recv() - self.persisted_size = response.persisted_size - return self.persisted_size + response = await self.write_obj_stream.recv() + self.persisted_size = response.persisted_size + return self.persisted_size def _on_open_error(self, exc): """Extracts routing token and write handle on redirect error during open.""" @@ -465,60 +480,66 @@ async def open( if self._is_stream_open: raise ValueError("Underlying bidi-gRPC stream is already open") - retry_policy = self._merge_retry_policy(retry_policy) + async with self._trace_span("Storage.AsyncAppendableObjectWriter.open"): + retry_policy = self._merge_retry_policy(retry_policy) + + async def _do_open(): + current_metadata = list(metadata) if metadata else [] + creating = not self.generation + + # Cleanup stream from previous failed attempt, if any. + if self.write_obj_stream: + if self.write_obj_stream.is_stream_open: + try: + await self.write_obj_stream.close() + except Exception as e: + logger.warning( + f"Error closing previous write stream during open retry. Got exception: {e}" + ) + self.write_obj_stream = None + self._is_stream_open = False - async def _do_open(): - current_metadata = list(metadata) if metadata else [] - creating = not self.generation + self.write_obj_stream = _AsyncWriteObjectStream( + client=self.client.grpc_client, + bucket_name=self.bucket_name, + object_name=self.object_name, + blob=self.blob, + generation_number=self.generation, + write_handle=self.write_handle, + routing_token=self._routing_token, + storage_class=self.storage_class, + preconditions=( + None if self._created_generation else self.preconditions + ), + ) - # Cleanup stream from previous failed attempt, if any. - if self.write_obj_stream: - if self.write_obj_stream.is_stream_open: - try: - await self.write_obj_stream.close() - except Exception as e: - logger.warning( - f"Error closing previous write stream during open retry. Got exception: {e}" + if self._routing_token: + current_metadata.append( + ( + "x-goog-request-params", + f"routing_token={self._routing_token}", ) - self.write_obj_stream = None - self._is_stream_open = False - - self.write_obj_stream = _AsyncWriteObjectStream( - client=self.client.grpc_client, - bucket_name=self.bucket_name, - object_name=self.object_name, - blob=self.blob, - generation_number=self.generation, - write_handle=self.write_handle, - routing_token=self._routing_token, - storage_class=self.storage_class, - preconditions=None if self._created_generation else self.preconditions, - ) + ) - if self._routing_token: - current_metadata.append( - ("x-goog-request-params", f"routing_token={self._routing_token}") + await self.write_obj_stream.open( + metadata=current_metadata if current_metadata else None ) - await self.write_obj_stream.open( - metadata=current_metadata if current_metadata else None - ) - - if self.write_obj_stream.generation_number: - if creating: - self._created_generation = True - self.generation = self.write_obj_stream.generation_number - if self.write_obj_stream.write_handle: - self.write_handle = self.write_obj_stream.write_handle - if self.write_obj_stream.persisted_size is not None: - self.persisted_size = self.write_obj_stream.persisted_size - # set offset while opening - self.offset = self.persisted_size + if self.write_obj_stream.generation_number: + if creating: + self._created_generation = True + self.generation = self.write_obj_stream.generation_number + if self.write_obj_stream.write_handle: + self.write_handle = self.write_obj_stream.write_handle + if self.write_obj_stream.persisted_size is not None: + self.persisted_size = self.write_obj_stream.persisted_size + # set offset while opening + self.offset = self.persisted_size - self._is_stream_open = True - self._routing_token = None + self._is_stream_open = True + self._routing_token = None - await retry_policy(_do_open)() + await retry_policy(_do_open)() async def append( self, @@ -557,100 +578,106 @@ async def append( logger.debug("No data provided to append; returning without action.") return - if retry_policy is None: - retry_policy = AsyncRetry(predicate=_is_write_retryable) - - strategy = _WriteResumptionStrategy() - buffer = io.BytesIO(data) - attempt_count = 0 - - def send_and_recv_generator( - requests: List[BidiWriteObjectRequest], - state: Dict[str, _WriteState], - metadata: Optional[List[Tuple[str, str]]] = None, + async with self._trace_span( + "Storage.AsyncAppendableObjectWriter.append", + attributes={"gcp.storage.payload.size_bytes": len(data)}, ): - async def generator(): - nonlocal attempt_count - nonlocal requests - attempt_count += 1 - resp = None - write_state = state["write_state"] - # If this is a retry or redirect, we must re-open the stream - if attempt_count > 1 or write_state.routing_token: - logger.info( - f"Re-opening the stream with attempt_count: {attempt_count}" - ) - - current_metadata = list(metadata) if metadata else [] - if write_state.routing_token: - current_metadata.append( - ( - "x-goog-request-params", - f"routing_token={write_state.routing_token}", - ) + if retry_policy is None: + retry_policy = AsyncRetry(predicate=_is_write_retryable) + + strategy = _WriteResumptionStrategy() + buffer = io.BytesIO(data) + attempt_count = 0 + + def send_and_recv_generator( + requests: List[BidiWriteObjectRequest], + state: Dict[str, _WriteState], + metadata: Optional[List[Tuple[str, str]]] = None, + ): + async def generator(): + nonlocal attempt_count + nonlocal requests + attempt_count += 1 + resp = None + write_state = state["write_state"] + # If this is a retry or redirect, we must re-open the stream + if attempt_count > 1 or write_state.routing_token: + logger.info( + f"Re-opening the stream with attempt_count: {attempt_count}" ) - self._routing_token = write_state.routing_token - - self._is_stream_open = False - await self.open(metadata=current_metadata) - - write_state.persisted_size = self.persisted_size - write_state.write_handle = self.write_handle - write_state.routing_token = None - - write_state.user_buffer.seek(write_state.persisted_size) - write_state.bytes_sent = write_state.persisted_size - write_state.bytes_since_last_flush = 0 - self.bytes_appended_since_last_flush = 0 - requests = strategy.generate_requests(state) - - for chunk_req in requests: - await self.write_obj_stream.send(chunk_req) - if chunk_req.flush: - self._flush_count += 1 - - resp = None - if chunk_req.state_lookup: - # TODO: if there's error, it'll raise error - # and will be handled by `recover_state_on_failure` - resp = await self.write_obj_stream.recv() - - if resp: - if resp.persisted_size is not None: - self.persisted_size = resp.persisted_size - state["write_state"].persisted_size = resp.persisted_size - self.offset = self.persisted_size - if resp.write_handle: - self.write_handle = resp.write_handle - state["write_state"].write_handle = resp.write_handle - - yield resp - - return generator() - - # State initialization - write_state = _WriteState( - _MAX_CHUNK_SIZE_BYTES, - buffer, - self.flush_interval, - enable_checksum=enable_checksum, - ) - write_state.write_handle = self.write_handle - write_state.persisted_size = self.persisted_size - # offset is set during `open()` call. - write_state.bytes_sent = self.offset or 0 - write_state.bytes_since_last_flush = self.bytes_appended_since_last_flush - - retry_manager = _BidiStreamRetryManager( - _WriteResumptionStrategy(), - lambda r, s: send_and_recv_generator(r, s, metadata), - ) - await retry_manager.execute({"write_state": write_state}, retry_policy) + current_metadata = list(metadata) if metadata else [] + if write_state.routing_token: + current_metadata.append( + ( + "x-goog-request-params", + f"routing_token={write_state.routing_token}", + ) + ) + self._routing_token = write_state.routing_token + + self._is_stream_open = False + await self.open(metadata=current_metadata) + + write_state.persisted_size = self.persisted_size + write_state.write_handle = self.write_handle + write_state.routing_token = None + + write_state.user_buffer.seek(write_state.persisted_size) + write_state.bytes_sent = write_state.persisted_size + write_state.bytes_since_last_flush = 0 + self.bytes_appended_since_last_flush = 0 + + requests = strategy.generate_requests(state) + + for chunk_req in requests: + await self.write_obj_stream.send(chunk_req) + if chunk_req.flush: + self._flush_count += 1 + + resp = None + if chunk_req.state_lookup: + # TODO: if there's error, it'll raise error + # and will be handled by `recover_state_on_failure` + resp = await self.write_obj_stream.recv() + + if resp: + if resp.persisted_size is not None: + self.persisted_size = resp.persisted_size + state[ + "write_state" + ].persisted_size = resp.persisted_size + self.offset = self.persisted_size + if resp.write_handle: + self.write_handle = resp.write_handle + state["write_state"].write_handle = resp.write_handle + + yield resp + + return generator() + + # State initialization + write_state = _WriteState( + _MAX_CHUNK_SIZE_BYTES, + buffer, + self.flush_interval, + enable_checksum=enable_checksum, + ) + write_state.write_handle = self.write_handle + write_state.persisted_size = self.persisted_size + # offset is set during `open()` call. + write_state.bytes_sent = self.offset or 0 + write_state.bytes_since_last_flush = self.bytes_appended_since_last_flush + + retry_manager = _BidiStreamRetryManager( + _WriteResumptionStrategy(), + lambda r, s: send_and_recv_generator(r, s, metadata), + ) + await retry_manager.execute({"write_state": write_state}, retry_policy) - # Sync local markers - self.bytes_appended_since_last_flush = write_state.bytes_since_last_flush - self.offset = write_state.bytes_sent + # Sync local markers + self.bytes_appended_since_last_flush = write_state.bytes_since_last_flush + self.offset = write_state.bytes_sent async def simple_flush(self) -> None: """Flushes the data to the server. @@ -664,12 +691,13 @@ async def simple_flush(self) -> None: if not self._is_stream_open: raise ValueError("Stream is not open. Call open() before simple_flush().") - await self.write_obj_stream.send( - _storage_v2.BidiWriteObjectRequest( - flush=True, + async with self._trace_span("Storage.AsyncAppendableObjectWriter.simpleFlush"): + await self.write_obj_stream.send( + _storage_v2.BidiWriteObjectRequest( + flush=True, + ) ) - ) - self.bytes_appended_since_last_flush = 0 + self.bytes_appended_since_last_flush = 0 async def flush(self) -> int: """Flushes the data to the server. @@ -683,17 +711,18 @@ async def flush(self) -> int: if not self._is_stream_open: raise ValueError("Stream is not open. Call open() before flush().") - await self.write_obj_stream.send( - _storage_v2.BidiWriteObjectRequest( - flush=True, - state_lookup=True, + async with self._trace_span("Storage.AsyncAppendableObjectWriter.flush"): + await self.write_obj_stream.send( + _storage_v2.BidiWriteObjectRequest( + flush=True, + state_lookup=True, + ) ) - ) - response = await self.write_obj_stream.recv() - self.persisted_size = response.persisted_size - self.offset = self.persisted_size - self.bytes_appended_since_last_flush = 0 - return self.persisted_size + response = await self.write_obj_stream.recv() + self.persisted_size = response.persisted_size + self.offset = self.persisted_size + self.bytes_appended_since_last_flush = 0 + return self.persisted_size async def close( self, @@ -746,43 +775,44 @@ async def close( "full_object_checksum can only be provided when finalize_on_close is True." ) - if finalize_on_close: - return await self.finalize( - full_object_checksum=full_object_checksum, - retry_policy=retry_policy, - ) + async with self._trace_span("Storage.AsyncAppendableObjectWriter.close"): + if finalize_on_close: + return await self.finalize( + full_object_checksum=full_object_checksum, + retry_policy=retry_policy, + ) - retry_policy = self._merge_retry_policy(retry_policy) + retry_policy = self._merge_retry_policy(retry_policy) - attempt_count = 0 - expected_offset = self.offset + attempt_count = 0 + expected_offset = self.offset - async def _do_close(): - nonlocal attempt_count - attempt_count += 1 + async def _do_close(): + nonlocal attempt_count + attempt_count += 1 - if attempt_count > 1: - logger.info( - f"Re-opening the stream for close retry attempt: {attempt_count}" - ) - self._is_stream_open = False - await self.open() - if ( - self.offset is not None - and expected_offset is not None - and self.offset != expected_offset - ): - raise exceptions.InternalServerError( - f"Unrecoverable data loss during reconnect. Expected offset {expected_offset}, got {self.offset}" + if attempt_count > 1: + logger.info( + f"Re-opening the stream for close retry attempt: {attempt_count}" ) + self._is_stream_open = False + await self.open() + if ( + self.offset is not None + and expected_offset is not None + and self.offset != expected_offset + ): + raise exceptions.InternalServerError( + f"Unrecoverable data loss during reconnect. Expected offset {expected_offset}, got {self.offset}" + ) - await self.write_obj_stream.close() - return self.persisted_size + await self.write_obj_stream.close() + return self.persisted_size - try: - return await retry_policy(_do_close)() - finally: - self._is_stream_open = False + try: + return await retry_policy(_do_close)() + finally: + self._is_stream_open = False async def finalize( self, @@ -844,48 +874,49 @@ async def finalize( ), ) - retry_policy = self._merge_retry_policy(retry_policy) + async with self._trace_span("Storage.AsyncAppendableObjectWriter.finalize"): + retry_policy = self._merge_retry_policy(retry_policy) - attempt_count = 0 - expected_offset = self.offset + attempt_count = 0 + expected_offset = self.offset - async def _do_finalize(): - nonlocal attempt_count - attempt_count += 1 + async def _do_finalize(): + nonlocal attempt_count + attempt_count += 1 - if attempt_count > 1: - logger.info( - f"Re-opening the stream for finalize retry attempt: {attempt_count}" - ) - self._is_stream_open = False - await self.open() - if ( - self.offset is not None - and expected_offset is not None - and self.offset != expected_offset - ): - raise exceptions.InternalServerError( - f"Unrecoverable data loss during reconnect. Expected offset {expected_offset}, got {self.offset}" + if attempt_count > 1: + logger.info( + f"Re-opening the stream for finalize retry attempt: {attempt_count}" ) + self._is_stream_open = False + await self.open() + if ( + self.offset is not None + and expected_offset is not None + and self.offset != expected_offset + ): + raise exceptions.InternalServerError( + f"Unrecoverable data loss during reconnect. Expected offset {expected_offset}, got {self.offset}" + ) - await self.write_obj_stream.send(finalize_req) - response = await self.write_obj_stream.recv() - self.object_resource = response.resource - self.persisted_size = self.object_resource.size - return self.object_resource - - try: - return await retry_policy(_do_finalize)() - finally: - if self.write_obj_stream: - try: - await self.write_obj_stream.close() - except Exception as e: - logger.debug( - f"Stream close during finalize cleanup resulted in: {e}" - ) - self._is_stream_open = False - self.offset = None + await self.write_obj_stream.send(finalize_req) + response = await self.write_obj_stream.recv() + self.object_resource = response.resource + self.persisted_size = self.object_resource.size + return self.object_resource + + try: + return await retry_policy(_do_finalize)() + finally: + if self.write_obj_stream: + try: + await self.write_obj_stream.close() + except Exception as e: + logger.debug( + f"Stream close during finalize cleanup resulted in: {e}" + ) + self._is_stream_open = False + self.offset = None @property def is_stream_open(self) -> bool: @@ -918,5 +949,8 @@ async def append_from_file( :param file_obj: A file handle opened in binary mode for reading. """ - while block := file_obj.read(block_size): - await self.append(block) + async with self._trace_span( + "Storage.AsyncAppendableObjectWriter.appendFromFile" + ): + while block := file_obj.read(block_size): + await self.append(block) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py index 21278cd624bf..e113943857a8 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py @@ -22,6 +22,9 @@ DEFAULT_CLIENT_INFO, ) from google.cloud.storage import __version__, _opentelemetry_metrics +from google.cloud.storage._bucket_metadata_cache import BucketMetadataCache +from google.cloud.storage._helpers import create_trace_span_helper +from google.cloud.storage.asyncio import _utils from google.cloud.storage.grpc_client import ( _DEFAULT_HOST, _resolve_direct_path_interconnect, @@ -107,6 +110,8 @@ def __init__( self._debug_metrics_enabled = _opentelemetry_metrics.is_debug_metrics_enabled( enable_debug_metrics ) + self._is_async_grpc_client = True + self._bucket_metadata_cache = BucketMetadataCache(self) if isinstance(credentials, auth_credentials.AnonymousCredentials): if client_options is None or client_options.api_endpoint is None: @@ -277,24 +282,31 @@ async def delete_object( :param retry: (Optional) Designation of what errors, if any, should be retried. """ _validate_metadata(metadata) - # The gRPC API requires the bucket name to be in the format "projects/_/buckets/bucket_name" - bucket_path = f"projects/_/buckets/{bucket_name}" - request = storage_v2.DeleteObjectRequest( - bucket=bucket_path, - object=object_name, - generation=generation, - if_generation_match=if_generation_match, - if_generation_not_match=if_generation_not_match, - if_metageneration_match=if_metageneration_match, - if_metageneration_not_match=if_metageneration_not_match, - **kwargs, - ) - await self._grpc_client.delete_object( - request=request, - metadata=metadata, - timeout=timeout, - retry=retry, - ) + async with create_trace_span_helper( + self, + bucket_name, + "Storage.AsyncGrpcClient.deleteObject", + rpc_system="grpc", + ): + # The gRPC API requires the bucket name to be in the format "projects/_/buckets/bucket_name" + bucket_path = f"projects/_/buckets/{bucket_name}" + request = storage_v2.DeleteObjectRequest( + bucket=bucket_path, + object=object_name, + generation=generation, + if_generation_match=if_generation_match, + if_generation_not_match=if_generation_not_match, + if_metageneration_match=if_metageneration_match, + if_metageneration_not_match=if_metageneration_not_match, + **kwargs, + ) + final_metadata = _utils.inject_traceparent_to_metadata(metadata) + await self._grpc_client.delete_object( + request=request, + metadata=final_metadata, + timeout=timeout, + retry=retry, + ) async def get_object( self, @@ -357,24 +369,31 @@ async def get_object( :returns: The object metadata resource. """ _validate_metadata(metadata) - bucket_path = f"projects/_/buckets/{bucket_name}" - - request = storage_v2.GetObjectRequest( - bucket=bucket_path, - object=object_name, - generation=generation, - if_generation_match=if_generation_match, - if_generation_not_match=if_generation_not_match, - if_metageneration_match=if_metageneration_match, - if_metageneration_not_match=if_metageneration_not_match, - soft_deleted=soft_deleted or False, - **kwargs, - ) + async with create_trace_span_helper( + self, + bucket_name, + "Storage.AsyncGrpcClient.getObject", + rpc_system="grpc", + ): + bucket_path = f"projects/_/buckets/{bucket_name}" + + request = storage_v2.GetObjectRequest( + bucket=bucket_path, + object=object_name, + generation=generation, + if_generation_match=if_generation_match, + if_generation_not_match=if_generation_not_match, + if_metageneration_match=if_metageneration_match, + if_metageneration_not_match=if_metageneration_not_match, + soft_deleted=soft_deleted or False, + **kwargs, + ) - # Calls the underlying GAPIC StorageAsyncClient.get_object method - return await self._grpc_client.get_object( - request=request, - metadata=metadata, - timeout=timeout, - retry=retry, - ) + final_metadata = _utils.inject_traceparent_to_metadata(metadata) + # Calls the underlying GAPIC StorageAsyncClient.get_object method + return await self._grpc_client.get_object( + request=request, + metadata=final_metadata, + timeout=timeout, + retry=retry, + ) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py index 697926d3f4e2..74b9581333d1 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py @@ -45,7 +45,7 @@ ) from google.cloud.storage.exceptions import DataCorruption -from ._utils import raise_if_no_fast_crc32c +from ._utils import create_top_level_trace_span, raise_if_no_fast_crc32c _MAX_READ_RANGES_PER_BIDI_READ_REQUEST = 100 _COMMON_READ_RETRYABLE_EXCEPTIONS = ( @@ -271,6 +271,20 @@ def __init__( self.is_finalized: bool = False self.full_obj_server_crc32c: Optional[int] = None + def _trace_span( + self, + name: str, + attributes: Optional[Dict[str, Union[str, int]]] = None, + ): + """Creates a T2 OpenTelemetry span only for top-level public API calls.""" + return create_top_level_trace_span( + self, + self.client, + self.bucket_name, + name, + attributes=attributes, + ) + async def __aenter__(self): """Opens the underlying bidi-gRPC connection to read from the object.""" await self.open() @@ -299,79 +313,83 @@ async def open( if self._is_stream_open: raise ValueError("Underlying bidi-gRPC stream is already open") - if retry_policy is None: + async with self._trace_span("Storage.AsyncMultiRangeDownloader.open"): + if retry_policy is None: - def on_error_wrapper(exc): - self._open_retries += 1 - self._on_open_error(exc) + def on_error_wrapper(exc): + self._open_retries += 1 + self._on_open_error(exc) - retry_policy = AsyncRetry( - predicate=_is_open_retryable, on_error=on_error_wrapper - ) - else: - original_on_error = retry_policy._on_error - - def combined_on_error(exc): - self._open_retries += 1 - self._on_open_error(exc) - if original_on_error: - original_on_error(exc) - - retry_policy = AsyncRetry( - predicate=_is_open_retryable, - initial=retry_policy._initial, - maximum=retry_policy._maximum, - multiplier=retry_policy._multiplier, - deadline=retry_policy._deadline, - on_error=combined_on_error, - ) + retry_policy = AsyncRetry( + predicate=_is_open_retryable, on_error=on_error_wrapper + ) + else: + original_on_error = retry_policy._on_error + + def combined_on_error(exc): + self._open_retries += 1 + self._on_open_error(exc) + if original_on_error: + original_on_error(exc) + + retry_policy = AsyncRetry( + predicate=_is_open_retryable, + initial=retry_policy._initial, + maximum=retry_policy._maximum, + multiplier=retry_policy._multiplier, + deadline=retry_policy._deadline, + on_error=combined_on_error, + ) - async def _do_open(): - current_metadata = list(metadata) if metadata else [] + async def _do_open(): + current_metadata = list(metadata) if metadata else [] - # Cleanup stream from previous failed attempt, if any. - if self.read_obj_str: - if self.read_obj_str.is_stream_open: - try: - await self.read_obj_str.close() - except exceptions.GoogleAPICallError as e: - logger.warning( - f"Failed to close existing stream during resumption: {e}" - ) - self.read_obj_str = None - self._is_stream_open = False + # Cleanup stream from previous failed attempt, if any. + if self.read_obj_str: + if self.read_obj_str.is_stream_open: + try: + await self.read_obj_str.close() + except exceptions.GoogleAPICallError as e: + logger.warning( + f"Failed to close existing stream during resumption: {e}" + ) + self.read_obj_str = None + self._is_stream_open = False + + self.read_obj_str = _AsyncReadObjectStream( + client=self.client.grpc_client, + bucket_name=self.bucket_name, + object_name=self.object_name, + generation_number=self.generation, + read_handle=self.read_handle, + ) - self.read_obj_str = _AsyncReadObjectStream( - client=self.client.grpc_client, - bucket_name=self.bucket_name, - object_name=self.object_name, - generation_number=self.generation, - read_handle=self.read_handle, - ) + if self._routing_token: + current_metadata.append( + ( + "x-goog-request-params", + f"routing_token={self._routing_token}", + ) + ) + self._routing_token = None - if self._routing_token: - current_metadata.append( - ("x-goog-request-params", f"routing_token={self._routing_token}") + await self.read_obj_str.open( + metadata=current_metadata if current_metadata else None ) - self._routing_token = None - - await self.read_obj_str.open( - metadata=current_metadata if current_metadata else None - ) - if self.read_obj_str.generation_number: - self.generation = self.read_obj_str.generation_number - if self.read_obj_str.read_handle: - self.read_handle = self.read_obj_str.read_handle - if self.read_obj_str.persisted_size is not None: - self.persisted_size = self.read_obj_str.persisted_size - self.is_finalized = self.read_obj_str.is_finalized - self.full_obj_server_crc32c = self.read_obj_str.full_obj_server_crc32c + if self.read_obj_str.generation_number: + self.generation = self.read_obj_str.generation_number + if self.read_obj_str.read_handle: + self.read_handle = self.read_obj_str.read_handle + if self.read_obj_str.persisted_size is not None: + self.persisted_size = self.read_obj_str.persisted_size + self.is_finalized = self.read_obj_str.is_finalized + self.full_obj_server_crc32c = self.read_obj_str.full_obj_server_crc32c - self._is_stream_open = True + self._is_stream_open = True - await retry_policy(_do_open)() - self._multiplexer = _StreamMultiplexer(self.read_obj_str) + await retry_policy(_do_open)() + self._multiplexer = _StreamMultiplexer(self.read_obj_str) def _create_stream_factory(self, state, metadata): """Create a factory that opens a new stream with current routing state.""" @@ -513,179 +531,188 @@ async def download_ranges( if not self._is_stream_open: raise ValueError("Underlying bidi-gRPC stream is not open") - if retry_policy is None: - retry_policy = AsyncRetry(predicate=_is_read_retryable) + async with self._trace_span( + "Storage.AsyncMultiRangeDownloader.downloadRanges", + attributes={"gcp.storage.range.count": len(read_ranges)}, + ): + if retry_policy is None: + retry_policy = AsyncRetry(predicate=_is_read_retryable) - # Initialize Global State for Retry Strategy - download_states = {} - read_ids = set() - try: - for read_range in read_ranges: - read_id = generate_random_56_bit_integer() - # Unpack tuple into self-documenting variable names to improve readability. - offset, length, user_buffer = read_range - - # Heuristic to detect full object reads: - # - Implicit full object read: start offset is 0 and length is 0 (read all). - # - Explicit full object read: start offset is 0 and length matches the exact persisted size. - is_full_object_read = (offset == 0 and length == 0) or ( - self.persisted_size is not None - and offset == 0 - and length == self.persisted_size - ) - download_states[read_id] = _DownloadState( - initial_offset=offset, - initial_length=length, - user_buffer=_wrap_user_buffer( - user_buffer, length, accepts_memoryview - ), - is_full_object_read=is_full_object_read, - ) + # Initialize Global State for Retry Strategy + download_states = {} + read_ids = set() + try: + for read_range in read_ranges: + read_id = generate_random_56_bit_integer() + # Unpack tuple into self-documenting variable names to improve readability. + offset, length, user_buffer = read_range + + # Heuristic to detect full object reads: + # - Implicit full object read: start offset is 0 and length is 0 (read all). + # - Explicit full object read: start offset is 0 and length matches the exact persisted size. + is_full_object_read = (offset == 0 and length == 0) or ( + self.persisted_size is not None + and offset == 0 + and length == self.persisted_size + ) + download_states[read_id] = _DownloadState( + initial_offset=offset, + initial_length=length, + user_buffer=_wrap_user_buffer( + user_buffer, length, accepts_memoryview + ), + is_full_object_read=is_full_object_read, + ) - initial_state = { - "download_states": download_states, - "read_handle": self.read_handle, - "routing_token": None, - "enable_checksum": enable_checksum, - "full_obj_server_crc32c": self.full_obj_server_crc32c, - } - - read_ids = set(download_states.keys()) - queue = self._multiplexer.register(read_ids) - strategy = _ReadResumptionStrategy() - - attempt_count = 0 - last_broken_generation = None - - def send_and_recv_via_multiplexer( - requests: List[_storage_v2.ReadRange], - state: Dict[str, Any], - ): - async def generator(): - nonlocal attempt_count, last_broken_generation - attempt_count += 1 - - if attempt_count > 1: - logger.info( - f"Resuming download (attempt {attempt_count}) for {len(requests)} ranges." - ) + initial_state = { + "download_states": download_states, + "read_handle": self.read_handle, + "routing_token": None, + "enable_checksum": enable_checksum, + "full_obj_server_crc32c": self.full_obj_server_crc32c, + } - # Reopen stream if needed - should_reopen = ( - attempt_count > 1 and last_broken_generation is not None - ) or (attempt_count == 1 and metadata is not None) - if should_reopen: - broken_gen = ( - last_broken_generation - if attempt_count > 1 - else self._multiplexer.stream_generation - ) - stream_factory = self._create_stream_factory(state, metadata) - await self._multiplexer.reopen_stream( - broken_gen, stream_factory - ) + read_ids = set(download_states.keys()) + queue = self._multiplexer.register(read_ids) + strategy = _ReadResumptionStrategy() - stream_generation = self._multiplexer.stream_generation + attempt_count = 0 + last_broken_generation = None - # Send Requests - pending_read_ids = {r.read_id for r in requests} - for i in range( - 0, len(requests), _MAX_READ_RANGES_PER_BIDI_READ_REQUEST - ): - batch = requests[i : i + _MAX_READ_RANGES_PER_BIDI_READ_REQUEST] - try: - await self._multiplexer.send( - _storage_v2.BidiReadObjectRequest(read_ranges=batch) + def send_and_recv_via_multiplexer( + requests: List[_storage_v2.ReadRange], + state: Dict[str, Any], + ): + async def generator(): + nonlocal attempt_count, last_broken_generation + attempt_count += 1 + + if attempt_count > 1: + logger.info( + f"Resuming download (attempt {attempt_count}) for {len(requests)} ranges." ) - except Exception: - last_broken_generation = stream_generation - raise - # Receive Responses - while pending_read_ids: - item = await queue.get() + # Reopen stream if needed + should_reopen = ( + attempt_count > 1 and last_broken_generation is not None + ) or (attempt_count == 1 and metadata is not None) + if should_reopen: + broken_gen = ( + last_broken_generation + if attempt_count > 1 + else self._multiplexer.stream_generation + ) + stream_factory = self._create_stream_factory( + state, metadata + ) + await self._multiplexer.reopen_stream( + broken_gen, stream_factory + ) - if isinstance(item, _StreamEnd): - if pending_read_ids: - last_broken_generation = stream_generation - raise exceptions.ServiceUnavailable( - "Stream ended with pending read_ids" + stream_generation = self._multiplexer.stream_generation + + # Send Requests + pending_read_ids = {r.read_id for r in requests} + for i in range( + 0, len(requests), _MAX_READ_RANGES_PER_BIDI_READ_REQUEST + ): + batch = requests[ + i : i + _MAX_READ_RANGES_PER_BIDI_READ_REQUEST + ] + try: + await self._multiplexer.send( + _storage_v2.BidiReadObjectRequest(read_ranges=batch) ) - break - - if isinstance(item, _StreamError): - if item.generation < stream_generation: - continue # stale error, skip - last_broken_generation = item.generation - raise item.exception - - # Track completion - if item.object_data_ranges: - for data_range in item.object_data_ranges: - if data_range.range_end: - pending_read_ids.discard( - data_range.read_range.read_id + except Exception: + last_broken_generation = stream_generation + raise + + # Receive Responses + while pending_read_ids: + item = await queue.get() + + if isinstance(item, _StreamEnd): + if pending_read_ids: + last_broken_generation = stream_generation + raise exceptions.ServiceUnavailable( + "Stream ended with pending read_ids" ) - yield item - # The strategy has consumed `item` by now: write the - # chunks whose checksums are done, and keep the - # backlog bounded. - await strategy.verify_pending_checksums( - state, _CRC32C_MAX_PENDING - ) + break + + if isinstance(item, _StreamError): + if item.generation < stream_generation: + continue # stale error, skip + last_broken_generation = item.generation + raise item.exception + + # Track completion + if item.object_data_ranges: + for data_range in item.object_data_ranges: + if data_range.range_end: + pending_read_ids.discard( + data_range.read_range.read_id + ) + yield item + # The strategy has consumed `item` by now: write the + # chunks whose checksums are done, and keep the + # backlog bounded. + await strategy.verify_pending_checksums( + state, _CRC32C_MAX_PENDING + ) - # Write (or reject) every held chunk before the attempt - # ends, so a mismatch is raised inside the retry loop, - # as an inline one is. - await strategy.verify_pending_checksums(state) + # Write (or reject) every held chunk before the attempt + # ends, so a mismatch is raised inside the retry loop, + # as an inline one is. + await strategy.verify_pending_checksums(state) - return generator() + return generator() - retry_manager = _BidiStreamRetryManager( - strategy, send_and_recv_via_multiplexer - ) + retry_manager = _BidiStreamRetryManager( + strategy, send_and_recv_via_multiplexer + ) - try: try: - await retry_manager.execute(initial_state, retry_policy) - except DataCorruption: - raise - except Exception: - # Chunks received before the stream error still have - # in-flight checksum futures in initial_state. Verify them - # before raising the stream error so an earlier corrupt - # chunk raises DataCorruption (matching inline verification) - # and valid chunks already received are written. + try: + await retry_manager.execute(initial_state, retry_policy) + except DataCorruption: + raise + except Exception: + # Chunks received before the stream error still have + # in-flight checksum futures in initial_state. Verify them + # before raising the stream error so an earlier corrupt + # chunk raises DataCorruption (matching inline verification) + # and valid chunks already received are written. + await strategy.verify_pending_checksums(initial_state) + raise await strategy.verify_pending_checksums(initial_state) + except DataCorruption: + if self.is_stream_open: + await self.close() raise - await strategy.verify_pending_checksums(initial_state) - except DataCorruption: - if self.is_stream_open: - await self.close() - raise + finally: + # Drop chunks still held, e.g. after cancellation. + strategy.discard_pending_checksums(initial_state) + + if initial_state.get("read_handle"): + self.read_handle = initial_state["read_handle"] + + for download_state in download_states.values(): + if ( + isinstance(download_state.user_buffer, _BufferWriter) + and download_state.bytes_written + != download_state.initial_length + ): + raise BufferError( + f"Buffer destination expected {download_state.initial_length} " + f"bytes at offset {download_state.initial_offset}, got " + f"{download_state.bytes_written}" + ) finally: - # Drop chunks still held, e.g. after cancellation. - strategy.discard_pending_checksums(initial_state) - - if initial_state.get("read_handle"): - self.read_handle = initial_state["read_handle"] - - for download_state in download_states.values(): - if ( - isinstance(download_state.user_buffer, _BufferWriter) - and download_state.bytes_written != download_state.initial_length - ): - raise BufferError( - f"Buffer destination expected {download_state.initial_length} " - f"bytes at offset {download_state.initial_offset}, got " - f"{download_state.bytes_written}" - ) - finally: - if self._multiplexer is not None: - self._multiplexer.unregister(read_ids) - for download_state in download_states.values(): - if isinstance(download_state.user_buffer, _BufferWriter): - download_state.user_buffer.close() + if self._multiplexer is not None: + self._multiplexer.unregister(read_ids) + for download_state in download_states.values(): + if isinstance(download_state.user_buffer, _BufferWriter): + download_state.user_buffer.close() async def close(self): """ @@ -694,17 +721,18 @@ async def close(self): if not self._is_stream_open: raise ValueError("Underlying bidi-gRPC stream is not open") - if self._multiplexer: - await self._multiplexer.close() - self._multiplexer = None + async with self._trace_span("Storage.AsyncMultiRangeDownloader.close"): + if self._multiplexer: + await self._multiplexer.close() + self._multiplexer = None - if self.read_obj_str: - try: - await self.read_obj_str.close() - except (asyncio.CancelledError, exceptions.GoogleAPICallError): - pass - self.read_obj_str = None - self._is_stream_open = False + if self.read_obj_str: + try: + await self.read_obj_str.close() + except (asyncio.CancelledError, exceptions.GoogleAPICallError): + pass + self.read_obj_str = None + self._is_stream_open = False @property def is_stream_open(self) -> bool: diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py index daea4ec6b8f0..a79f18f84c0d 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py @@ -19,7 +19,7 @@ from google.api_core.bidi_async import AsyncBidiRpc from google.cloud import _storage_v2 -from google.cloud.storage.asyncio import _zero_copy_bidi_read +from google.cloud.storage.asyncio import _utils, _zero_copy_bidi_read from google.cloud.storage.asyncio.async_abstract_object_stream import ( _AsyncAbstractObjectStream, ) @@ -126,6 +126,7 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None: current_metadata = other_metadata current_metadata.append(("x-goog-request-params", "&".join(request_params))) + current_metadata = _utils.inject_traceparent_to_metadata(current_metadata) self.socket_like_rpc = AsyncBidiRpc( self.rpc, diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py index ab1c57c44976..ab7ae73b08e0 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_write_object_stream.py @@ -207,6 +207,7 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None: final_metadata.append((key, value)) final_metadata.append(("x-goog-request-params", "&".join(request_param_values))) + final_metadata = _utils.inject_traceparent_to_metadata(final_metadata) self.socket_like_rpc = AsyncBidiRpc( self.rpc, diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_rapid_observability.py b/packages/google-cloud-storage/tests/unit/asyncio/test_rapid_observability.py new file mode 100644 index 000000000000..0d0dfcd57a99 --- /dev/null +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_rapid_observability.py @@ -0,0 +1,988 @@ +# Copyright 2026 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. + +"""Unit tests for OpenTelemetry tracing in Rapid Storage async gRPC classes.""" + +import asyncio +import importlib +from io import BytesIO +from unittest import mock + +import pytest +from google.api_core import exceptions as api_exceptions + +from google.cloud import _storage_v2 as storage_v2 +from google.cloud.storage import _opentelemetry_tracing +from google.cloud.storage.asyncio import _utils, async_grpc_client +from google.cloud.storage.asyncio.async_appendable_object_writer import ( + AsyncAppendableObjectWriter, +) +from google.cloud.storage.asyncio.async_multi_range_downloader import ( + AsyncMultiRangeDownloader, +) +from google.cloud.storage.asyncio.async_read_object_stream import ( + _AsyncReadObjectStream, +) +from google.cloud.storage.asyncio.async_write_object_stream import ( + _AsyncWriteObjectStream, +) + + +async def _noop_coro(value): + return value + + +def _custom_eager_task_factory(loop, coro, **kwargs): + """Minimal eager task factory for coroutines that finish without suspending.""" + result = None + try: + coro.send(None) + except StopIteration as stop: + result = stop.value + return asyncio.Task(_noop_coro(result), **kwargs) + + +@pytest.fixture +def exporter(monkeypatch): + """Set up OpenTelemetry InMemorySpanExporter and enable tracing.""" + try: + from opentelemetry import trace as trace_api + from opentelemetry.sdk.trace import TracerProvider, export + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + except ImportError: # pragma: NO COVER + pytest.skip("OpenTelemetry SDK packages are required for this test suite.") + + monkeypatch.setenv("ENABLE_GCS_PYTHON_CLIENT_OTEL_TRACES", "true") + monkeypatch.delenv("DISABLE_GCS_PYTHON_CLIENT_OTEL_BUCKET_METADATA", raising=False) + + tracer_provider = TracerProvider() + memory_exporter = InMemorySpanExporter() + span_processor = export.SimpleSpanProcessor(memory_exporter) + tracer_provider.add_span_processor(span_processor) + + with mock.patch.object( + trace_api, "get_tracer", side_effect=tracer_provider.get_tracer + ): + importlib.reload(_opentelemetry_tracing) + yield memory_exporter + memory_exporter.clear() + + monkeypatch.setenv("ENABLE_GCS_PYTHON_CLIENT_OTEL_TRACES", "false") + importlib.reload(_opentelemetry_tracing) + + +@pytest.fixture +def mock_client(): + """Create an AsyncGrpcClient with a mocked underlying GAPIC client.""" + with mock.patch("google.cloud._storage_v2.StorageAsyncClient"): + client = async_grpc_client.AsyncGrpcClient( + credentials=mock.Mock(), + ) + client._grpc_client = mock.AsyncMock() + # Pre-populate ACO bucket metadata cache for Rapid Storage bucket verification + client._bucket_metadata_cache.update_cache( + "my-rapid-bucket", + "//storage.googleapis.com/projects/123456789/buckets/my-rapid-bucket", + "us-east1-a", + ) + return client + + +def test_inject_traceparent_to_metadata_when_disabled(monkeypatch): + monkeypatch.setattr(_opentelemetry_tracing, "enable_otel_traces", False) + + orig = (("x-goog-request-params", "bucket=foo"),) + result = _utils.inject_traceparent_to_metadata(orig) + assert result == orig + + +def test_inject_traceparent_to_metadata_when_enabled(exporter): + with _opentelemetry_tracing.create_trace_span("Test.ParentSpan", rpc_system="grpc"): + orig = (("x-goog-request-params", "bucket=foo"),) + result = _utils.inject_traceparent_to_metadata(orig) + keys = [k for k, _ in result] + assert "traceparent" in keys + traceparent_val = dict(result)["traceparent"] + assert traceparent_val.startswith("00-") + + +@pytest.mark.asyncio +async def test_async_write_object_stream_open_injects_traceparent(exporter): + mock_grpc_client = mock.MagicMock() + mock_grpc_client._client._transport.bidi_write_object = mock.sentinel.write_rpc + mock_grpc_client._client._transport._wrapped_methods = { + mock.sentinel.write_rpc: mock.sentinel.wrapped_write_rpc + } + + with mock.patch( + "google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc" + ) as mock_bidi_rpc_cls: + mock_rpc = mock_bidi_rpc_cls.return_value + mock_rpc.open = mock.AsyncMock() + mock_rpc.recv = mock.AsyncMock( + return_value=storage_v2.BidiWriteObjectResponse(persisted_size=0) + ) + + stream = _AsyncWriteObjectStream(mock_grpc_client, "my-rapid-bucket", "obj1") + with _opentelemetry_tracing.create_trace_span( + "Test.WriteStreamOpen", rpc_system="grpc" + ): + await stream.open() + + metadata = dict(mock_bidi_rpc_cls.call_args.kwargs["metadata"]) + assert ( + metadata["x-goog-request-params"] + == "bucket=projects/_/buckets/my-rapid-bucket" + ) + assert metadata["traceparent"].startswith("00-") + + +@pytest.mark.asyncio +async def test_async_write_object_stream_open_skips_traceparent_when_disabled( + exporter, monkeypatch +): + from opentelemetry import trace as trace_api + + monkeypatch.setattr(_opentelemetry_tracing, "enable_otel_traces", False) + + mock_grpc_client = mock.MagicMock() + mock_grpc_client._client._transport.bidi_write_object = mock.sentinel.write_rpc + mock_grpc_client._client._transport._wrapped_methods = { + mock.sentinel.write_rpc: mock.sentinel.wrapped_write_rpc + } + + with mock.patch( + "google.cloud.storage.asyncio.async_write_object_stream.AsyncBidiRpc" + ) as mock_bidi_rpc_cls: + mock_rpc = mock_bidi_rpc_cls.return_value + mock_rpc.open = mock.AsyncMock() + mock_rpc.recv = mock.AsyncMock( + return_value=storage_v2.BidiWriteObjectResponse(persisted_size=0) + ) + + stream = _AsyncWriteObjectStream(mock_grpc_client, "my-rapid-bucket", "obj1") + with trace_api.get_tracer(__name__).start_as_current_span("AmbientSpan"): + await stream.open() + + metadata = dict(mock_bidi_rpc_cls.call_args.kwargs["metadata"]) + assert ( + metadata["x-goog-request-params"] + == "bucket=projects/_/buckets/my-rapid-bucket" + ) + assert "traceparent" not in metadata + + +@pytest.mark.asyncio +async def test_async_read_object_stream_open_injects_traceparent(exporter): + mock_grpc_client = mock.MagicMock() + mock_grpc_client._client._transport.bidi_read_object = mock.sentinel.read_rpc + mock_grpc_client._client._transport._wrapped_methods = { + mock.sentinel.read_rpc: mock.sentinel.wrapped_read_rpc + } + + with mock.patch( + "google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc" + ) as mock_bidi_rpc_cls: + mock_rpc = mock_bidi_rpc_cls.return_value + mock_rpc.open = mock.AsyncMock() + mock_rpc.recv = mock.AsyncMock( + return_value=storage_v2.BidiReadObjectResponse( + read_handle=storage_v2.BidiReadHandle(handle=b"rh-1") + ) + ) + + stream = _AsyncReadObjectStream(mock_grpc_client, "my-rapid-bucket", "obj1") + with _opentelemetry_tracing.create_trace_span( + "Test.ReadStreamOpen", rpc_system="grpc" + ): + await stream.open() + + metadata = dict(mock_bidi_rpc_cls.call_args.kwargs["metadata"]) + assert ( + metadata["x-goog-request-params"] + == "bucket=projects/_/buckets/my-rapid-bucket" + ) + assert metadata["traceparent"].startswith("00-") + + +@pytest.mark.asyncio +async def test_async_read_object_stream_open_skips_traceparent_when_disabled( + exporter, monkeypatch +): + from opentelemetry import trace as trace_api + + monkeypatch.setattr(_opentelemetry_tracing, "enable_otel_traces", False) + + mock_grpc_client = mock.MagicMock() + mock_grpc_client._client._transport.bidi_read_object = mock.sentinel.read_rpc + mock_grpc_client._client._transport._wrapped_methods = { + mock.sentinel.read_rpc: mock.sentinel.wrapped_read_rpc + } + + with mock.patch( + "google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc" + ) as mock_bidi_rpc_cls: + mock_rpc = mock_bidi_rpc_cls.return_value + mock_rpc.open = mock.AsyncMock() + mock_rpc.recv = mock.AsyncMock( + return_value=storage_v2.BidiReadObjectResponse( + read_handle=storage_v2.BidiReadHandle(handle=b"rh-1") + ) + ) + + stream = _AsyncReadObjectStream(mock_grpc_client, "my-rapid-bucket", "obj1") + with trace_api.get_tracer(__name__).start_as_current_span("AmbientSpan"): + await stream.open() + + metadata = dict(mock_bidi_rpc_cls.call_args.kwargs["metadata"]) + assert ( + metadata["x-goog-request-params"] + == "bucket=projects/_/buckets/my-rapid-bucket" + ) + assert "traceparent" not in metadata + + +@pytest.mark.asyncio +async def test_async_grpc_client_get_and_delete_object_spans(exporter, mock_client): + mock_client._grpc_client.get_object.return_value = storage_v2.Object( + name="obj1", bucket="projects/_/buckets/my-rapid-bucket" + ) + mock_client._grpc_client.delete_object.return_value = None + + await mock_client.get_object("my-rapid-bucket", "obj1") + await mock_client.delete_object("my-rapid-bucket", "obj1") + + spans = exporter.get_finished_spans() + assert len(spans) == 2 + + get_span, del_span = spans[0], spans[1] + assert get_span.name == "Storage.AsyncGrpcClient.getObject" + assert get_span.attributes["rpc.system"] == "grpc" + assert get_span.attributes["gcp.client.service"] == "storage" + assert ( + get_span.attributes["gcp.resource.destination.id"] + == "//storage.googleapis.com/projects/123456789/buckets/my-rapid-bucket" + ) + assert get_span.attributes["gcp.resource.destination.location"] == "us-east1-a" + + assert del_span.name == "Storage.AsyncGrpcClient.deleteObject" + assert del_span.attributes["rpc.system"] == "grpc" + + # Verify traceparent was injected into GAPIC call metadata + get_kwargs = mock_client._grpc_client.get_object.call_args.kwargs + metadata_keys = [k for k, _ in get_kwargs["metadata"]] + assert "traceparent" in metadata_keys + + +@pytest.mark.asyncio +async def test_async_appendable_object_writer_spans(exporter, mock_client): + writer = AsyncAppendableObjectWriter( + client=mock_client, + bucket_name="my-rapid-bucket", + object_name="append-obj", + ) + + with mock.patch( + "google.cloud.storage.asyncio.async_appendable_object_writer._AsyncWriteObjectStream" + ) as mock_stream_cls: + mock_stream = mock.AsyncMock() + mock_stream.generation_number = 1001 + mock_stream.write_handle = storage_v2.BidiWriteHandle(handle=b"handle-1") + mock_stream.persisted_size = 0 + mock_stream.recv.return_value = storage_v2.BidiWriteObjectResponse( + persisted_size=11 + ) + mock_stream_cls.return_value = mock_stream + + await writer.open() + await writer.append(b"hello world") + await writer.simple_flush() + await writer.flush() + await writer.close() + + spans = exporter.get_finished_spans() + span_names = [s.name for s in spans] + assert span_names == [ + "Storage.AsyncAppendableObjectWriter.open", + "Storage.AsyncAppendableObjectWriter.append", + "Storage.AsyncAppendableObjectWriter.simpleFlush", + "Storage.AsyncAppendableObjectWriter.flush", + "Storage.AsyncAppendableObjectWriter.close", + ] + + append_span = spans[1] + assert append_span.attributes["rpc.system"] == "grpc" + assert append_span.attributes["gcp.storage.payload.size_bytes"] == 11 + assert "gcp.storage.append.bytes" not in append_span.attributes + assert "gcp.storage.chunk.size" not in append_span.attributes + assert ( + append_span.attributes["gcp.resource.destination.id"] + == "//storage.googleapis.com/projects/123456789/buckets/my-rapid-bucket" + ) + assert append_span.attributes["gcp.resource.destination.location"] == "us-east1-a" + + +@pytest.mark.asyncio +async def test_async_appendable_object_writer_state_lookup_and_append_from_file_spans( + exporter, mock_client +): + writer = AsyncAppendableObjectWriter( + client=mock_client, + bucket_name="my-rapid-bucket", + object_name="append-obj", + ) + + with mock.patch( + "google.cloud.storage.asyncio.async_appendable_object_writer._AsyncWriteObjectStream" + ) as mock_stream_cls: + mock_stream = mock.AsyncMock() + mock_stream.generation_number = 1001 + mock_stream.write_handle = storage_v2.BidiWriteHandle(handle=b"handle-1") + mock_stream.persisted_size = 0 + mock_stream.recv.return_value = storage_v2.BidiWriteObjectResponse( + persisted_size=12 + ) + mock_stream_cls.return_value = mock_stream + + await writer.open() + persisted = await writer.state_lookup() + assert persisted == 12 + + # Multi-block file emits a single appendFromFile span and suppresses nested append spans + await writer.append_from_file(BytesIO(b"a" * 12), block_size=4) + # Empty file still emits a single appendFromFile span + await writer.append_from_file(BytesIO(b""), block_size=4) + + await writer.close() + + spans = exporter.get_finished_spans() + assert [s.name for s in spans] == [ + "Storage.AsyncAppendableObjectWriter.open", + "Storage.AsyncAppendableObjectWriter.stateLookup", + "Storage.AsyncAppendableObjectWriter.appendFromFile", + "Storage.AsyncAppendableObjectWriter.appendFromFile", + "Storage.AsyncAppendableObjectWriter.close", + ] + + +@pytest.mark.asyncio +async def test_async_appendable_object_writer_finalize_and_no_nested_spans( + exporter, mock_client +): + writer = AsyncAppendableObjectWriter( + client=mock_client, + bucket_name="my-rapid-bucket", + object_name="append-obj", + ) + + with mock.patch( + "google.cloud.storage.asyncio.async_appendable_object_writer._AsyncWriteObjectStream" + ) as mock_stream_cls: + mock_stream = mock.AsyncMock() + mock_stream.generation_number = 1001 + mock_stream.write_handle = storage_v2.BidiWriteHandle(handle=b"handle-1") + mock_stream.persisted_size = 0 + mock_stream.recv.return_value = storage_v2.BidiWriteObjectResponse( + resource=storage_v2.Object(size=5) + ) + mock_stream_cls.return_value = mock_stream + + # 1. Direct finalize() call emits Storage.AsyncAppendableObjectWriter.finalize + await writer.open() + await writer.finalize() + + # 2. close(finalize_on_close=True) emits only close(), without nested finalize() + await writer.open() + # Simulate a retry on finalize inside close(finalize_on_close=True) that re-opens stream + mock_stream.send.side_effect = [ + api_exceptions.ServiceUnavailable("transient error"), + None, + ] + await writer.close(finalize_on_close=True) + + spans = exporter.get_finished_spans() + span_names = [s.name for s in spans] + assert span_names == [ + "Storage.AsyncAppendableObjectWriter.open", + "Storage.AsyncAppendableObjectWriter.finalize", + "Storage.AsyncAppendableObjectWriter.open", + "Storage.AsyncAppendableObjectWriter.close", + ] + + +@pytest.mark.asyncio +async def test_async_multi_range_downloader_spans(exporter, mock_client): + mrd = AsyncMultiRangeDownloader( + client=mock_client, + bucket_name="my-rapid-bucket", + object_name="read-obj", + ) + + with ( + mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) as mock_stream_cls, + mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._StreamMultiplexer" + ) as mock_mux_cls, + mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._BidiStreamRetryManager" + ) as mock_retry_mgr_cls, + ): + mock_stream = mock.AsyncMock() + mock_stream.generation_number = 2002 + mock_stream.read_handle = storage_v2.BidiReadHandle(handle=b"rhandle-1") + mock_stream.persisted_size = 1024 + mock_stream.is_finalized = True + mock_stream.full_obj_server_crc32c = None + mock_stream_cls.return_value = mock_stream + + mock_mux = mock.AsyncMock() + mock_mux.register = mock.Mock(return_value=mock.AsyncMock()) + mock_mux.unregister = mock.Mock() + mock_mux_cls.return_value = mock_mux + + mock_retry_mgr = mock.AsyncMock() + mock_retry_mgr_cls.return_value = mock_retry_mgr + + await mrd.open() + buf1, buf2 = BytesIO(), BytesIO() + await mrd.download_ranges( + [(0, 100, buf1), (200, 300, buf2)], enable_checksum=False + ) + await mrd.close() + + spans = exporter.get_finished_spans() + span_names = [s.name for s in spans] + assert span_names == [ + "Storage.AsyncMultiRangeDownloader.open", + "Storage.AsyncMultiRangeDownloader.downloadRanges", + "Storage.AsyncMultiRangeDownloader.close", + ] + + download_span = spans[1] + assert download_span.attributes["rpc.system"] == "grpc" + assert download_span.attributes["gcp.storage.range.count"] == 2 + assert ( + download_span.attributes["gcp.resource.destination.id"] + == "//storage.googleapis.com/projects/123456789/buckets/my-rapid-bucket" + ) + assert download_span.attributes["gcp.resource.destination.location"] == "us-east1-a" + + +@pytest.mark.asyncio +async def test_async_multi_range_downloader_data_corruption_marks_span_error( + exporter, mock_client +): + from opentelemetry import trace as trace_api + + from google.cloud.storage.exceptions import DataCorruption + + mrd = AsyncMultiRangeDownloader( + client=mock_client, + bucket_name="my-rapid-bucket", + object_name="read-obj", + ) + + with ( + mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) as mock_stream_cls, + mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._StreamMultiplexer" + ) as mock_mux_cls, + mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._BidiStreamRetryManager" + ) as mock_retry_mgr_cls, + ): + mock_stream = mock.AsyncMock() + mock_stream.generation_number = 2002 + mock_stream.read_handle = storage_v2.BidiReadHandle(handle=b"rhandle-1") + mock_stream.persisted_size = 1024 + mock_stream.is_finalized = True + mock_stream.full_obj_server_crc32c = None + mock_stream_cls.return_value = mock_stream + + mock_mux = mock.AsyncMock() + mock_mux.register = mock.Mock(return_value=mock.AsyncMock()) + mock_mux.unregister = mock.Mock() + mock_mux_cls.return_value = mock_mux + + mock_retry_mgr = mock.AsyncMock() + mock_retry_mgr.execute.side_effect = DataCorruption( + mock.Mock(), "CRC32C checksum mismatch" + ) + mock_retry_mgr_cls.return_value = mock_retry_mgr + + await mrd.open() + with pytest.raises(DataCorruption): + await mrd.download_ranges([(0, 100, BytesIO())], enable_checksum=False) + + assert not mrd.is_stream_open + spans = exporter.get_finished_spans() + assert [s.name for s in spans] == [ + "Storage.AsyncMultiRangeDownloader.open", + "Storage.AsyncMultiRangeDownloader.downloadRanges", + ] + + download_span = spans[1] + assert download_span.status.status_code == trace_api.StatusCode.ERROR + assert len(download_span.events) == 1 + assert download_span.events[0].name == "exception" + + +@pytest.mark.asyncio +async def test_concurrent_tasks_emit_independent_t2_spans(exporter, mock_client): + """Concurrent tasks sharing an MRD or writer must not suppress each other's T2 spans.""" + mrd = AsyncMultiRangeDownloader( + client=mock_client, + bucket_name="my-rapid-bucket", + object_name="read-obj", + ) + writer = AsyncAppendableObjectWriter( + client=mock_client, + bucket_name="my-rapid-bucket", + object_name="append-obj", + ) + + with ( + mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) as mock_read_stream_cls, + mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._StreamMultiplexer" + ) as mock_mux_cls, + mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._BidiStreamRetryManager" + ) as mock_retry_mgr_cls, + mock.patch( + "google.cloud.storage.asyncio.async_appendable_object_writer._AsyncWriteObjectStream" + ) as mock_write_stream_cls, + ): + mock_read_stream = mock.AsyncMock() + mock_read_stream.generation_number = 2002 + mock_read_stream.read_handle = storage_v2.BidiReadHandle(handle=b"rhandle-1") + mock_read_stream.persisted_size = 1024 + mock_read_stream.is_finalized = True + mock_read_stream.full_obj_server_crc32c = None + mock_read_stream_cls.return_value = mock_read_stream + + mock_mux = mock.AsyncMock() + mock_mux.register = mock.Mock(return_value=mock.AsyncMock()) + mock_mux.unregister = mock.Mock() + mock_mux_cls.return_value = mock_mux + + async def yielding_execute(*args, **kwargs): + await asyncio.sleep(0) + + mock_retry_mgr = mock.AsyncMock() + mock_retry_mgr.execute.side_effect = yielding_execute + mock_retry_mgr_cls.return_value = mock_retry_mgr + + mock_write_stream = mock.AsyncMock() + mock_write_stream.generation_number = 1001 + mock_write_stream.write_handle = storage_v2.BidiWriteHandle(handle=b"wh-1") + mock_write_stream.persisted_size = 0 + + async def yielding_recv(*args, **kwargs): + await asyncio.sleep(0) + return storage_v2.BidiWriteObjectResponse(persisted_size=8) + + mock_write_stream.recv.side_effect = yielding_recv + mock_write_stream_cls.return_value = mock_write_stream + + await mrd.open() + await asyncio.gather( + mrd.download_ranges([(0, 10, BytesIO())], enable_checksum=False), + mrd.download_ranges([(10, 20, BytesIO())], enable_checksum=False), + ) + await mrd.close() + + await writer.open() + await asyncio.gather(writer.flush(), writer.flush()) + await writer.close() + + span_names = [s.name for s in exporter.get_finished_spans()] + assert span_names.count("Storage.AsyncMultiRangeDownloader.downloadRanges") == 2 + assert span_names.count("Storage.AsyncAppendableObjectWriter.flush") == 2 + + +@pytest.mark.asyncio +async def test_top_level_span_context_resets_owner_on_aenter_error( + exporter, mock_client +): + owner = object() + failing_cm = mock.MagicMock() + failing_cm.__aenter__ = mock.AsyncMock(side_effect=RuntimeError("aenter failed")) + + with mock.patch.object(_utils, "create_trace_span_helper", return_value=failing_cm): + with pytest.raises(RuntimeError, match="aenter failed"): + async with _utils.create_top_level_trace_span( + owner, mock_client, "my-rapid-bucket", "Test.AenterFailure" + ): + pass # pragma: NO COVER + + assert _utils._ACTIVE_T2_SPAN_OWNERS.get() == frozenset() + + +@pytest.mark.asyncio +async def test_bucket_metadata_cache_async_grpc_fetch(mock_client): + cache = mock_client._bucket_metadata_cache + cache.clear() + + mock_client._grpc_client.get_bucket.return_value = storage_v2.Bucket( + name="projects/_/buckets/new-rapid-bucket", + project="projects/987654321", + location="US-WEST1-B", + location_type="zone", + ) + + # First call misses cache and schedules _fetch_background_async on event loop + res = cache.get_or_queue_fetch("new-rapid-bucket") + assert res is None + + # Deterministically await tracked background tasks + await asyncio.gather(*cache._background_tasks) + + cached = cache.get("new-rapid-bucket") + assert cached == ( + "//storage.googleapis.com/projects/987654321/buckets/new-rapid-bucket", + "us-west1-b", + ) + + +@pytest.mark.asyncio +async def test_bucket_metadata_cache_async_grpc_403_fallback(mock_client): + cache = mock_client._bucket_metadata_cache + cache.clear() + + mock_client._grpc_client.get_bucket.side_effect = api_exceptions.Forbidden( + "storage.buckets.get denied" + ) + + res = cache.get_or_queue_fetch("restricted-bucket") + assert res is None + + await asyncio.gather(*cache._background_tasks) + + cached = cache.get("restricted-bucket") + assert cached == ( + "//storage.googleapis.com/projects/_/buckets/restricted-bucket", + "global", + ) + + +@pytest.mark.asyncio +async def test_bucket_metadata_cache_async_grpc_404_eviction_and_refresh(mock_client): + cache = mock_client._bucket_metadata_cache + + # 1. If bucket still exists on 404 check, refresh cache entry from GetBucket response + mock_client._grpc_client.get_bucket.return_value = storage_v2.Bucket( + name="projects/_/buckets/my-rapid-bucket", + project="projects/123456789", + location="US-EAST1-B", + location_type="zone", + ) + cache.check_and_evict("my-rapid-bucket") + await asyncio.gather(*cache._background_tasks) + assert cache.get("my-rapid-bucket") == ( + "//storage.googleapis.com/projects/123456789/buckets/my-rapid-bucket", + "us-east1-b", + ) + + # 2. If GetBucket returns 404 NotFound, evict bucket from cache + mock_client._grpc_client.get_bucket.side_effect = api_exceptions.NotFound( + "Bucket not found" + ) + cache.check_and_evict("my-rapid-bucket") + await asyncio.gather(*cache._background_tasks) + assert cache.get("my-rapid-bucket") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "task_factory", + [ + _custom_eager_task_factory, + pytest.param( + getattr(asyncio, "eager_task_factory", None), + marks=pytest.mark.skipif( + not hasattr(asyncio, "eager_task_factory"), + reason="asyncio.eager_task_factory requires Python 3.12+", + ), + ), + ], +) +async def test_bucket_metadata_cache_eager_task_no_deadlock(mock_client, task_factory): + """Verify eager task execution does not deadlock on BucketMetadataCache._lock.""" + cache = mock_client._bucket_metadata_cache + cache.clear() + loop = asyncio.get_running_loop() + orig_factory = loop.get_task_factory() + + async def immediate_get_bucket(**kwargs): + assert not cache._lock.locked(), "Lock must be released before spawning task" + return storage_v2.Bucket( + name="projects/_/buckets/eager-bucket", + project="projects/555555", + location="US-CENTRAL1-A", + location_type="zone", + ) + + async def immediate_missing_bucket(**kwargs): + assert not cache._lock.locked(), "Lock must be released before spawning task" + raise api_exceptions.NotFound("Bucket deleted") + + loop.set_task_factory(task_factory) + try: + mock_client._grpc_client.get_bucket = immediate_get_bucket + assert cache.get_or_queue_fetch("eager-bucket") is None + assert cache.get("eager-bucket") == ( + "//storage.googleapis.com/projects/555555/buckets/eager-bucket", + "us-central1-a", + ) + assert "eager-bucket" not in cache._inflight_fetches + + mock_client._grpc_client.get_bucket = immediate_missing_bucket + cache.check_and_evict("eager-bucket") + assert cache.get("eager-bucket") is None + assert "eager-bucket" not in cache._inflight_checks + finally: + loop.set_task_factory(orig_factory) + + await asyncio.gather(*cache._background_tasks) + await asyncio.sleep(0) + assert len(cache._background_tasks) == 0 + + +@pytest.mark.asyncio +async def test_bucket_metadata_cache_background_get_bucket_detaches_trace_context( + exporter, mock_client +): + """Background GetBucket tasks must run in an empty contextvars.Context.""" + from opentelemetry import trace as trace_api + + cache = mock_client._bucket_metadata_cache + observed_span_valid = [] + + async def inspect_context_get_bucket(**kwargs): + span_ctx = trace_api.get_current_span().get_span_context() + observed_span_valid.append(span_ctx.is_valid) + return storage_v2.Bucket( + name="projects/_/buckets/detached-bucket", + project="projects/123456789", + location="US-EAST1-A", + location_type="zone", + ) + + mock_client._grpc_client.get_bucket.side_effect = inspect_context_get_bucket + mock_client._grpc_client.get_object.side_effect = api_exceptions.NotFound( + "Object not found" + ) + + with trace_api.get_tracer(__name__).start_as_current_span("CallerParentSpan"): + cache.get_or_queue_fetch("detached-bucket") + await asyncio.gather(*cache._background_tasks) + + cache.check_and_evict("my-rapid-bucket") + await asyncio.gather(*cache._background_tasks) + + finished_counts_at_not_found = [] + orig_check_and_evict = cache.check_and_evict + + def spy_check_and_evict(bucket_name): + finished_counts_at_not_found.append(len(exporter.get_finished_spans())) + return orig_check_and_evict(bucket_name) + + with mock.patch.object(cache, "check_and_evict", side_effect=spy_check_and_evict): + with pytest.raises(api_exceptions.NotFound): + await mock_client.get_object("my-rapid-bucket", "missing-obj") + await asyncio.gather(*cache._background_tasks) + + assert observed_span_valid == [False, False, False] + # CallerParentSpan + getObject span both finished before _handle_not_found ran + assert finished_counts_at_not_found == [2] + + +def test_bucket_metadata_cache_async_client_no_running_loop_fallback(mock_client): + """When no event loop is running, _spawn_async_task closes coro and clears inflight.""" + cache = mock_client._bucket_metadata_cache + assert cache.get_or_queue_fetch("sync-miss-bucket") is None + assert "sync-miss-bucket" not in cache._inflight_fetches + + cache.check_and_evict("my-rapid-bucket") + assert "my-rapid-bucket" not in cache._inflight_checks + assert len(cache._background_tasks) == 0 + + +def test_trace_span_helper_exit_suppressed_and_uninitialized(exporter, mock_client): + """Verify _TraceSpanHelperContext.__exit__ handles uninitialized and suppressing base CMs.""" + from google.cloud.storage import _helpers + + uninit_cm = _helpers._TraceSpanHelperContext( + mock_client, "my-rapid-bucket", "Test.Uninitialized" + ) + assert uninit_cm.__exit__(None, None, None) is False + + suppressing_base_cm = mock.MagicMock() + suppressing_base_cm.__enter__.return_value = object() + suppressing_base_cm.__exit__.return_value = True + + with ( + mock.patch.object( + _helpers, "_base_create_trace_span", return_value=suppressing_base_cm + ), + mock.patch.object( + mock_client._bucket_metadata_cache, "check_and_evict" + ) as mock_evict, + ): + with _helpers.create_trace_span_helper( + mock_client, "my-rapid-bucket", "Test.Suppressed404" + ): + raise api_exceptions.NotFound("suppressed 404") + + mock_evict.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_grpc_error_records_span_exception_once_and_evicts_on_404( + exporter, mock_client +): + from opentelemetry import trace as trace_api + + mock_client._grpc_client.get_object.side_effect = api_exceptions.NotFound( + "Object or bucket not found" + ) + mock_client._grpc_client.get_bucket.side_effect = api_exceptions.NotFound( + "Bucket not found" + ) + + with pytest.raises(api_exceptions.NotFound): + await mock_client.get_object("my-rapid-bucket", "missing-obj") + + await asyncio.gather(*mock_client._bucket_metadata_cache._background_tasks) + + # Verify bucket was evicted after 404 background existence check + assert mock_client._bucket_metadata_cache.get("my-rapid-bucket") is None + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "Storage.AsyncGrpcClient.getObject" + assert span.status.status_code == trace_api.StatusCode.ERROR + # Exception event should be recorded exactly once (no duplicate event) + exception_events = [e for e in span.events if e.name == "exception"] + assert len(exception_events) == 1 + + +@pytest.mark.asyncio +async def test_create_trace_span_helper_with_sync_mock_context_manager(mock_client): + """Verify fallback when _base_create_trace_span returns a synchronous context manager.""" + from contextlib import contextmanager + + from google.cloud.storage import _helpers + + fake_span = object() + + @contextmanager + def sync_cm(*args, **kwargs): + yield fake_span + + with mock.patch.object(_helpers, "_base_create_trace_span", side_effect=sync_cm): + async with _helpers.create_trace_span_helper( + mock_client, "my-rapid-bucket", "Test.SyncFallback" + ) as span: + assert span is fake_span + + +@pytest.mark.asyncio +async def test_noop_trace_span_fast_path_when_disabled(monkeypatch, mock_client): + """Verify zero-allocation singleton context manager is returned when tracing is disabled.""" + from google.cloud.storage import _helpers + + monkeypatch.setattr(_opentelemetry_tracing, "enable_otel_traces", False) + + cm1 = _opentelemetry_tracing.create_trace_span("Test.NoOp") + cm2 = _helpers.create_trace_span_helper(mock_client, "my-rapid-bucket", "Test.NoOp") + writer = AsyncAppendableObjectWriter( + client=mock_client, + bucket_name="my-rapid-bucket", + object_name="append-obj", + ) + mrd = AsyncMultiRangeDownloader( + client=mock_client, + bucket_name="my-rapid-bucket", + object_name="read-obj", + ) + cm3 = writer._trace_span("Test.NoOp") + cm4 = mrd._trace_span("Test.NoOp") + + assert cm1 is _opentelemetry_tracing._NOOP_TRACE_SPAN_CONTEXT + assert cm2 is _opentelemetry_tracing._NOOP_TRACE_SPAN_CONTEXT + assert cm3 is _opentelemetry_tracing._NOOP_TRACE_SPAN_CONTEXT + assert cm4 is _opentelemetry_tracing._NOOP_TRACE_SPAN_CONTEXT + + with cm1 as s1: + assert s1 is None + async with cm2 as s2: + assert s2 is None + async with cm3 as s3: + assert s3 is None + async with cm4 as s4: + assert s4 is None + + +@pytest.mark.asyncio +async def test_async_grpc_client_tracing_coexists_with_metrics_and_interconnect_options( + exporter, +): + """Verify tracing, ACO cache, OTel metrics gating, and DirectPath Interconnect compose cleanly.""" + from google.cloud.storage import _opentelemetry_metrics + + with ( + mock.patch("google.cloud._storage_v2.StorageAsyncClient"), + mock.patch.multiple( + _opentelemetry_metrics, + HAS_OPENTELEMETRY_METRICS=True, + _ENABLE_METRICS_DEV_GATE=True, + ), + ): + client = async_grpc_client.AsyncGrpcClient( + credentials=mock.Mock(), + attempt_direct_path_xds_over_interconnect=True, + enable_metrics=True, + enable_debug_metrics=True, + ) + client._grpc_client = mock.AsyncMock() + client._grpc_client.get_object.return_value = storage_v2.Object( + name="obj1", bucket="projects/_/buckets/my-rapid-bucket" + ) + client._bucket_metadata_cache.update_cache( + "my-rapid-bucket", + "//storage.googleapis.com/projects/123456789/buckets/my-rapid-bucket", + "us-east1-a", + ) + + assert client.metrics_enabled is True + assert client.debug_metrics_enabled is True + assert client._is_async_grpc_client is True + + await client.get_object("my-rapid-bucket", "obj1") + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "Storage.AsyncGrpcClient.getObject" + assert spans[0].attributes["gcp.resource.destination.location"] == "us-east1-a" diff --git a/packages/google-cloud-storage/tests/unit/test__bucket_metadata_cache.py b/packages/google-cloud-storage/tests/unit/test__bucket_metadata_cache.py index 25077d720ac5..f8e7fe7fcfff 100644 --- a/packages/google-cloud-storage/tests/unit/test__bucket_metadata_cache.py +++ b/packages/google-cloud-storage/tests/unit/test__bucket_metadata_cache.py @@ -210,3 +210,20 @@ def test_verify_existence_background_deleted(self): self.assertNotIn("b1", cache._cache) self.assertNotIn("b1", cache._inflight_checks) + + @mock.patch("threading.Thread") + def test_thread_spawned_outside_lock(self, mock_thread): + client = mock.Mock() + cache = BucketMetadataCache(client) + lock_states_at_start = [] + + def record_lock_state(): + lock_states_at_start.append(cache._lock.locked()) + + mock_thread.return_value.start.side_effect = record_lock_state + + cache.get_or_queue_fetch("b1") + cache.update_cache("b2", "dest2", "loc2") + cache.check_and_evict("b2") + + self.assertEqual(lock_states_at_start, [False, False]) diff --git a/packages/google-cloud-storage/tests/unit/test__opentelemetry_tracing.py b/packages/google-cloud-storage/tests/unit/test__opentelemetry_tracing.py index 70722bb0b0f8..4e8801aec3b8 100644 --- a/packages/google-cloud-storage/tests/unit/test__opentelemetry_tracing.py +++ b/packages/google-cloud-storage/tests/unit/test__opentelemetry_tracing.py @@ -39,9 +39,13 @@ def setup(): memory_exporter = InMemorySpanExporter() span_processor = export.SimpleSpanProcessor(memory_exporter) tracer_provider.add_span_processor(span_processor) - trace_api.set_tracer_provider(tracer_provider) + with mock.patch.object( + trace_api, "get_tracer", side_effect=tracer_provider.get_tracer + ): + importlib.reload(_opentelemetry_tracing) + yield memory_exporter + memory_exporter.clear() importlib.reload(_opentelemetry_tracing) - yield memory_exporter @pytest.fixture() @@ -56,6 +60,9 @@ def setup_optin(mock_os_environ): """Mock envar to opt-in tracing for storage client.""" mock_os_environ["ENABLE_GCS_PYTHON_CLIENT_OTEL_TRACES"] = True importlib.reload(_opentelemetry_tracing) + yield + mock_os_environ["ENABLE_GCS_PYTHON_CLIENT_OTEL_TRACES"] = "False" + importlib.reload(_opentelemetry_tracing) @pytest.fixture() @@ -63,6 +70,8 @@ def setup_optout(mock_os_environ): """Mock envar to opt-in tracing for storage client.""" mock_os_environ["ENABLE_GCS_PYTHON_CLIENT_OTEL_TRACES"] = "False" importlib.reload(_opentelemetry_tracing) + yield + importlib.reload(_opentelemetry_tracing) def test_opentelemetry_not_installed(setup, monkeypatch): @@ -140,6 +149,74 @@ def test_enable_trace_error(setup, setup_optin): assert span.name == "OtelTracing.Test" raise NotFound("Test catching NotFound error in trace span.") + spans = setup.get_finished_spans() + assert len(spans) == 1 + finished_span = spans[0] + assert finished_span.status.status_code == trace_api.StatusCode.ERROR + assert len(finished_span.events) == 1 + assert finished_span.events[0].name == "exception" + + +def test_enable_trace_non_api_error_sync(setup, setup_optin): + from opentelemetry import trace as trace_api + + from google.cloud.storage.exceptions import DataCorruption + + with pytest.raises(DataCorruption): + with _opentelemetry_tracing.create_trace_span("OtelTracing.DataCorruptionSync"): + raise DataCorruption(mock.Mock(), "CRC32C checksum mismatch") + + spans = setup.get_finished_spans() + assert len(spans) == 1 + finished_span = spans[0] + assert finished_span.name == "OtelTracing.DataCorruptionSync" + assert finished_span.status.status_code == trace_api.StatusCode.ERROR + assert len(finished_span.events) == 1 + assert finished_span.events[0].name == "exception" + + +@pytest.mark.asyncio +async def test_enable_trace_errors_async(setup, setup_optin): + from google.cloud.exceptions import NotFound + from opentelemetry import trace as trace_api + + from google.cloud.storage.exceptions import DataCorruption + + with pytest.raises(GoogleAPICallError): + async with _opentelemetry_tracing.create_trace_span( + "OtelTracing.AsyncApiError" + ) as span: + assert span.kind == trace_api.SpanKind.CLIENT + raise NotFound("Async NotFound error") + + with pytest.raises(DataCorruption): + async with _opentelemetry_tracing.create_trace_span( + "OtelTracing.AsyncDataCorruption" + ): + raise DataCorruption(mock.Mock(), "Async CRC32C checksum mismatch") + + spans = setup.get_finished_spans() + assert len(spans) == 2 + + api_err_span, corruption_span = spans[0], spans[1] + assert api_err_span.name == "OtelTracing.AsyncApiError" + assert api_err_span.status.status_code == trace_api.StatusCode.ERROR + assert len(api_err_span.events) == 1 + assert api_err_span.events[0].name == "exception" + + assert corruption_span.name == "OtelTracing.AsyncDataCorruption" + assert corruption_span.status.status_code == trace_api.StatusCode.ERROR + assert len(corruption_span.events) == 1 + assert corruption_span.events[0].name == "exception" + + +def test_trace_span_context_disabled_before_enter(setup, setup_optin, monkeypatch): + cm = _opentelemetry_tracing.create_trace_span("OtelTracing.DisabledBeforeEnter") + monkeypatch.setattr(_opentelemetry_tracing, "enable_otel_traces", False) + with cm as span: + assert span is None + assert len(setup.get_finished_spans()) == 0 + def test_get_final_attributes(setup, setup_optin): from google.api_core import retry as api_retry