mongo/buildscripts/resmokelib/utils/batched_baggage_span_proces...

29 lines
1.1 KiB
Python

from typing import Optional
from opentelemetry.baggage import get_all as get_all_baggage
from opentelemetry.context import Context
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
from opentelemetry.trace import Span
class BatchedBaggageSpanProcessor(BatchSpanProcessor):
"""
The BatchedBaggageSpanProcessor is combination between a baggage propagator and the BatchSpanProcessor.
The BatchedBaggageSpanProcessor reads entries stored in Baggage
from the parent context and adds the baggage entries' keys and
values to the span as attributes on span start.
Use this to propogate attributes you want in every nested span.
"""
def __init__(self, span_exporter: SpanExporter):
# Lower the max_export_batch_size because we were hitting sizes too big for evergreen
super().__init__(span_exporter, max_export_batch_size=100)
def on_start(self, span: "Span", parent_context: Optional[Context] = None) -> None:
baggage = get_all_baggage(parent_context)
for key, value in baggage.items():
span.set_attribute(key, value)
super().on_start(span, parent_context)