cleanup

Author: dimaqq

ops/_main.py

@@ -555,14 +555,11 @@ def main(charm_class: Type[_charm.CharmBase], use_juju_for_storage: Optional[boo
 
     See `ops.main() <#ops-main-entry-point>`_ for details.
     """
-    _tracing.setup_tracing(charm_class.__name__)
-
-    try:
-        with tracer.start_as_current_span('ops.main'):
-            manager = _Manager(charm_class, use_juju_for_storage=use_juju_for_storage)
+    with _tracing.setup_tracing(charm_class.__name__):
+        try:
+            with tracer.start_as_current_span('ops.main'):
+                manager = _Manager(charm_class, use_juju_for_storage=use_juju_for_storage)
 
             manager.run()
-    except _Abort as e:
-        sys.exit(e.exit_code)
-    finally:
-        _tracing.shutdown_tracing()
+        except _Abort as e:
+            sys.exit(e.exit_code)

ops/_tracing/__init__.py

@@ -18,6 +18,9 @@
 
 from __future__ import annotations
 
+from contextlib import contextmanager
+from typing import Generator
+
 import opentelemetry.trace
 
 import ops.version
@@ -30,19 +33,17 @@
         mark_observed,
         set_tracing_destination,
         setup_tracing,
-        shutdown_tracing,
     )
 except ImportError:
 
     def mark_observed() -> None: ...
     def set_tracing_destination(*, url: str | None, ca: str | None = None) -> None: ...
-    def setup_tracing(charm_class_name: str) -> None: ...
-    def shutdown_tracing() -> None: ...
+    @contextmanager
+    def setup_tracing(charm_class_name: str) -> Generator[None, None, None]: yield
 
 
 __all__ = [
     'mark_observed',
     'set_tracing_destination',
     'setup_tracing',
-    'shutdown_tracing',
 ]

ops/_tracing/export.py

@@ -22,7 +22,7 @@
 import urllib.error
 import urllib.request
 from pathlib import Path
-from typing import Sequence
+from typing import Generator, Sequence
 
 import otlp_json
 from opentelemetry.instrumentation.urllib import URLLibInstrumentor
@@ -200,7 +200,8 @@ def suppress_juju_log_handler():
         juju_log_handler.drop.reset(token)
 
 
-def setup_tracing(charm_class_name: str) -> None:
+@contextlib.contextmanager
+def setup_tracing(charm_class_name: str) -> Generator[None, None, None]:
     global _exporter
     # FIXME would it be better to pass Juju context explicitly?
     juju_context = ops.jujucontext._JujuContext.from_dict(os.environ)
@@ -224,6 +225,10 @@ def setup_tracing(charm_class_name: str) -> None:
     span_processor = BatchSpanProcessor(_exporter)
     provider.add_span_processor(span_processor)
     set_tracer_provider(provider)
+    try:
+        yield
+    finally:
+        shutdown_tracing()
     # FIXME: in testing with tracing, we need a hack.
     # OpenTelemetry disallows setting the tracer provider twice,
     # a warning is issued and new provider is ignored.