From ac63a3066d3ee78897778ee7162d9ea4d6ef71a6 Mon Sep 17 00:00:00 2001 From: Lubomír Sedlář Date: Apr 29 2025 07:21:30 +0000 Subject: Add OpenTelemetry instrumentation module This module allows Koji library users to enable instrumentation by wrapping the ClientSession. It will then create a span for each API call. The spans store which method was called and what the arguments were. This is completely transparent to the server. The instrumentation does not forward any trace IDs to the hub (unless requests are instrumented, in which case the traceparent header is added to all HTTP calls). --- diff --git a/koji/opentelemetry/instrumentation.py b/koji/opentelemetry/instrumentation.py new file mode 100644 index 0000000..3c338e6 --- /dev/null +++ b/koji/opentelemetry/instrumentation.py @@ -0,0 +1,188 @@ +from contextlib import ExitStack +import itertools + +import koji + +from opentelemetry import trace + + +def _name(session, name): + """Helper function to generate nicely readable names for individual spans.""" + return "%s.%s" % (session.__class__.__name__, name) + + +def _format_args(args, kwargs): + """Format API call arguments as a string so that they can be attached to + spans. + """ + return ", ".join( + itertools.chain( + (repr(arg) for arg in args), + (f"{key}={value!r}" for key, value in kwargs.items()), + ) + ) + + +def _instrument_method(tracer, span_name, callable): + """Wrap an API method call in a span, and attach the arguments as attribute + of the span. + """ + def wrapper(*args, **kwargs): + with tracer.start_as_current_span(span_name) as span: + span.set_attribute("arguments", _format_args(args, kwargs)) + return callable(*args, **kwargs) + + return wrapper + + +class InstrumentedMulticallHack: + def __init__(self, tracer, multicall): + self._tracer = tracer + self._multicall = multicall + + def __call__(self, *args, **kwargs): + return InstrumentedMulticallSession( + self._tracer, self._multicall(*args, **kwargs) + ) + + +class InstrumentedMulticallSession: + def __init__(self, tracer, session): + self._tracer = tracer + self._session = session + self._stack = ExitStack() + self._span = None + + def __getattr__(self, name): + if not self._span: + self._span = self._stack.enter_context( + self._tracer.start_as_current_span("multicall") + ) + return _instrument_method( + self._tracer, _name(self._session, name), getattr(self._session, name) + ) + + def call_all(self, *args, **kwargs): + result = self._session.call_all(*args, **kwargs) + self._stack.close() + self._span = None + return result + + def __enter__(self): + self._span = self._stack.enter_context( + self._tracer.start_as_current_span("multicall") + ) + self._stack.enter_context(self._session) + return self + + def __exit__(self, exc_type, exc_value, tb): + self._stack.close() + self._span = None + + +class InstrumentedSystem: + def __init__(self, tracer, system): + self._tracer = tracer + self._system = system + + def __getattr__(self, name): + return _instrument_method( + self._tracer, _name(self._system, name), getattr(self._system, name) + ) + + +class instrument_client_session: + """Main entry point of this module. + Calling this with a ClientSession will add instrumentation to all API calls + done by that session. + """ + def __init__(self, session): + self._tracer = trace.get_tracer("koji.opentelemetry.instrumentation") + self._session = session + self._stack = ExitStack() + + @property + def system(self): + return InstrumentedSystem(self._tracer, self._session.system) + + def __getattr__(self, name): + return _instrument_method( + self._tracer, _name(self._session, name), getattr(self._session, name) + ) + + @property + def multicall(self): + return InstrumentedMulticallHack(self._tracer, self._session._multicall) + + @multicall.setter + def multicall(self, value): + self._stack.enter_context(self._tracer.start_as_current_span("multicall")) + self._session.multicall = value + + def multiCall(self, *args, **kwargs): + result = self._session.multiCall(*args, **kwargs) + self._stack.close() + return result + + +if __name__ == "__main__": + # This code doesn't really do anything useful other than exercise the tracing. + + # Imports + boilerplate + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter, + ) + from opentelemetry.instrumentation.requests import RequestsInstrumentor + + provider = TracerProvider( + resource=Resource(attributes={"service.name": "instrumentation-demo"}) + ) + processor = BatchSpanProcessor(OTLPSpanExporter()) + provider.add_span_processor(processor) + trace.set_tracer_provider(provider) + + # Application tracer + tracer = trace.get_tracer(__name__) + + # This is optional: by default instrumenting requests is quite detrimental + # to tracing koji as it generates span for each request, but with no way of + # telling what it actually does. + RequestsInstrumentor().instrument() + + # Create our client and run some operations with it. + m = koji.get_profile_module("koji") + c = instrument_client_session(koji.ClientSession(m.config.server)) + + with tracer.start_as_current_span("tracing-demo"): + print(len(c.system.listMethods())) + + with tracer.start_as_current_span("separate-calls"): + # Just two API calls in sequence + print(c.getBuild(1)["nvr"]) + print(c.getBuild(2, strict=True)["nvr"]) + + with tracer.start_as_current_span("multicall-session"): + # Two calls with an ugly multicall interface. + m = c.multicall() + b1 = m.getBuild(1) + b2 = m.getBuild(2, strict=True) + m.call_all() + print(b1.result["nvr"], b2.result["nvr"]) + + with tracer.start_as_current_span("multicall-contextmanager"): + # Two calls with a good multicall interface. + with c.multicall() as m: + b1 = m.getBuild(1) + b2 = m.getBuild(2, strict=True) + print(b1.result["nvr"], b2.result["nvr"]) + + with tracer.start_as_current_span("multicall-modal"): + # Two calls with bad multicall interface. + c.multicall = True + c.getBuild(1) + c.getBuild(2, strict=True) + res = c.multiCall() + print([x[0]["nvr"] for x in res])