Coverage for gco / services / tracing.py: 100.00%
290 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""OpenTelemetry tracing for the GCO API services, exported to AWS X-Ray.
3The four FastAPI services (``health-monitor``, ``manifest-processor``,
4``inference-proxy``, ``cost-monitor``) record a server span per request and a
5client span per in-cluster hop they make, and send them straight to X-Ray's
6OTLP endpoint, where Transaction Search indexes them into the ``aws/spans`` log
7group. There is no collector to deploy or keep patched: each process signs its
8own OTLP/HTTP exports with the pod's AWS credentials.
10Design constraints, and how this module meets them:
12* **Inert unless asked.** Nothing happens unless ``GCO_TRACING_ENABLED`` is
13 ``true``, so tests, local runs and images without AWS credentials never
14 export. GCO deployments trace by default only because the manifests always
15 set the variable from ``cdk.json``.
16* **Import-light.** OpenTelemetry, botocore and httpx2 are imported inside the
17 functions that need them, and never while tracing is off. The
18 inference-monitor and queue-processor images ship without OpenTelemetry and
19 still import this module through :mod:`gco.services.structured_logging`.
20* **Never breaks the service.** A missing region, credential or package logs
21 one warning and leaves tracing off. Exports run on the SDK's background
22 thread, never raise, and rate-limit their warnings.
23* **Callers cannot steer sampling.** Roots and remote parents alike are sampled
24 by trace-id ratio, so an incoming ``traceparent`` flag neither forces nor
25 suppresses recording. Every GCO hop applies the same ratio to the same trace
26 id, so a trace is kept or dropped as a whole.
27* **No self-tracing.** The exporter's HTTP client is never instrumented, so
28 exporting spans cannot produce spans.
29"""
31from __future__ import annotations
33import gzip
34import logging
35import os
36import re
37import threading
38import time
39from collections.abc import Callable, Sequence
40from dataclasses import dataclass
41from typing import TYPE_CHECKING, Any, Protocol, cast
42from urllib.parse import urlsplit
44if TYPE_CHECKING: # pragma: no cover - typing only
45 import httpx2
46 from fastapi import FastAPI
47 from opentelemetry.metrics import MeterProvider
48 from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
49 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
51 # The exporter implements the SDK's SpanExporter interface but inherits
52 # from it only for the type checker: a real base class would import the
53 # SDK whenever this module loads, which the OpenTelemetry-free images
54 # cannot do. The SDK never isinstance-checks its exporters.
55 _SpanExporterBase = SpanExporter
57 class _CredentialSource(Protocol):
58 """What the exporter needs from botocore credentials.
60 Refreshable credentials (Pod Identity, IRSA) renew themselves inside
61 ``get_frozen_credentials``, so asking again for every attempt is how
62 a long-lived exporter keeps signing with current keys.
63 """
65 def get_frozen_credentials(self) -> object:
66 """Return the credentials to sign one request with, refreshed first if due."""
68else:
69 _SpanExporterBase = object
71logger = logging.getLogger(__name__)
73TRACING_ENABLED_ENV = "GCO_TRACING_ENABLED"
74TRACING_SAMPLE_RATIO_ENV = "GCO_TRACING_SAMPLE_RATIO"
75TRACING_ENDPOINT_ENV = "GCO_TRACING_ENDPOINT"
76DEFAULT_SAMPLE_RATIO = 0.05
77EXCLUDED_PATHS = ("/healthz", "/readyz", "/metrics", "/api/v1/health")
79#: ``excluded_urls`` for the FastAPI instrumentation, which ``re.search``-es
80#: each pattern against ``scheme://host[:port]/path`` (query already removed).
81#: Anchoring both ends, with an authority that cannot contain ``/``, excludes
82#: exactly these paths: the ``/metrics`` scrape is skipped while the
83#: health-monitor's ``/api/v1/metrics`` API route is still traced.
84_EXCLUDED_URLS = ",".join(
85 rf"^[a-z][a-z0-9+.-]*://[^/]*{re.escape(path)}\Z" for path in EXCLUDED_PATHS
86)
88#: The region becomes part of the default endpoint's host name.
89_REGION_RE = re.compile(r"[a-z]+(?:-[a-z]+)+-\d+")
90_XRAY_SIGNING_NAME = "xray"
91#: Per-export budget: the first attempt's timeout, and the deadline every
92#: retry and backoff must also fit inside.
93_EXPORT_TIMEOUT_SECONDS = 10.0
94_MAX_EXPORT_ATTEMPTS = 4
95_RETRY_BASE_DELAY_SECONDS = 0.5
96_WARNING_INTERVAL_SECONDS = 60.0
97#: How much of an error response a warning quotes.
98_BODY_PREVIEW_CHARS = 200
101class _Throttle:
102 """Admits one event per interval and counts the ones it holds back.
104 Export failures repeat on every batch (a missing permission, an
105 unreachable endpoint), so one warning a minute that carries the number of
106 repeats it stands for says everything without flooding the service log.
107 """
109 def __init__(self, interval: float, clock: Callable[[], float] = time.monotonic) -> None:
110 self._interval = interval
111 self._clock = clock
112 self._lock = threading.Lock()
113 self._last: float | None = None
114 self._held = 0
116 def admit(self) -> int | None:
117 """Return how many events were held back since the last admitted one,
118 or ``None`` when this event must be held back too."""
119 with self._lock:
120 now = self._clock()
121 if self._last is not None and now - self._last < self._interval:
122 self._held += 1
123 return None
124 self._last = now
125 held, self._held = self._held, 0
126 return held
129def _throttled_warning(throttle: _Throttle, message: str, *args: object) -> None:
130 held = throttle.admit()
131 if held is None:
132 return
133 if held:
134 message += " (%d similar warnings suppressed)"
135 args = (*args, held)
136 logger.warning(message, *args)
139def _is_retryable_status(status_code: int) -> bool:
140 """Throttling and server-side failures are worth another attempt."""
141 return status_code == 429 or 500 <= status_code <= 599
144def _new_export_client() -> httpx2.Client:
145 """Build the exporter's own client.
147 It trusts the public CA store (httpx2's default ``verify``, backed by
148 truststore), ignores proxy and CA variables from the environment, and is
149 never wrapped by the tracing transport.
150 """
151 import httpx2
153 return httpx2.Client(timeout=_EXPORT_TIMEOUT_SECONDS, trust_env=False, follow_redirects=False)
156class _XRaySpanExporter(_SpanExporterBase):
157 """Sends span batches to X-Ray's OTLP endpoint as SigV4-signed HTTPS posts.
159 Each batch is encoded once, as gzip-compressed OTLP protobuf, and posted
160 through one shared client. Every attempt is signed afresh, so
161 ``X-Amz-Date`` stays current across retries and rotated credentials are
162 picked up. Throttling, server errors and transport errors are retried with
163 exponential backoff while the next attempt still fits the export deadline;
164 a backoff wait ends early when the exporter shuts down.
165 """
167 def __init__(
168 self,
169 *,
170 endpoint: str,
171 region: str,
172 credentials: _CredentialSource,
173 client: httpx2.Client | None = None,
174 timeout: float = _EXPORT_TIMEOUT_SECONDS,
175 clock: Callable[[], float] = time.monotonic,
176 wait: Callable[[float], bool] | None = None,
177 ) -> None:
178 self._endpoint = endpoint
179 self._region = region
180 self._credentials = credentials
181 self._client = client if client is not None else _new_export_client()
182 self._timeout = timeout
183 self._clock = clock
184 self._stopped = threading.Event()
185 # Returns True when shutdown cut the wait short.
186 self._wait = wait if wait is not None else self._stopped.wait
187 self._warnings = _Throttle(_WARNING_INTERVAL_SECONDS, clock)
189 def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
190 from opentelemetry.sdk.trace.export import SpanExportResult
192 if self._stopped.is_set():
193 return SpanExportResult.FAILURE
194 if not spans:
195 return SpanExportResult.SUCCESS
196 try:
197 body = self._encode(spans)
198 except Exception as exc:
199 self._warn(
200 "Dropped %d span(s): they could not be encoded (%s: %r)",
201 len(spans),
202 type(exc).__name__,
203 str(exc),
204 )
205 return SpanExportResult.FAILURE
206 if self._post(body, len(spans)):
207 return SpanExportResult.SUCCESS
208 return SpanExportResult.FAILURE
210 def shutdown(self) -> None:
211 if self._stopped.is_set():
212 return
213 self._stopped.set()
214 self._client.close()
216 def force_flush(self, timeout_millis: int = 30000) -> bool:
217 # Nothing is buffered here: the batch span processor owns the queue.
218 return True
220 @staticmethod
221 def _encode(spans: Sequence[ReadableSpan]) -> bytes:
222 from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans
224 return gzip.compress(encode_spans(spans).SerializeToString(), compresslevel=6)
226 def _post(self, body: bytes, span_count: int) -> bool:
227 """Deliver one encoded batch; True once X-Ray accepted it."""
228 import httpx2
230 deadline = self._clock() + self._timeout
231 attempts = 0
232 failure = "the export deadline passed before the first attempt"
233 for attempt in range(_MAX_EXPORT_ATTEMPTS):
234 if attempt:
235 delay = _RETRY_BASE_DELAY_SECONDS * 2 ** (attempt - 1)
236 if delay >= deadline - self._clock() or self._wait(delay):
237 break
238 remaining = deadline - self._clock()
239 if remaining <= 0:
240 break
241 attempts += 1
242 try:
243 response = self._send(body, remaining)
244 except httpx2.TransportError as exc:
245 failure = f"{type(exc).__name__}: {str(exc)!r}"
246 continue
247 except Exception as exc:
248 # Signing or refreshing credentials failed; an instant retry cannot help.
249 failure = f"{type(exc).__name__}: {str(exc)!r}"
250 break
251 if response.is_success:
252 self._report_partial_success(response, span_count)
253 return True
254 failure = f"HTTP {response.status_code}: {response.text[:_BODY_PREVIEW_CHARS]!r}"
255 if not _is_retryable_status(response.status_code):
256 break
257 self._warn(
258 "Dropped %d span(s): X-Ray export failed after %d attempt(s): %s",
259 span_count,
260 attempts,
261 failure,
262 )
263 return False
265 def _send(self, body: bytes, timeout: float) -> httpx2.Response:
266 from botocore.auth import SigV4Auth
267 from botocore.awsrequest import AWSRequest
269 request = AWSRequest(
270 method="POST",
271 url=self._endpoint,
272 data=body,
273 headers={"Content-Type": "application/x-protobuf", "Content-Encoding": "gzip"},
274 )
275 SigV4Auth(
276 self._credentials.get_frozen_credentials(), _XRAY_SIGNING_NAME, self._region
277 ).add_auth(request)
278 # The body and every signed header go out exactly as they were signed.
279 headers = {str(name): str(value) for name, value in request.headers.items()}
280 return self._client.post(self._endpoint, content=body, headers=headers, timeout=timeout)
282 def _report_partial_success(self, response: httpx2.Response, span_count: int) -> None:
283 """Warn when X-Ray accepted the request but rejected some of its spans."""
284 from google.protobuf.message import DecodeError
285 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
286 ExportTraceServiceResponse,
287 )
289 if not response.headers.get("content-type", "").startswith("application/x-protobuf"):
290 return
291 try:
292 partial = ExportTraceServiceResponse.FromString(response.content).partial_success
293 except DecodeError:
294 return
295 if partial.rejected_spans:
296 self._warn(
297 "X-Ray rejected %d of %d exported span(s): %r",
298 partial.rejected_spans,
299 span_count,
300 partial.error_message[:_BODY_PREVIEW_CHARS],
301 )
303 def _warn(self, message: str, *args: object) -> None:
304 _throttled_warning(self._warnings, message, *args)
307@dataclass(frozen=True, slots=True)
308class _ActiveTracing:
309 """What the other entry points use once :func:`configure_tracing` succeeded."""
311 provider: TracerProvider
312 #: Instrumentation histograms go nowhere: GCO's service metrics are Prometheus.
313 meter_provider: MeterProvider
316_lock = threading.Lock()
317_active: _ActiveTracing | None = None
318#: Set by the first enabled configure call whatever its outcome, so a failed
319#: setup warns once per process and a shut-down process stays untraced. Only
320#: read and set while holding ``_lock``.
321_configure_attempted = threading.Event()
322_instrumentation_warnings = _Throttle(_WARNING_INTERVAL_SECONDS)
325def tracing_enabled() -> bool:
326 """True only when ``GCO_TRACING_ENABLED`` is ``true`` (case-insensitive)."""
327 return os.getenv(TRACING_ENABLED_ENV, "").strip().lower() == "true"
330def sample_ratio() -> float:
331 """Head-sampling ratio in [0, 1]; invalid or unset values use the default."""
332 raw = os.getenv(TRACING_SAMPLE_RATIO_ENV, "").strip()
333 if not raw:
334 return DEFAULT_SAMPLE_RATIO
335 try:
336 value = float(raw)
337 except ValueError:
338 pass
339 else:
340 if 0.0 <= value <= 1.0:
341 return value
342 logger.warning(
343 "Ignoring %s=%r: expected a number from 0 to 1; sampling %s of traces instead",
344 TRACING_SAMPLE_RATIO_ENV,
345 raw,
346 DEFAULT_SAMPLE_RATIO,
347 )
348 return DEFAULT_SAMPLE_RATIO
351def _configured_region() -> str:
352 for name in ("REGION", "AWS_REGION", "AWS_DEFAULT_REGION"):
353 value = os.getenv(name, "").strip()
354 if value:
355 return value
356 return ""
359def _default_endpoint(region: str) -> str:
360 """X-Ray's OTLP traces URL in ``region``, on that Region's partition.
362 botocore's endpoint metadata knows each partition's DNS suffix (for
363 example ``amazonaws.com.cn`` in the China Regions), so the URL names the
364 host the X-Ray API itself answers on in that Region; no request is made
365 and no credentials are needed. An unexpected lookup failure falls back to
366 the commercial suffix, which ``GCO_TRACING_ENDPOINT`` can override.
367 """
368 try:
369 import botocore.session
371 base = (
372 botocore.session.Session().create_client("xray", region_name=region).meta.endpoint_url
373 )
374 except Exception:
375 return f"https://xray.{region}.amazonaws.com/v1/traces"
376 return f"{str(base).rstrip('/')}/v1/traces"
379def _resource_attributes(service_name: str, region: str) -> dict[str, str]:
380 """Attributes naming the service and where it runs; unset ones are omitted."""
381 from gco import __version__
383 attributes = {
384 "service.name": service_name,
385 "service.namespace": "gco",
386 "service.version": __version__,
387 "cloud.provider": "aws",
388 "cloud.platform": "aws_eks",
389 "cloud.region": region,
390 }
391 for attribute, env_name in (
392 ("k8s.cluster.name", "CLUSTER_NAME"),
393 ("k8s.namespace.name", "POD_NAMESPACE"),
394 ("k8s.pod.name", "POD_NAME"),
395 ):
396 value = os.getenv(env_name, "").strip()
397 if value:
398 attributes[attribute] = value
399 return attributes
402def _setup(service_name: str) -> _ActiveTracing | None:
403 """Build and install the process tracer provider, or log why not and return None."""
404 region = _configured_region()
405 if not region:
406 logger.warning(
407 "Tracing disabled: %s is true but no AWS region is set "
408 "(REGION, AWS_REGION or AWS_DEFAULT_REGION)",
409 TRACING_ENABLED_ENV,
410 )
411 return None
412 if _REGION_RE.fullmatch(region) is None:
413 logger.warning("Tracing disabled: %r is not an AWS region name", region)
414 return None
415 endpoint = os.getenv(TRACING_ENDPOINT_ENV, "").strip() or _default_endpoint(region)
416 parts = urlsplit(endpoint)
417 if parts.scheme != "https" or not parts.hostname:
418 # Signed exports carry a session token, which must not travel in clear text.
419 logger.warning(
420 "Tracing disabled: %s must be an https:// URL with a host (got scheme %r)",
421 TRACING_ENDPOINT_ENV,
422 parts.scheme,
423 )
424 return None
426 try:
427 # Everything used later (by the exporter, instrument_fastapi_app and
428 # the transport wrappers) is imported now too, so a package missing
429 # from the image costs this one warning instead of tracing that only
430 # half works.
431 import botocore.auth
432 import botocore.awsrequest
433 import botocore.session
434 import httpx2 # noqa: F401
435 import opentelemetry.exporter.otlp.proto.common.trace_encoder
436 import opentelemetry.instrumentation.fastapi
437 import opentelemetry.instrumentation.httpx
438 import opentelemetry.proto.collector.trace.v1.trace_service_pb2 # noqa: F401
439 from opentelemetry import trace
440 from opentelemetry.metrics import NoOpMeterProvider
441 from opentelemetry.propagate import set_global_textmap
442 from opentelemetry.sdk.resources import Resource
443 from opentelemetry.sdk.trace import TracerProvider
444 from opentelemetry.sdk.trace.export import BatchSpanProcessor
445 from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased
446 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
447 except ImportError as exc:
448 logger.warning(
449 "Tracing disabled: a tracing package is missing from this image (%s)", exc.name or exc
450 )
451 return None
453 credentials = botocore.session.Session().get_credentials()
454 if credentials is None:
455 logger.warning("Tracing disabled: no AWS credentials are available to sign X-Ray exports")
456 return None
458 ratio = sample_ratio()
459 sampler = TraceIdRatioBased(ratio)
460 provider = TracerProvider(
461 resource=Resource.create(_resource_attributes(service_name, region)),
462 sampler=ParentBased(
463 root=sampler, remote_parent_sampled=sampler, remote_parent_not_sampled=sampler
464 ),
465 )
466 exporter = _XRaySpanExporter(endpoint=endpoint, region=region, credentials=credentials)
467 provider.add_span_processor(BatchSpanProcessor(exporter))
468 trace.set_tracer_provider(provider)
469 # W3C trace context only: GCO propagates no baggage.
470 set_global_textmap(TraceContextTextMapPropagator())
471 logger.info(
472 "Tracing enabled for %s: sampling %s of traces to %s", service_name, ratio, parts.hostname
473 )
474 return _ActiveTracing(provider=provider, meter_provider=NoOpMeterProvider())
477def configure_tracing(service_name: str) -> bool:
478 """Install the process-wide tracer provider once; True when tracing is active.
480 Safe to call from every module that builds an app: a service module runs
481 twice per process (``python -m`` and Uvicorn's import string), and the
482 second call reuses the first call's provider. Never raises.
483 """
484 global _active
485 if not tracing_enabled():
486 return False
487 with _lock:
488 if not _configure_attempted.is_set():
489 _configure_attempted.set()
490 try:
491 _active = _setup(service_name)
492 except Exception:
493 logger.warning("Tracing disabled: setup failed", exc_info=True)
494 return _active is not None
497def instrument_fastapi_app(app: FastAPI) -> None:
498 """Add server spans to one FastAPI app when tracing is active.
500 Probe and scrape paths (:data:`EXCLUDED_PATHS`) are not traced, the ASGI
501 ``receive``/``send`` sub-spans are dropped and no headers are captured.
502 Instrumenting the same app again does nothing.
503 """
504 active = _active
505 if active is None or getattr(app, "_is_instrumented_by_opentelemetry", False):
506 return
507 try:
508 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
510 FastAPIInstrumentor.instrument_app(
511 app,
512 tracer_provider=active.provider,
513 meter_provider=active.meter_provider,
514 excluded_urls=_EXCLUDED_URLS,
515 exclude_spans=["receive", "send"],
516 )
517 except Exception:
518 _throttled_warning(
519 _instrumentation_warnings,
520 "Tracing: could not instrument the FastAPI app; serving it without server spans",
521 )
524def wrap_async_transport(transport: httpx2.AsyncBaseTransport) -> httpx2.AsyncBaseTransport:
525 """Wrap an async httpx2 transport with client spans when tracing is active.
527 The wrapper also injects ``traceparent`` into each request, so the next
528 GCO hop continues the trace.
529 """
530 active = _active
531 if active is None:
532 return transport
533 try:
534 from opentelemetry.instrumentation.httpx import AsyncOpenTelemetryTransportHttpx2
536 return AsyncOpenTelemetryTransportHttpx2(
537 # The instrumentation annotates the wrapped transport with httpx's
538 # class; this wrapper is its httpx2 flavour, same interface.
539 cast(Any, transport),
540 tracer_provider=active.provider,
541 meter_provider=active.meter_provider,
542 )
543 except Exception:
544 _throttled_warning(
545 _instrumentation_warnings,
546 "Tracing: could not wrap an HTTP transport; its requests are not traced",
547 )
548 return transport
551def wrap_sync_transport(transport: httpx2.BaseTransport) -> httpx2.BaseTransport:
552 """Wrap a sync httpx2 transport with client spans when tracing is active."""
553 active = _active
554 if active is None:
555 return transport
556 try:
557 from opentelemetry.instrumentation.httpx import SyncOpenTelemetryTransportHttpx2
559 return SyncOpenTelemetryTransportHttpx2(
560 cast(Any, transport),
561 tracer_provider=active.provider,
562 meter_provider=active.meter_provider,
563 )
564 except Exception:
565 _throttled_warning(
566 _instrumentation_warnings,
567 "Tracing: could not wrap an HTTP transport; its requests are not traced",
568 )
569 return transport
572def current_trace_fields() -> dict[str, object]:
573 """Trace/span ids of the current span for log correlation, or ``{}``.
575 Runs for every JSON log line, so while tracing is inactive it answers from
576 one global read. OpenTelemetry is only consulted after
577 :func:`configure_tracing` imported it, which images without it never do.
578 Unsampled requests keep their ids (``trace_sampled`` is false), so their
579 log lines still correlate across services. Never raises.
580 """
581 if _active is None:
582 return {}
583 try:
584 from opentelemetry import trace
586 context = trace.get_current_span().get_span_context()
587 if not context.is_valid:
588 return {}
589 return {
590 "trace_id": format(context.trace_id, "032x"),
591 "span_id": format(context.span_id, "016x"),
592 "trace_sampled": context.trace_flags.sampled,
593 }
594 except Exception:
595 return {}
598def shutdown_tracing() -> None:
599 """Flush and shut down the tracer provider; idempotent and never raises.
601 Called from each service's lifespan shutdown: queued spans are exported,
602 within the export deadline, before the exporter's client closes. The
603 process does not trace again afterwards.
604 """
605 global _active
606 with _lock:
607 active, _active = _active, None
608 if active is None:
609 return
610 try:
611 active.provider.shutdown()
612 except Exception:
613 logger.warning(
614 "Tracing: shutdown did not complete; queued spans may be lost", exc_info=True
615 )