diff --git a/src/nu/lang/runtime/runtime.py b/src/nu/lang/runtime/runtime.py index 7f85faeb..1b2af841 100755 --- a/src/nu/lang/runtime/runtime.py +++ b/src/nu/lang/runtime/runtime.py @@ -57,6 +57,18 @@ _RT_CTX: contextvars.ContextVar[Context] = contextvars.ContextVar("nu_rt_ctx") +def _carry_ctx() -> Callable[..., object]: + """Return a runner that calls a function inside a copy of the caller's context. + + A worker thread starts with an empty contextvars context, so `_RT_CTX` is + unset there and `Runtime.ctx` raises `LookupError`. Taking the copy on the + calling side and submitting `copy.run` keeps the Context resolvable on the + worker. A fresh copy per branch keeps a `.set()` inside that branch local to + it, the same copy-on-write rule an asyncio.Task gets. + """ + return contextvars.copy_context().run + + class Runtime: """Per-drive Runtime. Owns a Program, a per-task Context, and a Budget.""" @@ -118,7 +130,7 @@ def eval_parallel(self, nids: Iterable[int]) -> list: if self.budget.max_parallel == 1 or self.budget.thread_pool is None: return self.eval_each(nids) pool = self.budget.thread_pool - futures = [pool.submit(self.eval, n) for n in nids] + futures = [pool.submit(_carry_ctx(), self.eval, n) for n in nids] return [f.result() for f in futures] def _drive_async(self, nids: list[int]) -> list: @@ -147,7 +159,7 @@ async def place(n: int) -> object: async with sem: if on_loop_col[n]: return await self.aeval(n) - return await loop.run_in_executor(pool, self.eval, n) + return await loop.run_in_executor(pool, _carry_ctx(), self.eval, n) return [place(n) for n in nids] @@ -261,7 +273,7 @@ def drain(n: int) -> None: finally: q.put(_DONE) - futures = [pool.submit(drain, n) for n in nids] + futures = [pool.submit(_carry_ctx(), drain, n) for n in nids] remaining = len(futures) try: while remaining > 0: @@ -306,7 +318,7 @@ def in_thread(self, fn: Callable, *args: object, **kwargs: object) -> Future: if self.budget.thread_pool is None: msg = "in_thread requires max_parallel > 1" raise RuntimeError(msg) - return self.budget.thread_pool.submit(fn, *args, **kwargs) + return self.budget.thread_pool.submit(_carry_ctx(), fn, *args, **kwargs) async def a_in_thread(self, fn: Callable, *args: object, **kwargs: object) -> object: """Await a blocking call on the Budget's thread pool.""" @@ -314,12 +326,13 @@ async def a_in_thread(self, fn: Callable, *args: object, **kwargs: object) -> ob msg = "a_in_thread requires max_parallel > 1" raise RuntimeError(msg) loop = asyncio.get_running_loop() + run = _carry_ctx() if kwargs: return await loop.run_in_executor( self.budget.thread_pool, - lambda: fn(*args, **kwargs), + lambda: run(fn, *args, **kwargs), ) - return await loop.run_in_executor(self.budget.thread_pool, fn, *args) + return await loop.run_in_executor(self.budget.thread_pool, run, fn, *args) # --- sentinel-propagating evaluation ----------------------------------- @@ -417,7 +430,7 @@ async def run_child(n: int) -> None: await q.put(v) else: async with sem: - await loop.run_in_executor(pool, _drain_sync, n, loop) + await loop.run_in_executor(pool, _carry_ctx(), _drain_sync, n, loop) finally: await q.put(_DONE) diff --git a/tests/nu/lang/runtime/test_runtime.py b/tests/nu/lang/runtime/test_runtime.py index 2d5715ec..d768e460 100644 --- a/tests/nu/lang/runtime/test_runtime.py +++ b/tests/nu/lang/runtime/test_runtime.py @@ -512,6 +512,67 @@ def test_ctx_attrs_writes_are_visible_through_runtime() -> None: assert ctx.attrs["y"] == "hello" +# --- context across the thread boundary ----------------------------------- +# +# A pool worker starts with an empty contextvars context, so every hand-off to +# a thread has to carry the caller's copy or `rt.ctx` raises LookupError there. +# One case per dispatch site. + + +def test_ctx_resolves_on_eval_parallel_workers() -> None: + program = _fake_program(thunks=[lambda rt: rt.ctx]) + ctx = Context() + with Budget(max_parallel=2) as budget: + rt = Runtime(program, ctx, budget=budget) + assert rt.eval_parallel([0, 0]) == [ctx, ctx] + + +async def test_ctx_resolves_on_async_placement_workers() -> None: + program = _fake_program(thunks=[lambda rt: rt.ctx], athunks=[None], on_loop=[False]) + ctx = Context() + with Budget(max_parallel=2, async_mode=True) as budget: + rt = Runtime(program, ctx, budget=budget) + assert await rt.aeval_parallel([0, 0]) == [ctx, ctx] + + +def test_ctx_resolves_on_merge_workers() -> None: + def gen(rt: Runtime) -> object: + return iter([rt.ctx]) + + program = _fake_program(thunks=[gen, gen]) + ctx = Context() + with Budget(max_parallel=2) as budget: + rt = Runtime(program, ctx, budget=budget) + assert list(rt.merge([0, 1])) == [ctx, ctx] + + +async def test_ctx_resolves_on_amerge_workers() -> None: + def gen(rt: Runtime) -> object: + return iter([rt.ctx]) + + program = _fake_program(thunks=[gen], athunks=[None], on_loop=[False]) + ctx = Context() + with Budget(max_parallel=2, async_mode=True) as budget: + rt = Runtime(program, ctx, budget=budget) + assert [v async for v in rt.amerge([0])] == [ctx] + + +def test_ctx_resolves_inside_in_thread() -> None: + program = compile(Literal(1)) + ctx = Context() + with Budget(max_parallel=2) as budget: + rt = Runtime(program, ctx, budget=budget) + assert rt.in_thread(lambda: rt.ctx).result() is ctx + + +async def test_ctx_resolves_inside_a_in_thread() -> None: + program = compile(Literal(1)) + ctx = Context() + with Budget(max_parallel=2, async_mode=True) as budget: + rt = Runtime(program, ctx, budget=budget) + assert await rt.a_in_thread(lambda: rt.ctx) is ctx + + # --- budget lifecycle -----------------------------------------------------