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 Tracer
 
 
 @duron.op
@@ -34,7 +36,13 @@ 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 (
+        greeting_flow.invoke(
+            log,
+            tracer=Tracer(str(uuid.uuid4())),
+        ) 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

@@ -10,11 +10,11 @@
     TypedDict,
     TypeVar,
     assert_never,
+    assert_type,
     final,
     overload,
 )
 
-from duron._core.config import config
 from duron._core.context import Context
 from duron._core.ops import (
     Barrier,
@@ -28,7 +28,8 @@
 from duron._core.stream import ObserverStream, Stream, StreamWriter
 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.log import derive_id, is_entry, random_id, set_metadata
+from duron.tracing import Tracer, use_tracer
 from duron.typing import Unspecified, inspect_function
 
 if TYPE_CHECKING:
@@ -86,8 +87,9 @@ def __init__(
     def invoke(
         fn: Fn[_P, _T_co],
         log: LogStorage,
+        tracer: Tracer | None,
     ) -> contextlib.AbstractAsyncContextManager[Invoke[_P, _T_co]]:
-        return _InvokeGuard(Invoke(fn, log))
+        return _InvokeGuard(Invoke(fn, log), tracer)
 
     async def start(self, *args: _P.args, **kwargs: _P.kwargs) -> None:
         def get_init() -> InitParams:
@@ -106,7 +108,6 @@ def get_init() -> InitParams:
             self._log,
             codec,
             watchers=self._watchers,
-            debug=config.debug,
         )
         await self._run.resume()
 
@@ -122,7 +123,6 @@ def cb() -> InitParams:
             self._log,
             self._fn.codec,
             watchers=self._watchers,
-            debug=config.debug,
         )
         await self._run.resume()
 
@@ -217,12 +217,14 @@ def get_run(self) -> _InvokeRun:
 
 @final
 class _InvokeGuard(Generic[_P, _T_co]):
-    __slots__ = ("_job",)
+    __slots__ = ("_job", "_tracer")
 
-    def __init__(self, job: Invoke[_P, _T_co]) -> None:
+    def __init__(self, job: Invoke[_P, _T_co], tracer: Tracer | None) -> None:
         self._job = job
+        self._tracer = use_tracer(tracer)
 
     async def __aenter__(self) -> Invoke[_P, _T_co]:
+        self._tracer.__enter__()
         return self._job
 
     async def __aexit__(
@@ -232,6 +234,7 @@ async def __aexit__(
         traceback: TracebackType | None,
     ) -> None:
         await self._job.close()
+        self._tracer.__exit__(exc_type, exc_value, traceback)
 
 
 class InitParams(TypedDict):
@@ -286,7 +289,7 @@ async def _invoke_prelude(
 class _InvokeRun:
     __slots__ = (
         "_codec",
-        "_debug",
+        "_lease",
         "_log",
         "_loop",
         "_now",
@@ -297,6 +300,7 @@ class _InvokeRun:
         "_streams",
         "_task",
         "_tasks",
+        "_tracer",
         "_watchers",
     )
 
@@ -310,15 +314,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 +338,12 @@ def __init__(
             ],
         ] = {}
         self._watchers = watchers or []
-        self._debug: dict[str, JSONValue] | None = (
-            {"run.id": random_id()} if debug else None
-        )
+        self._tracer: Tracer | None = Tracer.current()
 
     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 +362,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 +382,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:
@@ -495,24 +495,23 @@ async def handle_message(
             id_ = e["id"]
             self._loop.post_completion(id_, result=offset)
             self._pending_ops.discard(id_)
+        elif e["type"] == "trace":
+            pass
+        else:
+            assert_type(e["type"], Literal["promise/create"])
 
     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:
-                    entry["debug"] = {
-                        **entry["debug"],
-                        **self._debug,
-                    }
-                else:
-                    entry["debug"] = self._debug
-            else:
-                _ = entry.pop("debug", None)
-            offset = await self._log.append(self._running, entry)
+            if self._tracer:
+                self._tracer.attach_metadata(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:
@@ -524,14 +523,18 @@ async def enqueue_op(self, id_: str, fut: OpFuture[object]) -> None:
                     "id": id_,
                     "type": "promise/create",
                 }
-                if op.metadata:
-                    promise_create_entry["metadata"] = op.metadata
-                if self._debug:
-                    promise_create_entry["debug"] = {
+
+                set_metadata(
+                    promise_create_entry,
+                    op.metadata,
+                    {
                         "fn.name": str(
                             getattr(op.callable, "__qualname__", op.callable)
                         ),
                     }
+                    if self._tracer
+                    else None,
+                )
                 await self.enqueue_log(promise_create_entry)
 
                 async def cb() -> None:
@@ -590,8 +593,7 @@ def done(f: OpFuture[object]) -> None:
                     "id": stream_id,
                     "type": "stream/create",
                 }
-                if op.metadata:
-                    stream_create_entry["metadata"] = op.metadata
+                set_metadata(stream_create_entry, op.metadata)
                 await self.enqueue_log(stream_create_entry)
 
             case StreamEmit():
@@ -634,8 +636,7 @@ def done(f: OpFuture[object]) -> None:
                     "id": id_,
                     "type": "promise/create",
                 }
-                if op.metadata:
-                    promise_create_entry["metadata"] = op.metadata
+                set_metadata(promise_create_entry, op.metadata)
                 self._tasks[id_] = (asyncio.Future(), op.return_type)
                 await self.enqueue_log(promise_create_entry)
             case _:

src/duron/_decorator/fn.py

@@ -29,6 +29,7 @@
     from duron._core.context import Context
     from duron.codec import Codec
     from duron.log import LogStorage
+    from duron.tracing import Tracer
     from duron.typing import TypeHint
 
 
@@ -50,8 +51,10 @@ def __call__(
     ) -> Coroutine[Any, Any, _T_co]:
         return self.fn(ctx, *args, **kwargs)
 
-    def invoke(self, log: LogStorage) -> AsyncContextManager[Invoke[_P, _T_co]]:
-        return Invoke[_P, _T_co].invoke(self, log)
+    def invoke(
+        self, log: LogStorage, /, *, tracer: Tracer | None = None
+    ) -> AsyncContextManager[Invoke[_P, _T_co]]:
+        return Invoke[_P, _T_co].invoke(self, log, tracer)
 
 
 @overload

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: