feat: add time and random for context

Author: brian14708

src/duron/__init__.py

@@ -1,3 +1,2 @@
 from duron.context import Context as Context
-from duron.context import get_context as get_context
 from duron.fn import durable as durable

src/duron/context.py

@@ -1,36 +1,68 @@
 from __future__ import annotations
 
-import asyncio
-from typing import TYPE_CHECKING, TypeVar, cast
+from contextvars import ContextVar
+from random import Random
+from typing import TYPE_CHECKING, TypeVar, cast, final
 
 from typing_extensions import overload
 
-from duron.event_loop import EventLoop
 from duron.ops import FnCall
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
+    from contextvars import Token
+    from inspect import Traceback
+
+    from duron.event_loop import EventLoop
 
     _T = TypeVar("_T")
 
+_context: ContextVar[Context | None] = ContextVar("duron_context", default=None)
+
 
+@final
 class Context:
-    def __init__(self):
-        loop = asyncio.get_event_loop()
-        assert isinstance(loop, EventLoop)
+    def __init__(self, loop: EventLoop) -> None:
         self._loop: EventLoop = loop
-        pass
+        self._token: Token[Context | None] | None = None
+
+    def __enter__(self) -> Context:
+        token = _context.set(self)
+        self._token = token
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Traceback | None,
+    ):
+        if self._token:
+            _context.reset(self._token)
+
+    @classmethod
+    def current(cls) -> Context:
+        ctx = _context.get()
+        if ctx is None:
+            raise RuntimeError("No duron context is active")
+        return ctx
 
     @overload
-    async def run(self, fn: Callable[[], Awaitable[_T]]) -> _T: ...
+    async def run(self, fn: Callable[[], Awaitable[_T]], /) -> _T: ...
     @overload
-    async def run(self, fn: Callable[[], _T]) -> _T: ...
-    async def run(self, fn: Callable[[], Awaitable[_T] | _T]) -> _T:
+    async def run(self, fn: Callable[[], _T], /) -> _T: ...
+    async def run(
+        self,
+        fn: Callable[[], Awaitable[_T] | _T],
+        /,
+    ) -> _T:
         return cast("_T", await self._loop.create_op(FnCall(callable=fn)))
 
+    def time(self) -> float:
+        return self._loop.time()
+
+    def time_ns(self) -> int:
+        return self._loop.time_ns()
 
-def get_context() -> Context:
-    """
-    Get the current duron execution context.
-    """
-    return Context()
+    def random(self) -> Random:
+        return Random(self._loop.generate_op_id())

src/duron/event_loop.py

@@ -87,7 +87,7 @@ def __init__(self, seed: bytes) -> None:
 
         self._timers: list[TimerHandle] = []
 
-    def _generate_id(self) -> bytes:
+    def generate_op_id(self) -> bytes:
         ctx = _task_ctx.get(self._ctx)
         ctx.seq += 1
         return _mix_id(ctx.parent_id, ctx.seq - 1)
@@ -150,6 +150,9 @@ def call_later(
     def time(self) -> float:
         return self._now_ns / 1e9
 
+    def time_ns(self) -> int:
+        return self._now_ns
+
     def tick(self, time: int) -> None:
         self._now_ns = time
 
@@ -166,7 +169,7 @@ def create_task(
         context: Context | None = None,
         **kwargs: Any,
     ) -> Task[_T]:
-        ctx = self._context_new_task(context, self._generate_id())
+        ctx = self._context_new_task(context, self.generate_op_id())
         return ctx.run(
             cast("type[Task[_T]]", Task), coro, name=name, loop=self, **kwargs
         )
@@ -217,7 +220,7 @@ def poll_completion(self, task: Future[_T]) -> WaitSet | None:
             events._set_running_loop(old)
 
     def create_op(self, params: object) -> OpFuture:
-        id = self._generate_id()
+        id = self.generate_op_id()
         s = OpFuture(id, params, self)
         self._ops[id] = s
         return s

src/duron/fn.py

@@ -34,9 +34,11 @@ class DurableFn(Generic[_P, _T_co]):
     fn: Callable[Concatenate[Context, _P], Coroutine[Any, Any, _T_co]]
 
     def __call__(
-        self,
-        log: LogStorage[_TOffset, _TLease],
-    ) -> TaskGuard[_P, _T_co]:
+        self, ctx: Context, *args: _P.args, **kwargs: _P.kwargs
+    ) -> Coroutine[Any, Any, _T_co]:
+        return self.fn(ctx, *args, **kwargs)
+
+    def create_task(self, log: LogStorage[_TOffset, _TLease]) -> TaskGuard[_P, _T_co]:
         return TaskGuard(Task(self, cast("LogStorage[object, object]", log)))
 
 

src/duron/task.py

@@ -17,8 +17,8 @@
     final,
 )
 
-from duron.context import get_context
-from duron.event_loop import create_loop
+from duron.context import Context
+from duron.event_loop import EventLoop, create_loop
 from duron.log import is_entry
 from duron.ops import FnCall, TaskRun
 
@@ -113,14 +113,16 @@ async def _task_prelude(
     task_fn: DurableFn[..., object],
     init: Callable[[], TaskInitParams],
 ) -> object:
-    ctx = get_context()
-    init_params = await ctx.run(init)
-    if init_params["version"] != _CURRENT_VERSION:
-        raise Exception("version mismatch")
-    codec = task_fn.codec
-    args = (codec.decode_json(arg) for arg in init_params["args"])
-    kwargs = {k: codec.decode_json(v) for k, v in init_params["kwargs"].items()}
-    return await task_fn.fn(get_context(), *args, **kwargs)
+    loop = asyncio.get_event_loop()
+    assert isinstance(loop, EventLoop)
+    with Context(loop) as ctx:
+        init_params = await ctx.run(init)
+        if init_params["version"] != _CURRENT_VERSION:
+            raise Exception("version mismatch")
+        codec = task_fn.codec
+        args = (codec.decode_json(arg) for arg in init_params["args"])
+        kwargs = {k: codec.decode_json(v) for k, v in init_params["kwargs"].items()}
+        return await task_fn.fn(ctx, *args, **kwargs)
 
 
 @final

tests/test_task.py

@@ -47,18 +47,18 @@ async def activity(ctx: Context, i: str) -> str:
     }
 
     log = MemoryLogStorage()
-    async with activity(log) as t:
+    async with activity.create_task(log) as t:
         await t.start("test")
         a = await t.wait()
     assert set(e["id"] for e in await log.entries()) == IDS
 
-    async with activity(log) as t:
+    async with activity.create_task(log) as t:
         await t.start("test")
         b = await t.wait()
     assert a == b
 
     log2 = MemoryLogStorage((await log.entries())[:-2])
-    async with activity(log2) as t:
+    async with activity.create_task(log2) as t:
         await t.start("test")
         c = await t.wait()
     assert a == c
@@ -78,11 +78,11 @@ async def error():
 
     log = MemoryLogStorage()
     with pytest.raises(check=lambda v: "test error" in str(v)):
-        async with activity(log) as t:
+        async with activity.create_task(log) as t:
             await t.start()
             await t.wait()
     with pytest.raises(check=lambda v: "test error" in str(v)):
-        async with activity(log) as t:
+        async with activity.create_task(log) as t:
             await t.start()
             await t.wait()
 
@@ -97,12 +97,12 @@ async def activity(ctx: Context, s: str) -> str:
         return s
 
     log = MemoryLogStorage()
-    async with activity(log) as t:
+    async with activity.create_task(log) as t:
         await t.start("hello")
         with pytest.raises(asyncio.TimeoutError):
             _ = await asyncio.wait_for(t.wait(), 0.1)
 
-    async with activity(log) as t:
+    async with activity.create_task(log) as t:
         sleep = 0
         await t.resume()
         x = await t.wait()
@@ -119,7 +119,7 @@ async def activity(ctx: Context, s: str) -> str:
         return s
 
     log = MemoryLogStorage()
-    async with activity(log) as t:
+    async with activity.create_task(log) as t:
         await t.start("hello")
         try:
             _ = await t.wait()
@@ -143,8 +143,28 @@ async def activity(ctx: Context) -> CustomPoint:
         return CustomPoint(x=pt.x + 5, y=pt.y + 10)
 
     log = MemoryLogStorage()
-    async with activity(log) as t:
+    async with activity.create_task(log) as t:
         await t.start()
         a = await t.wait()
         assert type(a) is CustomPoint
         assert a.x == 6 and a.y == 12
+
+
+@pytest.mark.asyncio
+async def test_random():
+    @durable()
+    async def activity(ctx: Context) -> int:
+        assert ctx.time_ns() == ctx.time_ns()
+        return ctx.random().randint(1, 100)
+
+    log = MemoryLogStorage()
+    async with activity.create_task(log) as t:
+        await t.start()
+        a = await t.wait()
+
+    log = MemoryLogStorage()
+    async with activity.create_task(log) as t:
+        await t.start()
+        b = await t.wait()
+
+    assert a == b