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

@@ -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 current_tracer
 from duron.typing import Unspecified, inspect_function
 
 if TYPE_CHECKING:
@@ -53,6 +54,7 @@
         StreamCreateEntry,
         StreamEmitEntry,
     )
+    from duron.tracing import Tracer
     from duron.typing import FunctionType, TypeHint
 
 
@@ -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()
 
@@ -286,7 +286,7 @@ async def _invoke_prelude(
 class _InvokeRun:
     __slots__ = (
         "_codec",
-        "_debug",
+        "_lease",
         "_log",
         "_loop",
         "_now",
@@ -297,6 +297,7 @@ class _InvokeRun:
         "_streams",
         "_task",
         "_tasks",
+        "_tracer",
         "_watchers",
     )
 
@@ -310,15 +311,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 +335,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 = current_tracer()
 
     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 +359,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 +379,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 +492,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 +520,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 +590,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 +633,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

@@ -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/log.py

@@ -18,7 +18,7 @@ class _BaseEntry(TypedDict):
     id: str
     # Unix timestamp in microseconds
     ts: int
-    debug: NotRequired[dict[str, JSONValue]]
+    metadata: NotRequired[dict[str, JSONValue]]
 
 
 class ErrorInfo(TypedDict):
@@ -28,7 +28,6 @@ class ErrorInfo(TypedDict):
 
 class PromiseCreateEntry(_BaseEntry):
     type: Literal["promise/create"]
-    metadata: NotRequired[dict[str, JSONValue]]
 
 
 class PromiseCompleteEntry(_BaseEntry):
@@ -40,7 +39,6 @@ class PromiseCompleteEntry(_BaseEntry):
 
 class StreamCreateEntry(_BaseEntry):
     type: Literal["stream/create"]
-    metadata: NotRequired[dict[str, JSONValue]]
 
 
 class StreamEmitEntry(_BaseEntry):
@@ -59,6 +57,10 @@ class BarrierEntry(_BaseEntry):
     type: Literal["barrier"]
 
 
+class TraceEntry(_BaseEntry):
+    type: Literal["trace"]
+
+
 class AnyEntry(_BaseEntry):
     type: str
 
@@ -70,6 +72,7 @@ class AnyEntry(_BaseEntry):
     | StreamEmitEntry
     | StreamCompleteEntry
     | BarrierEntry
+    | TraceEntry
 )
 
 
@@ -84,6 +87,16 @@ def is_entry(entry: Entry | AnyEntry) -> TypeGuard[Entry]:
     }
 
 
+def set_metadata(entry: Entry, *metadata: dict[str, JSONValue] | None) -> None:
+    m = entry.get("metadata")
+    for md in metadata:
+        if md:
+            if m is None:
+                entry["metadata"] = m = {**md}
+            else:
+                m.update(md)
+
+
 class LogStorage(Protocol):
     __slots__: tuple[str, ...] = ()