mongo/buildscripts/resmokelib/utils/otel_thread_pool_executor.py
Juan Gu 855dfadef0 SERVER-94077 Use isort in Ruff configs (#27865)
GitOrigin-RevId: e793d662774ccd3ab6c3f356c2287cf1f7ff9805
2024-10-10 19:33:49 +00:00

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)