feat: use tracer instead of debug

Author: brian14708

examples/hello_world.py

@@ -3,10 +3,12 @@
 import asyncio
 import random
 import sys
+import uuid
 from pathlib import Path
 
 import duron
 from duron.contrib.storage import FileLogStorage
+from duron.tracing import LogTracer, use_tracer
 
 
 @duron.op
@@ -34,7 +36,11 @@ async def greeting_flow(ctx: duron.Context, name: str) -> str:
 
 
 async def run_workflow(name: str, log_file: Path) -> str:
-    async with greeting_flow.invoke(FileLogStorage(log_file)) as job:
+    log = FileLogStorage(log_file)
+    async with (
+        use_tracer(LogTracer(str(uuid.uuid4()), log)),
+        greeting_flow.invoke(log) as job,
+    ):
         await job.start(name)
         return await job.wait()
 

src/duron/_core/config.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -13,21 +12,16 @@
 @dataclass(slots=True)
 class _Config:
     codec: Codec
-    debug: bool
 
 
 config = _Config(
     codec=DefaultCodec(),
-    debug=os.getenv("DURON_DEBUG", "0").lower() in {"1", "true"},
 )
 
 
 def set_config(
     *,
     codec: Codec | None = None,
-    debug: bool | None = None,
 ) -> None:
     if codec is not None:
         config.codec = codec
-    if debug is not None:
-        config.debug = debug

src/duron/_core/invoke.py

@@ -14,7 +14,6 @@
     overload,
 )
 
-from duron._core.config import config
 from duron._core.context import Context
 from duron._core.ops import (
     Barrier,
@@ -29,6 +28,7 @@
 from duron._loop import EventLoop, create_loop
 from duron.codec import Codec, JSONValue
 from duron.log import derive_id, is_entry, random_id
+from duron.tracing import current_tracer
 from duron.typing import Unspecified, inspect_function
 
 if TYPE_CHECKING:
@@ -106,7 +106,6 @@ def get_init() -> InitParams:
             self._log,
             codec,
             watchers=self._watchers,
-            debug=config.debug,
         )
         await self._run.resume()
 
@@ -122,7 +121,6 @@ def cb() -> InitParams:
             self._log,
             self._fn.codec,
             watchers=self._watchers,
-            debug=config.debug,
         )
         await self._run.resume()
 
@@ -287,6 +285,7 @@ class _InvokeRun:
     __slots__ = (
         "_codec",
         "_debug",
+        "_lease",
         "_log",
         "_loop",
         "_now",
@@ -310,15 +309,13 @@ def __init__(
             tuple[Callable[[dict[str, JSONValue]], bool], StreamObserver[object]]
         ]
         | None = None,
-        debug: bool = False,
     ) -> None:
         self._loop = create_loop(asyncio.get_running_loop())
-        if debug:
-            self._loop.set_debug(True)
         self._task = self._loop.create_task(task)
         self._log = log
         self._codec = codec
-        self._running: bytes | None = None
+        self._running: bool = False
+        self._lease: bytes | None = None
         self._pending_msg: list[Entry] = []
         self._pending_task: dict[
             str,
@@ -336,11 +333,15 @@ def __init__(
             ],
         ] = {}
         self._watchers = watchers or []
+        tracer = current_tracer()
         self._debug: dict[str, JSONValue] | None = (
-            {"run.id": random_id()} if debug else None
+            {"trace.id": tracer.id()} if tracer else None
         )
 
     async def close(self) -> None:
+        if self._lease:
+            await self._log.release_lease(self._lease)
+            self._lease = None
         for task, _ in self._tasks.values():
             _ = task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
@@ -359,6 +360,7 @@ def now(self) -> int:
         return self._now
 
     async def resume(self) -> None:
+        self._lease = await self._log.acquire_lease()
         recvd_msgs: set[str] = set()
         async for o, entry in self._log.stream(None, live=False):
             ts = entry["ts"]
@@ -378,21 +380,17 @@ async def run(self) -> object:
         if self._task.done():
             return self._task.result()
 
-        self._running = await self._log.acquire_lease()
-        try:
-            for msg in self._pending_msg:
-                await self.enqueue_log(msg)
-            self._pending_msg.clear()
-            for key, (task_fn, return_type) in self._pending_task.items():
-                self._tasks[key] = (asyncio.create_task(task_fn()), return_type)
-            self._pending_task.clear()
-
-            while waitset := await self._step():
-                await waitset.block(self.now())
-            return self._task.result()
-        finally:
-            await self._log.release_lease(self._running)
-            self._running = None
+        self._running = True
+        for msg in self._pending_msg:
+            await self.enqueue_log(msg)
+        self._pending_msg.clear()
+        for key, (task_fn, return_type) in self._pending_task.items():
+            self._tasks[key] = (asyncio.create_task(task_fn()), return_type)
+        self._pending_task.clear()
+
+        while waitset := await self._step():
+            await waitset.block(self.now())
+        return self._task.result()
 
     async def _step(self) -> WaitSet | None:
         while True:
@@ -499,6 +497,9 @@ async def handle_message(
     async def enqueue_log(self, entry: Entry, *, flush: bool = False) -> None:
         if not self._running:
             self._pending_msg.append(entry)
+        elif self._lease is None:
+            # closed
+            return
         else:
             if self._debug:
                 if "debug" in entry:
@@ -510,9 +511,9 @@ async def enqueue_log(self, entry: Entry, *, flush: bool = False) -> None:
                     entry["debug"] = self._debug
             else:
                 _ = entry.pop("debug", None)
-            offset = await self._log.append(self._running, entry)
+            offset = await self._log.append(self._lease, entry)
             if flush:
-                await self._log.flush(self._running)
+                await self._log.flush(self._lease)
             await self.handle_message(offset, entry)
 
     async def enqueue_op(self, id_: str, fut: OpFuture[object]) -> None:

src/duron/_decorator/fn.py

@@ -50,7 +50,11 @@ def __call__(
     ) -> Coroutine[Any, Any, _T_co]:
         return self.fn(ctx, *args, **kwargs)
 
-    def invoke(self, log: LogStorage) -> AsyncContextManager[Invoke[_P, _T_co]]:
+    def invoke(
+        self,
+        log: LogStorage,
+        /,
+    ) -> AsyncContextManager[Invoke[_P, _T_co]]:
         return Invoke[_P, _T_co].invoke(self, log)
 
 

src/duron/_loop.py

@@ -83,7 +83,6 @@ class EventLoop(asyncio.AbstractEventLoop):
     __slots__: tuple[str, ...] = (
         "_closed",
         "_ctx",
-        "_debug",
         "_event",
         "_exc_handler",
         "_host",
@@ -95,7 +94,6 @@ class EventLoop(asyncio.AbstractEventLoop):
 
     def __init__(self, host: asyncio.AbstractEventLoop) -> None:
         self._ready: deque[asyncio.Handle] = deque()
-        self._debug: bool = False
         self._host: asyncio.AbstractEventLoop = host
         self._exc_handler: (
             Callable[[asyncio.AbstractEventLoop, dict[str, object]], object] | None
@@ -317,11 +315,7 @@ def close(self) -> None:
 
     @override
     def get_debug(self) -> bool:
-        return self._debug
-
-    @override
-    def set_debug(self, enabled: bool) -> None:
-        self._debug = enabled
+        return False
 
     @override
     def default_exception_handler(self, context: dict[str, object]) -> None:

src/duron/tracing/__init__.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Final, Protocol, final
+
+if TYPE_CHECKING:
+    from contextvars import Token
+    from types import TracebackType
+
+    from duron.log import LogStorage
+
+_current_tracer: ContextVar[Tracer | None] = ContextVar("duron_tracer", default=None)
+
+
+class Tracer(Protocol):
+    def id(self) -> str: ...
+
+    def enabled(self) -> bool: ...
+    def __bool__(self) -> bool: ...
+
+
+@final
+class _NoopTracer:
+    __slots__: tuple[str, ...] = ()
+
+    @staticmethod
+    def id() -> str:
+        return ""
+
+    @staticmethod
+    def enabled() -> bool:
+        return False
+
+    @staticmethod
+    def __bool__() -> bool:
+        return False
+
+
+_Noop: Final = _NoopTracer()
+
+
+class LogTracer:
+    __slots__: tuple[str, ...] = ("_id", "_log")
+
+    def __init__(
+        self,
+        id_: str,
+        log: LogStorage,
+        /,
+    ) -> None:
+        self._log: LogStorage = log
+        self._id: str = id_
+
+    def id(self) -> str:
+        return self._id
+
+    @staticmethod
+    def enabled() -> bool:
+        return True
+
+    @staticmethod
+    def __bool__() -> bool:
+        return True
+
+
+def current_tracer() -> Tracer:
+    return _current_tracer.get() or _Noop
+
+
+class _TracerGuard:
+    __slots__ = ("_token", "_value")
+
+    def __init__(self, value: Tracer | None) -> None:
+        self._value = value
+        self._token: Token[Tracer | None] | None = None
+
+    def __enter__(self) -> None:
+        self._token = _current_tracer.set(self._value)
+
+    def __exit__(
+        self,
+        _exc_type: type[BaseException] | None,
+        _exc_value: BaseException | None,
+        _traceback: TracebackType | None,
+    ) -> None:
+        if self._token is not None:
+            _current_tracer.reset(self._token)
+            self._token = None
+
+    async def __aenter__(self) -> None:
+        self._token = _current_tracer.set(self._value)
+
+    async def __aexit__(
+        self,
+        _exc_type: type[BaseException] | None,
+        _exc_value: BaseException | None,
+        _traceback: TracebackType | None,
+    ) -> None:
+        if self._token is not None:
+            _current_tracer.reset(self._token)
+            self._token = None
+
+
+def use_tracer(tracer: Tracer | None) -> _TracerGuard:
+    return _TracerGuard(tracer)