Coverage for app/backend/src/couchers/tracing.py: 62%

1from collections.abc import Sequence 

2 

3from opentelemetry import trace 

4from opentelemetry.context import Context 

5from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter 

6from opentelemetry.instrumentation.grpc import GrpcInstrumentorServer 

7from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor 

8from opentelemetry.instrumentation.threading import ThreadingInstrumentor 

9from opentelemetry.sdk.resources import Resource 

10from opentelemetry.sdk.trace import TracerProvider 

11from opentelemetry.sdk.trace.export import BatchSpanProcessor 

12from opentelemetry.sdk.trace.sampling import ParentBased, Sampler, SamplingResult, TraceIdRatioBased 

13from opentelemetry.trace import Link, SpanKind 

14from opentelemetry.trace.span import TraceState 

15from opentelemetry.util.types import Attributes 

16 

17from couchers.config import config 

18from couchers.db import _get_base_engine 

19from couchers.experimentation import get_global_float_value 

20 

21 

22class FeatureFlagRatioSampler(Sampler): 

23 """ 

24 Samples the fraction of traces given by the `trace_sample_ratio` flag (default 0). Wired as the 

25 ParentBased root, so it's consulted only for root spans - one flag read per RPC, not per span. 

26 """ 

27 

28 def should_sample( 

29 self, 

30 parent_context: Context | None, 

31 trace_id: int, 

32 name: str, 

33 kind: SpanKind | None = None, 

34 attributes: Attributes = None, 

35 links: Sequence[Link] | None = None, 

36 trace_state: TraceState | None = None, 

37 ) -> SamplingResult: 

38 ratio = get_global_float_value("trace_sample_ratio", 0.0) 

39 return TraceIdRatioBased(ratio).should_sample( 

40 parent_context, trace_id, name, kind, attributes, links, trace_state 

41 ) 

42 

43 def get_description(self) -> str: 

44 return "FeatureFlagRatioSampler{flag=trace_sample_ratio}" 

45 

46 

47def setup_tracing() -> None: 

48 if config.OPENTELEMETRY_ENDPOINT != "": 

49 ThreadingInstrumentor().instrument() 

50 

51 grpc_server_instrumentor = GrpcInstrumentorServer() # type: ignore[no-untyped-call] 

52 grpc_server_instrumentor.instrument() 

53 SQLAlchemyInstrumentor().instrument(engine=_get_base_engine(), enable_commenter=True, commenter_options={}) 

54 

55 tracer = TracerProvider( 

56 resource=Resource(attributes={"service.name": "backend"}), 

57 sampler=ParentBased(root=FeatureFlagRatioSampler()), 

58 ) 

59 tracer.add_span_processor( 

60 BatchSpanProcessor(OTLPSpanExporter(endpoint=config.OPENTELEMETRY_ENDPOINT, insecure=True)) 

61 ) 

62 

63 trace.set_tracer_provider(tracer)