feat: move id encoding to log

Author: brian14708

src/duron/_core/context.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import asyncio
-import binascii
 from contextlib import contextmanager
 from contextvars import ContextVar
 from random import Random
@@ -19,6 +18,7 @@
 from duron._core.signal import create_signal
 from duron._core.stream import create_stream, run_stream
 from duron._decorator.op import CheckpointOp, Op
+from duron.log import encode_id
 from duron.typing import inspect_function
 
 if TYPE_CHECKING:
@@ -185,7 +185,7 @@ async def create_promise(
             ExternalPromiseCreate(metadata=self._get_metadata(None), return_type=dtype),
         )
         return (
-            binascii.b2a_base64(fut.id, newline=False).decode(),
+            encode_id(fut.id),
             cast("asyncio.Future[_T]", fut),
         )
 

src/duron/_core/invoke.py

@@ -1,11 +1,8 @@
 from __future__ import annotations
 
 import asyncio
-import binascii
 import contextlib
-import os
 import time
-from hashlib import blake2b
 from typing import TYPE_CHECKING, Generic, Literal, cast
 from typing_extensions import (
     Any,
@@ -31,7 +28,7 @@
 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 is_entry
+from duron.log import decode_id, derive_id, encode_id, is_entry, random_id
 from duron.typing import Unspecified, inspect_function
 
 if TYPE_CHECKING:
@@ -257,7 +254,7 @@ async def _invoke_prelude(
         if init_params["version"] != _CURRENT_VERSION:
             msg = "version mismatch"
             raise RuntimeError(msg)
-        loop.set_key(_decode_id(init_params["seed"]))
+        loop.set_key(decode_id(init_params["seed"]))
         extra_kwargs: dict[str, object] = {}
         for name, type_, dtype in job_fn.inject:
             with ctx.metadata({"param.name": name}):
@@ -423,7 +420,7 @@ async def handle_message(
             pending_info = self._pending_task.pop(e["promise_id"], None)
             task_info = self._tasks.get(e["promise_id"], None)
 
-            id_ = _decode_id(e["promise_id"])
+            id_ = decode_id(e["promise_id"])
 
             return_type: TypeHint[Any] = Unspecified
             if pending_info is not None:
@@ -454,7 +451,7 @@ async def handle_message(
                 msg = f"Invalid promise/complete entry: {e!r}"
                 raise ValueError(msg)
         elif e["type"] == "stream/create":
-            id_ = _decode_id(e["id"])
+            id_ = decode_id(e["id"])
             if e["id"] not in self._streams:
                 self._loop.post_completion(
                     id_, exception=ValueError("Stream not found")
@@ -463,7 +460,7 @@ async def handle_message(
                 self._loop.post_completion(id_, result=e["id"])
             self._pending_ops.discard(id_)
         elif e["type"] == "stream/emit":
-            id_ = _decode_id(e["id"])
+            id_ = decode_id(e["id"])
             if e["stream_id"] not in self._streams:
                 self._loop.post_completion(
                     id_, exception=ValueError("Stream not found")
@@ -478,7 +475,7 @@ async def handle_message(
                 self._loop.post_completion(id_, result=None)
             self._pending_ops.discard(id_)
         elif e["type"] == "stream/complete":
-            id_ = _decode_id(e["id"])
+            id_ = decode_id(e["id"])
             if e["stream_id"] not in self._streams:
                 self._loop.post_completion(
                     id_, exception=ValueError("Stream not found")
@@ -495,7 +492,7 @@ async def handle_message(
                 _ = self._streams.pop(e["stream_id"], None)
             self._pending_ops.discard(id_)
         elif e["type"] == "barrier":
-            id_ = _decode_id(e["id"])
+            id_ = decode_id(e["id"])
             self._loop.post_completion(id_, result=offset)
             self._pending_ops.discard(id_)
 
@@ -524,7 +521,7 @@ async def enqueue_op(self, id_: bytes, fut: OpFuture[object]) -> None:
             case FnCall():
                 promise_create_entry: PromiseCreateEntry = {
                     "ts": self.now(),
-                    "id": _encode_id(id_),
+                    "id": encode_id(id_),
                     "type": "promise/create",
                 }
                 if op.metadata:
@@ -538,11 +535,12 @@ async def enqueue_op(self, id_: bytes, fut: OpFuture[object]) -> None:
                 await self.enqueue_log(promise_create_entry)
 
                 async def cb() -> None:
+                    now_us = self.now()
                     entry: PromiseCompleteEntry = {
-                        "ts": self.now(),
-                        "id": _encode_id(id_, ack=True),
+                        "ts": now_us,
+                        "id": encode_id(derive_id(id_)),
                         "type": "promise/complete",
-                        "promise_id": _encode_id(id_),
+                        "promise_id": encode_id(id_),
                     }
                     try:
                         result = op.callable(*op.args, **op.kwargs)
@@ -558,7 +556,7 @@ async def cb() -> None:
 
                 def done(f: OpFuture[object]) -> None:
                     if f.cancelled():
-                        sid = _encode_id(f.id)
+                        sid = encode_id(f.id)
                         if self._pending_task.get(sid, None):
                             # pending task cancelled
                             pass
@@ -568,14 +566,14 @@ def done(f: OpFuture[object]) -> None:
                                 _ = task.get_loop().call_soon(task.cancel)
 
                 fut.add_done_callback(done)
-                sid = _encode_id(id_)
+                sid = encode_id(id_)
                 if self._running:
                     self._tasks[sid] = (asyncio.create_task(cb()), op.return_type)
                 else:
                     self._pending_task[sid] = (cb, op.return_type)
 
             case StreamCreate():
-                stream_id = _encode_id(id_)
+                stream_id = encode_id(id_)
 
                 # Determine which observer to use
                 ob = [op.observer] if op.observer else []
@@ -599,7 +597,7 @@ def done(f: OpFuture[object]) -> None:
             case StreamEmit():
                 stream_emit_entry: StreamEmitEntry = {
                     "ts": self.now(),
-                    "id": _encode_id(id_),
+                    "id": encode_id(id_),
                     "stream_id": op.stream_id,
                     "type": "stream/emit",
                     "value": self._codec.encode_json(op.value),
@@ -609,7 +607,7 @@ def done(f: OpFuture[object]) -> None:
                 if op.exception:
                     stream_close_entry_err: StreamCompleteEntry = {
                         "ts": self.now(),
-                        "id": _encode_id(id_),
+                        "id": encode_id(id_),
                         "stream_id": op.stream_id,
                         "type": "stream/complete",
                         "error": _encode_error(op.exception),
@@ -618,27 +616,27 @@ def done(f: OpFuture[object]) -> None:
                 else:
                     stream_close_entry: StreamCompleteEntry = {
                         "ts": self.now(),
-                        "id": _encode_id(id_),
+                        "id": encode_id(id_),
                         "stream_id": op.stream_id,
                         "type": "stream/complete",
                     }
                     await self.enqueue_log(stream_close_entry)
             case Barrier():
                 barrier_entry: BarrierEntry = {
                     "ts": self.now(),
-                    "id": _encode_id(id_),
+                    "id": encode_id(id_),
                     "type": "barrier",
                 }
                 await self.enqueue_log(barrier_entry, flush=True)
             case ExternalPromiseCreate():
                 promise_create_entry = {
                     "ts": self.now(),
-                    "id": _encode_id(id_),
+                    "id": encode_id(id_),
                     "type": "promise/create",
                 }
                 if op.metadata:
                     promise_create_entry["metadata"] = op.metadata
-                self._tasks[_encode_id(id_)] = (asyncio.Future(), op.return_type)
+                self._tasks[encode_id(id_)] = (asyncio.Future(), op.return_type)
                 await self.enqueue_log(promise_create_entry)
             case _:
                 assert_never(op)
@@ -653,9 +651,10 @@ async def complete_external_promise(
         if id_ not in self._tasks:
             msg = "Promise not found"
             raise ValueError(msg)
+        now_us = self.now()
         entry: PromiseCompleteEntry = {
-            "ts": self.now(),
-            "id": _encode_id(_decode_id(id_), ack=True),
+            "ts": now_us,
+            "id": encode_id(derive_id(decode_id(id_))),
             "type": "promise/complete",
             "promise_id": id_,
         }
@@ -723,21 +722,8 @@ async def close_stream(
         return cnt
 
 
-def _encode_id(id_: bytes, *, ack: bool = False) -> str:
-    if ack:
-        id_ = blake2b(
-            id_,
-            digest_size=12,
-        ).digest()
-    return binascii.b2a_base64(id_, newline=False).decode()
-
-
-def _decode_id(encoded: str) -> bytes:
-    return binascii.a2b_base64(encoded)
-
-
 def _generate_id() -> str:
-    return _encode_id(os.urandom(12))
+    return encode_id(random_id())
 
 
 def _encode_error(error: BaseException) -> ErrorInfo:

src/duron/_loop.py

@@ -4,11 +4,9 @@
 import contextlib
 import contextvars
 import logging
-import os
 from asyncio import events, tasks
 from collections import deque
 from dataclasses import dataclass
-from hashlib import blake2b
 from heapq import heappop, heappush
 from typing import TYPE_CHECKING, Generic
 from typing_extensions import (
@@ -20,6 +18,8 @@
     override,
 )
 
+from duron.log import derive_id, random_id
+
 _T = TypeVar("_T")
 
 if TYPE_CHECKING:
@@ -108,15 +108,14 @@ def __init__(self, host: asyncio.AbstractEventLoop) -> None:
         self._timers: list[asyncio.TimerHandle] = []
 
     def set_key(self, key: bytes) -> None:
-        self._key = key
+        self._key = derive_id(key, key=self._key)
 
     def generate_op_id(self) -> bytes:
         ctx = _task_ctx.get(self._ctx)
         ctx.seq += 1
-        return _mix_id(ctx.parent_id, self._key, ctx.seq - 1)
-
-    def host_loop(self) -> asyncio.AbstractEventLoop:
-        return self._host
+        return derive_id(
+            ctx.parent_id, context=(ctx.seq - 1).to_bytes(4, "big"), key=self._key
+        )
 
     @override
     def call_soon(
@@ -256,7 +255,7 @@ def poll_completion(self, task: Future[_T]) -> WaitSet | None:
 
     def create_op(self, params: object, *, external: bool = False) -> OpFuture[object]:
         if external:
-            id_ = os.urandom(12)
+            id_ = random_id()
             self._event.set()
         else:
             id_ = self.generate_op_id()
@@ -288,7 +287,7 @@ def post_completion(
         if op := self._ops.pop(id_, None):
             if op.done():
                 return
-            tid = _mix_id(op.id, self._key, -1)
+            tid = derive_id(op.id, key=self._key)
             token = _task_ctx.set(_TaskCtx(parent_id=tid))
             if exception is None:
                 _ = self.call_soon(op.set_result, result)
@@ -356,12 +355,6 @@ def _timer_handle_cancelled(self, _th: asyncio.TimerHandle) -> None:
         pass
 
 
-def _mix_id(a: bytes, key: bytes, b: int) -> bytes:
-    if b == -1:
-        return blake2b(a, key=key, digest_size=12).digest()
-    return blake2b(b.to_bytes(4, "little") + a, key=key, digest_size=12).digest()
-
-
 def create_loop(
     parent_loop: asyncio.AbstractEventLoop,
 ) -> EventLoop:

src/duron/log.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import binascii
+import os
 from abc import ABC, abstractmethod
+from hashlib import blake2b
 from typing import TYPE_CHECKING, Literal
 from typing_extensions import NotRequired, TypedDict
 
@@ -102,3 +105,19 @@ async def append(self, token: bytes, entry: Entry, /) -> int: ...
 
     @abstractmethod
     async def flush(self, token: bytes, /) -> None: ...
+
+
+def encode_id(raw: bytes) -> str:
+    return binascii.b2a_base64(raw, newline=False).decode()
+
+
+def decode_id(encoded: str) -> bytes:
+    return binascii.a2b_base64(encoded)
+
+
+def random_id() -> bytes:
+    return os.urandom(12)
+
+
+def derive_id(base: bytes, *, context: bytes = b"", key: bytes = b"") -> bytes:
+    return blake2b(base, salt=context, key=key, digest_size=12).digest()

tests/test_id.py

@@ -0,0 +1,28 @@
+import pytest
+
+from duron.log import (
+    decode_id,
+    derive_id,
+    encode_id,
+    random_id,
+)
+
+
+def test_generates_unique_ids() -> None:
+    ids = {random_id() for _ in range(1000)}
+    assert len(ids) == 1000
+
+
+def test_derive_id_deterministic() -> None:
+    base = b"test base"
+    key = b"test key"
+
+    id1 = derive_id(base, key=key)
+    id2 = derive_id(base, key=key)
+    assert id1 == id2
+    assert decode_id(encode_id(id1)) == id2
+
+
+@pytest.mark.benchmark
+def test_bench_derive_id() -> None:
+    _ = decode_id(encode_id(derive_id(b"hello", context=b"key")))