use context manager @pcrespov @sanderegg

Author: bisgaard-itis

packages/service-library/src/servicelib/async_utils.py

@@ -164,19 +164,17 @@ async def worker(in_q: Queue[QueueElement], out_q: Queue) -> None:
                     while True:
                         element = await in_q.get()
                         in_q.task_done()
-                        tracing.attach_context(element.tracing_context)
-                        # check if requested to shutdown
-                        try:
-                            do_profile = element.do_profile
-                            awaitable = element.input
-                            if awaitable is None:
-                                break
-                            with profile_context(do_profile):
-                                result = await awaitable
-                        except Exception as e:  # pylint: disable=broad-except
-                            result = e
-                        finally:
-                            tracing.detach_context(element.tracing_context)
+                        with tracing.use_tracing_context(element.tracing_context):
+                            # check if requested to shutdown
+                            try:
+                                do_profile = element.do_profile
+                                awaitable = element.input
+                                if awaitable is None:
+                                    break
+                                with profile_context(do_profile):
+                                    result = await awaitable
+                            except Exception as e:  # pylint: disable=broad-except
+                                result = e
                         await out_q.put(result)
 
                     logging.info(

packages/service-library/src/servicelib/tracing.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from typing import TypeAlias
 
 from opentelemetry import context as otcontext
@@ -18,14 +19,15 @@ def get_context() -> TracingContext:
     return otcontext.get_current()
 
 
-def attach_context(context: TracingContext) -> None:
+@contextmanager
+def use_tracing_context(context: TracingContext):
     if context is not None:
         otcontext.attach(context)
-
-
-def detach_context(context: TracingContext) -> None:
-    if context is not None:
-        otcontext.detach(context)
+    try:
+        yield
+    finally:
+        if context is not None:
+            otcontext.detach(context)
 
 
 def setup_log_tracing(tracing_settings: TracingSettings):