Repository navigation
Add await_io_and_run() and add_object() to the Rust io-context crate - #8030
Conversation
Add the io-context crate, beside //src/workerd/io in a package of its own because //src/workerd/io depends on jsg. It has IoContext::current() and try_current(), now(), the request's clock, and write_logfwdr(). A handle holds the IoContext's id (IoContext::getId()), not a pointer. Each call looks up IoContext::tryCurrent() and fails unless the id matches, so a handle can never reach another request's IoContext or a destroyed one, and lib.rs needs no unsafe. Outside a request, current() fails with the error C++ IoContext::current() throws; the Rust copy of its message and kAsyncIoErrorMessage name each other. In a Durable Object, write_logfwdr() waits for the output gate, as AnalyticsEngine::writeDataPoint() does, and it rejects anything but exactly one serialized message. The shims use no C++ JSG and are checked by cpp-jsg-allowlist-test. io-context-test runs the crate in a real worker through TestFixture, calling Rust through a test-only bridge: inside and outside a request, logfwdr delivery and validation, the output gate, and stale ids.
Take the jsg::Error by reference in to_io(), and return errors rather than unwrapping when converting the time to nanoseconds.
Each copy now names the other, and the test that checks they match.
Either module registry can run a Rust builtin's add_builtin closure during a request: the legacy registry the first time a request resolves the module, for example by reading the process global, and the new registry when a request resolves the synthetic module directly. The closure then saw that request's IoContext. The adapters in rust/jsg/jsg.h now construct a CallScope around each call of a Rust module callback, and the Node.js registration passes SuppressIoContextScope. jsg cannot depend on the IoContext, so the scope is a template parameter with no default, which makes every registration choose it. io-context-test resolves a probe module during a request through each registry and checks that its closure saw no IoContext.
…-context # Conflicts: # src/rust/jsg/modules.rs
lib-test.rs was listed in the crate's sources, but lib.rs never declared it as a module, so it was never compiled and io-context_test passed without running any tests.
The shim now returns the current IoContext's id, from IoContext::tryGetCurrentId(), instead of a reference that Rust read the id through. That drops the opaque IoContext type and all hand-written unsafe code from the crate. The message of the error IoContext::current() throws outside a request becomes IoContext::NO_REQUEST_MESSAGE, which Rust reads through the shim instead of keeping a copy, so the test that the copies match goes away. Also builds the test bridge against main, which removed jsg::Number.
…nto biwanczuk/rust-jsg-io-context
Rust APIs need it for failures that only a bug causes, such as an IoContext that is no longer current. Throwing the error gives JavaScript a generic "internal error; reference = <id>" Error and logs the message with that ID. A new test checks the redaction.
IoContext::await_io() waits for a Rust future as part of the request, then runs a callback with its output in the IoContext. It goes through IoContext::awaitIoAndRun(), a new public variant of awaitIo() for callers that settle no jsg::Promise. Both now share the I/O half of awaitIoImpl(), split out as resumeAfterIo(), so the callback runs like awaitIo()'s does: holding a pending event, in the critical section and async context frame current when the wait began, and recorded for async tracing. If the request ends first, the future and callback are dropped without the callback running. IoContext::add_object() gives the IoContext ownership of a value, as C++ addObject() does. It returns IoOwned<T>, from which the value can only be taken out, rather than used in place as IoOwn<T> is.
05d31f8 to
42efd1c
Compare
…text-await-io # Conflicts: # src/rust/jsg/BUILD.bazel # src/workerd/io/io-context/BUILD.bazel # src/workerd/io/io-context/ffi.c++ # src/workerd/io/io-context/ffi.h # src/workerd/io/io-context/ffi.rs # src/workerd/io/io-context/io-context-test.c++ # src/workerd/io/io-context/lib.rs # src/workerd/io/io-context/test_bridge.rs
An await_io() callback now returns jsg::Result<()>. If it fails, IoContext::run() logs its exception and fails the task, as for an addTask() task. A panic in the callback, which terminates execution, fails the task with a KJ exception, and a panic in the future fails the wait instead of aborting the process. IoOwned::take() takes a lock and fails, as dereferencing an IoOwn does, unless the IoContext that owns the value is current. Remove jsg::Lock::await_io(), an unimplemented stub, in favour of IoContext::await_io(), whose compile tests now pin that the future is Send and 'static, that the callback is 'static, and that the callback's handles cannot escape. The future stays Send: it is polled without the isolate lock. new_internal_error() names its errors Error, the type JavaScript gets. Also test await_io() in a Durable Object behind the input gate, against hang detection, and with failing or panicking callbacks and futures; make the cross-thread test wait on the event loop instead of polling; and correct resumeAfterIo()'s comments.
Throw JsExceptionThrown directly in io_context_await_io(), as a reviewed exception counted in the C++ JSG allowlist, rather than through a helper in src/rust/jsg/jsg.h that kept the use out of the count.
It wraps C++ IoContext::awaitIoAndRun(), not awaitIo(): it returns no promise, and its callback settles any promise itself, so a failure rejects nothing. Name it after what it wraps, and leave await_io free for a promise-returning version that would match awaitIo(). The FFI function and the compile tests are renamed to match, and the doc comment now says how it differs from awaitIo().
bartlomieju
changed the title
Rust code that resumes JavaScript after a KJ event must run in the async context frame that started the work, as IoContext::awaitIo does in C++. Add jsg::AsyncContextFrame, which captures the current frame and later runs a closure in it with enter(), restoring the previous frame however the closure ends. TracedAsyncContextFrame is its GC-traced counterpart for resource fields and Function::new data; to_async_context_frame() converts back. Both read and write V8's continuation-preserved embedder data directly. C++ AsyncContextFrame::current() treats embedder data that is not a frame wrapper as the root frame. Rust cannot recognize the wrapper without C++ JSG, so it treats anything but undefined as a frame. The two agree because nothing else is ever set: C++ sets only a frame's wrapper or undefined, and Rust can set only undefined or a value it read from the embedder data, which a private handle type enforces. The C++ setter documents the invariant, and a test checks that Rust's is_root() agrees with C++ in every state Rust can leave the isolate in. The jsg-test shims create real C++ jsg::AsyncContextFrames the way AsyncLocalStorage.run() does. Compile tests check that enter() holds the lock and does not let it escape, that AsyncContextFrame cannot be traced, and that TracedAsyncContextFrame is not Sync. No JS-visible change. Nothing outside tests uses it yet: IoContext::await_io (#8030) enters the frame on the C++ side. Reading and writing AsyncLocalStorage values and AsyncContextFrame::wrap are left for later.
![clanky-the-clanker[bot]](https://avatars.githubusercontent.com/in/5306670?s=60&v=4){kind=small}
There was a problem hiding this comment.
Superseded by the review of
744da07a.
This review, of 5275f1f0
💬 Approved with comments — Nothing blocks this pull request and no WARNING remains. Two small documentation follow-ups: one INFO about a stale compile-test pointer in the jsg README and one SUGGESTION about the Rust AGENTS.md guidance.
Findings: 1 info · 1 suggestion
Not opened by the correctness review: 3 changed files it was meant to read
src/rust/jsg/cpp-jsg-allowlist.txtsrc/workerd/io/io-context/BUILD.bazelsrc/workerd/io/io-context/test-bridge.h
The decision is advisory: this review is a comment and never approves, blocks or requests changes on the pull request.
📖 AGENTS.md: 💬 1 suggestion
The Rust agent guidance doesn't mention the new allowlisted C++ JSG use
Recommendation: 💬 Optional Nits
- [SUGGESTION] Mention the allowlisted
jsg::JsExceptionThrownuse inio_context_await_io_and_run()insrc/rust/AGENTS.mdrule 3, matching the README. (src/workerd/io/io-context/ffi.c++:47)
View full review
src/rust/AGENTS.md:50(safety rule 3) says the C++ shims call V8 directly, not C++ JSG.src/workerd/io/io-context/ffi.c++now throwsjsg::JsExceptionThrown, a reviewed and allowlisted exception thatsrc/rust/jsg/README.mdalready explains. Add a short note to the AGENTS.md rule, or a pointer to the README, so the guidance doesn't look absolute.
📝 Documentation: 💬 1 info
The jsg README still lists the await-io-* compile tests under src/rust/jsg-compile-tests
Recommendation: 💬 Optional Nits
- [INFO]
src/rust/jsg/README.md:169listsawait-io-*compile tests insrc/rust/jsg-compile-tests. They are removed there and now live insrc/workerd/io/io-context/compile-testsasawait-io-and-run-*. (src/rust/jsg/README.md:169)
View full review
src/rust/jsg/README.md:168-170: the "Tested by" line under §4 still listsawait-io-*among the compile tests insrc/rust/jsg-compile-tests. This pull request deletes those tests and addsawait-io-and-run-*tests undersrc/workerd/io/io-context/compile-tests. Point readers to the new location.
✅ All other checks passed
🔒 Security: ✅ LGTM
No security issues found. IoOwned::take() checks that the owning IoContext is current before handing out the value, using the same error as an I/O object used by another request.
📋 Code Quality: ✅ LGTM
No issues. resumeAfterIo() is shared by awaitIoImpl() and awaitIoAndRun(), which keeps awaitIo()'s behaviour the same. The FFI shim handles all three callback outcomes (succeeded, failed, panicked), and a panicking future fails the task. IoOwned's slot is kept behind an Arc<Mutex<..>>, which fits the C++ IoOwn deferred-drop rules.
⚡ Performance: ✅ LGTM
No measurable performance impact. Each call does a few small allocations and creates one mutex.
Reviewed commit: 5275f1f0
There was a problem hiding this comment.
Superseded by the review of
8bc245e0.
This review, of 744da07a
✅ Approved — Both documentation notes from the last review are fixed, and nothing changed since then raises a new issue.
Not opened by the correctness review: 3 changed files it was meant to read
src/rust/jsg/cpp-jsg-allowlist.txtsrc/workerd/io/io-context/BUILD.bazelsrc/workerd/io/io-context/test-bridge.h
The decision is advisory: this review is a comment and never approves, blocks or requests changes on the pull request.
✅ All checks passed
🔒 Security: ✅ LGTM
No security issues found. IoOwned::take() checks that the owning IoContext is current before it hands out the value.
📖 AGENTS.md: ✅ LGTM
src/rust/AGENTS.md safety rule 3 now names io_context_await_io_and_run() as the one reviewed exception, the shim that throws jsg::JsExceptionThrown. It matches the README and the allowlist.
📋 Code Quality: ✅ LGTM
No issues. resumeAfterIo() is shared by awaitIoImpl() and awaitIoAndRun(), so awaitIo() behaves the same as before. The BUILD file adds the compile-test target and the cxx header dependency these changes need.
⚡ Performance: ✅ LGTM
No measurable performance impact. Each call makes a few small allocations, and none of this is on a hot path.
📝 Documentation: ✅ LGTM
The README's "Tested by" line now points to src/workerd/io/io-context/compile-tests (await-io-and-run-*).
Reviewed commit: 744da07a
test_bridge.rs calls record_async_context() and record_await_context(), which were defined in io-context-test.c++. The rust_test that wd_rust_crate generates for io-context-test-bridge links the bridge without that file. ELF and Mach-O linkers drop the unreachable shims before reporting undefined symbols, but lld-link keeps them alive through the crate's extern "Rust" exports, so the Windows build failed to link. Move the callbacks and the state they record into a testonly test-bridge library in the crate's link_deps, so every binary that links the crate resolves them.
There was a problem hiding this comment.
Superseded by the review of
49d533e3.
This review, of 8bc245e0
✅ Approved — No issues found. Since the last review, the only change moves the test bridge's C++ callbacks into a library that every binary linking the bridge crate links too, and that change is correct.
Not opened by the correctness review: 1 changed file it was meant to read
src/rust/jsg/cpp-jsg-allowlist.txt
The decision is advisory: this review is a comment and never approves, blocks or requests changes on the pull request.
✅ All checks passed
🔒 Security: ✅ LGTM
No security issues found. IoOwned::take() still checks that the owning IoContext is current before it hands out the value.
📖 AGENTS.md: ✅ LGTM
Safety rule 3 in src/rust/AGENTS.md names io_context_await_io_and_run() as the one reviewed exception, which matches the README and the allowlist.
📋 Code Quality: ✅ LGTM
The new :test-bridge library defines record_await_context() and record_async_context(). The bridge crate links it through link_deps, so the generated io-context-test-bridge_test can resolve both symbols. The shared test state is declared extern in test-bridge-state.h and defined once in test-bridge.c++, so io-context-test.c++ and the bridge use the same state. resumeAfterIo() is still shared by awaitIoImpl() and awaitIoAndRun(), so awaitIo() behaves as before. The allowlist entry for src/workerd/io/io-context/ffi.c++ JsExceptionThrown 1 matches the single reviewed exception in that shim.
⚡ Performance: ✅ LGTM
No measurable performance impact. Each call makes only a few small allocations, none on a hot path, and the latest change touches only test code.
📝 Documentation: ✅ LGTM
The README, the AGENTS.md safety rule and the allowlist all describe the same exception, io_context_await_io_and_run(), and they agree with each other.
Reviewed commit: 8bc245e0
There was a problem hiding this comment.
Superseded by the review of
3ec326e1.
This review, of 49d533e3
✅ Approved — No issues found. Since the last review, the only change moves the test bridge's C++ callbacks into a library that every binary linking the bridge crate links too, and that change is correct.
Not opened by the correctness review: 1 changed file it was meant to read
src/rust/jsg/cpp-jsg-allowlist.txt
The decision is advisory: this review is a comment and never approves, blocks or requests changes on the pull request.
✅ All checks passed
🔒 Security: ✅ LGTM
No security issues found. IoOwned::take() still checks that the owning IoContext is current before it hands out the value.
📖 AGENTS.md: ✅ LGTM
Safety rule 3 in src/rust/AGENTS.md names io_context_await_io_and_run() as the one reviewed exception, which matches the README and the allowlist.
📋 Code Quality: ✅ LGTM
The new :test-bridge library defines record_await_context() and record_async_context(). The bridge crate links it through link_deps, so the generated io-context-test-bridge_test can resolve both symbols. The shared test state is declared extern in test-bridge-state.h and defined once in test-bridge.c++, so io-context-test.c++ and the bridge use the same state. resumeAfterIo() is still shared by awaitIoImpl() and awaitIoAndRun(), so awaitIo() behaves as before. The allowlist entry for src/workerd/io/io-context/ffi.c++ JsExceptionThrown 1 matches the single reviewed exception in that shim.
⚡ Performance: ✅ LGTM
No measurable performance impact. Each call makes only a few small allocations, none on a hot path, and the latest change touches only test code.
📝 Documentation: ✅ LGTM
The README, the AGENTS.md safety rule and the allowlist all describe the same exception, io_context_await_io_and_run(), and they agree with each other.
Reviewed commit: 49d533e3 — same changes as 8bc245e0, so this is that commit's review (no new review was run; line numbers are as of 8bc245e0)
There was a problem hiding this comment.
Superseded by the review of
fc74e465.
This review, of 3ec326e1
✅ Approved — No issues found. Since the last review, the only change moves the test bridge's C++ callbacks into a library that every binary linking the bridge crate links too, and that change is correct.
Not opened by the correctness review: 1 changed file it was meant to read
src/rust/jsg/cpp-jsg-allowlist.txt
The decision is advisory: this review is a comment and never approves, blocks or requests changes on the pull request.
✅ All checks passed
🔒 Security: ✅ LGTM
No security issues found. IoOwned::take() still checks that the owning IoContext is current before it hands out the value.
📖 AGENTS.md: ✅ LGTM
Safety rule 3 in src/rust/AGENTS.md names io_context_await_io_and_run() as the one reviewed exception, which matches the README and the allowlist.
📋 Code Quality: ✅ LGTM
The new :test-bridge library defines record_await_context() and record_async_context(). The bridge crate links it through link_deps, so the generated io-context-test-bridge_test can resolve both symbols. The shared test state is declared extern in test-bridge-state.h and defined once in test-bridge.c++, so io-context-test.c++ and the bridge use the same state. resumeAfterIo() is still shared by awaitIoImpl() and awaitIoAndRun(), so awaitIo() behaves as before. The allowlist entry for src/workerd/io/io-context/ffi.c++ JsExceptionThrown 1 matches the single reviewed exception in that shim.
⚡ Performance: ✅ LGTM
No measurable performance impact. Each call makes only a few small allocations, none on a hot path, and the latest change touches only test code.
📝 Documentation: ✅ LGTM
The README, the AGENTS.md safety rule and the allowlist all describe the same exception, io_context_await_io_and_run(), and they agree with each other.
Reviewed commit: 3ec326e1 — same changes as 8bc245e0, so this is that commit's review (no new review was run; line numbers are as of 8bc245e0)
There was a problem hiding this comment.
Superseded by the review of
f042712c.
This review, of fc74e465
✅ Approved — No issues found. Since the last review, the only change moves the test bridge's C++ callbacks into a library that every binary linking the bridge crate links too, and that change is correct.
Not opened by the correctness review: 1 changed file it was meant to read
src/rust/jsg/cpp-jsg-allowlist.txt
The decision is advisory: this review is a comment and never approves, blocks or requests changes on the pull request.
✅ All checks passed
🔒 Security: ✅ LGTM
No security issues found. IoOwned::take() still checks that the owning IoContext is current before it hands out the value.
📖 AGENTS.md: ✅ LGTM
Safety rule 3 in src/rust/AGENTS.md names io_context_await_io_and_run() as the one reviewed exception, which matches the README and the allowlist.
📋 Code Quality: ✅ LGTM
The new :test-bridge library defines record_await_context() and record_async_context(). The bridge crate links it through link_deps, so the generated io-context-test-bridge_test can resolve both symbols. The shared test state is declared extern in test-bridge-state.h and defined once in test-bridge.c++, so io-context-test.c++ and the bridge use the same state. resumeAfterIo() is still shared by awaitIoImpl() and awaitIoAndRun(), so awaitIo() behaves as before. The allowlist entry for src/workerd/io/io-context/ffi.c++ JsExceptionThrown 1 matches the single reviewed exception in that shim.
⚡ Performance: ✅ LGTM
No measurable performance impact. Each call makes only a few small allocations, none on a hot path, and the latest change touches only test code.
📝 Documentation: ✅ LGTM
The README, the AGENTS.md safety rule and the allowlist all describe the same exception, io_context_await_io_and_run(), and they agree with each other.
Reviewed commit: fc74e465 — same changes as 8bc245e0, so this is that commit's review (no new review was run; line numbers are as of 8bc245e0)
guybedford
left a comment
guybedford
left a comment
There was a problem hiding this comment.
LGTM, thanks. A couple of small nits inline.
Its only caller, the Rust io-context shim, rethrew the exception itself: a Rust future fails only by panicking, and there is nothing for the callback to do then. awaitIoAndRun() now rethrows it inside run() and passes `func` just the value, so ExceptionOr stays private and the shim no longer switches on it. The task still fails the same way, and async tracing still records the wait as failed. Also move the comment about rejecting with the caught JS exception in awaitIoImpl() onto the tryCatch() handler it describes.
There was a problem hiding this comment.
💬 Approved with comments — No correctness, security or performance issues. One INFO: the description still says ExceptionOr<T> becomes public, which the latest commit reverted.
Findings: 1 info
Not opened by the correctness review: 1 changed file it was meant to read
src/rust/jsg/cpp-jsg-allowlist.txt
The decision is advisory: this review is a comment and never approves, blocks or requests changes on the pull request.
📝 Documentation: 💬 1 info
Pull request description is stale on ExceptionOr visibility
Recommendation: 💬 Optional Nits
- [INFO] The pull request description says "ExceptionOr<T> becomes public, because awaitIoAndRun()'s callback receives one". Since f042712 it stays private and the callback receives only the value. Consider updating the description so the merged commit message is accurate.
View full review
The README, src/rust/AGENTS.md and cpp-jsg-allowlist.txt all describe the same single reviewed exception, io_context_await_io_and_run(), and they agree with each other. The pull request description still says that ExceptionOr<T> becomes public and that the callback receives one. Since f042712, ExceptionOr is private (io-context.h:1496) and func receives only the value.
✅ All other checks passed
🔒 Security: ✅ LGTM
No security issues found. IoOwned::take() checks that the owning IoContext is current before it hands out the value. Because the future and its output must be Send, isolate-bound values cannot get into a future that is polled without the isolate lock.
📖 AGENTS.md: ✅ LGTM
Safety rule 3 in src/rust/AGENTS.md names io_context_await_io_and_run() as the one reviewed exception, which matches the README and the allowlist.
📋 Code Quality: ✅ LGTM
- Failure path (
io-context.h:1921-1950): when the promise fails,awaitIoAndRun()callskj::throwFatalExceptioninsiderun(). The task fails and is logged as anaddTask()task would be, andfuncdoes not run. - Same continuation as
awaitIo(): the callback holds a pending event (promise.attach(registerPendingEvent())) and runs undergetCriticalSection().requireCurrent()is called beforegetCurrentLock(). awaitIoImpl(): after the split, the promise is still destroyed beforefuncifrequireCurrent()ornewPromiseAndResolverthrows.- Callback outcomes in
ffi.c++: all three cases are handled.FailedthrowsJsExceptionThrownforrun()to log.Panickedfails with a KJ assertion, so no JavaScript runs while termination is pending.
⚡ Performance: ✅ LGTM
No performance issues. Each awaitIoAndRun() call costs about what an awaitIo() call costs, and the shared resumeAfterIo() adds no work to the awaitIo() path.
Reviewed commit: f042712c
Gives Rust APIs the trace spans that porting `MemoryCache.read()` to Rust needs, so the Rust port records the same spans and tags as the C++ implementation: ```rust let mut span = ctx.make_trace_span(c"memory_cache_read")?; span.set_tag(c"cache_hit", false); let mut child = span.new_child(c"memory_cache_fallback"); let user_span = ctx.make_user_trace_span(c"memory_cache_read")?; ``` - `IoContext::make_trace_span()` and `make_user_trace_span()` call C++ `makeTraceSpan()` and `makeUserTraceSpan()`, so spans are children of the span current in the async context frame. - `SpanBuilder` and `TraceContext` hold the C++ objects and end the spans when dropped. Like the C++ types, they stay on the `IoContext`'s thread: neither is `Send` or `Sync`. - `SpanBuilder::set_tag()`, `SpanBuilder::new_child()`, `TraceContext::set_tag()` and `TraceContext::detach_async()` mirror their C++ counterparts. As in C++, setting a tag twice replaces its value but is treated as a bug: it is logged as a warning, and as an error in predictable mode. C++ expects span names and tag keys to be string literals, which it does not copy. The Rust API takes them as `&'static CStr` (`c"..."` literals), which cross the bridge with their NUL, so C++ wraps them in a `kj::ConstString` with a null disposer, as it does `_kjc` literals. A tag value (`TagValue`) can be a bool, an integer, a float, a literal (not copied either), or a borrowed string, which is copied only if the span is recorded. As in C++, a user trace span started just before `await_io_and_run()` is an operation that the wait is reported under, instead of a bridge of its own. This works because `await_io_and_run()` shares `awaitIo()`'s implementation (#8030). `TraceContext::set_tag()` annotates that operation even when the spans are not recorded, and `detach_async()` stops a later, unrelated wait in the same turn from taking it over. This also fixes a span tag bug in the C++ `MemoryCache.read()`, which the Rust port has to reproduce tag for tag. The `memory_cache_fallback` span was tagged `fallback_success=true` before the fallback's result was validated; when validation then threw (a NaN expiration, or a value that cannot be serialized), the catch handler overwrote the tag with `false`. The final tags were right, but `trace.c++` treats overwriting a tag as a programming error and logs it at ERROR in predictable mode. The tag is now set once the result is accepted. `RecordedSpans` in test-fixture records the spans of the requests a `TestFixture` runs; io-context-test and memory-cache-test share it. A new test in `memory-cache-test.c++` records the spans of two rejected results and checks their tags; without the fix it fails on the overwrite's log.
Give Rust APIs the two IoContext operations that porting
MemoryCache.read() to Rust (#8066) needs: await_io_and_run(), to wait
for I/O and then resume the request, and add_object(), to tie a value's
lifetime to the request.
await_io_and_run() wraps a new C++ IoContext::awaitIoAndRun(), which
shares its implementation with awaitIo() but returns no promise: the
callback settles whatever it needs to itself, and failures fail the
task as for addTask(). add_object() wraps addObject() and returns an
IoOwned with IoOwn's lifetime rules. These replace the
unimplemented jsg::Lock::await_io() stub, and jsg::Error's
new_internal_error() becomes public for the port to use.
Tests cover both operations in requests to a real worker, including
Durable Objects, failures and panics. No JS-visible change; nothing
outside tests uses these yet.