mongo/buildscripts/resmokelib/utils/otel_thread_pool_executor.py

18 lines
528 B
Python

import concurrent.futures
import opentelemetry.context
class OtelThreadPoolExecutor(concurrent.futures.ThreadPoolExecutor):
def submit(self, fn, *args, **kwargs):
context = opentelemetry.context.get_current()
if context:
def attach_context(*args, **kwargs):
opentelemetry.context.attach(context)
return fn(*args, **kwargs)
return super().submit(attach_context, *args, **kwargs)
else:
return super().submit(fn, *args, **kwargs)