Repository navigation
Add trace spans to the Rust io-context crate - #8034
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 generated output of |
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().
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.
IoContext::make_trace_span() and make_user_trace_span() start spans as C++ makeTraceSpan() and makeUserTraceSpan() do, returning a SpanBuilder or TraceContext that holds the C++ object and ends the span when dropped. SpanBuilder::set_tag(), SpanBuilder::new_child() and TraceContext::set_tag() mirror their C++ counterparts. Span names and tag keys are &'static CStr, written as c"..." literals. C++ wraps them without copying, as it does its own string literals. A tag value can be a bool, an integer, a float, a literal, or a borrowed string, which is copied only if the span is observed. A user trace span started before await_io_and_run() is adopted by it for async tracing, as for awaitIo(), and a test checks this.
The memory_cache_fallback span was tagged fallback_success=true before the fallback's result was validated. If validation then threw, as for a NaN expiration or a value that cannot be serialized, the catch handler overwrote the tag with false. The final tags were right, but overwriting a tag is treated as a programming error, logged at ERROR in predictable mode. The tag is now set once the result has been accepted. A test records the spans of two rejected results and checks their tags, failing on the overwrite's log without the fix.
73517e9 to
0a14de1
Compare
guybedford
left a comment
There was a problem hiding this comment.
A few points, mainly on matching the C++ behaviour exactly, since the MemoryCache port depends on that. The MemoryCache fallback_success fix looks right: makeFallback catches its own errors, so nothing after the true tag can send the catch path into overwriting it.
AI generated review
…text-trace-spans # Conflicts: # src/workerd/io/io-context.h # 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
io-context-test and memory-cache-test each had their own SpanObserver that recorded a request's spans and their tags. RecordedSpans in test-fixture now does it for both: it records spans in the order they open, with their parents and sorted tags, and supplies a requestObserverFactory for TestFixture.
C++ TraceContext::setTag() annotates the async trace operation before it checks whether the spans are observed, but the io-context shim returned early for unobserved spans of both kinds. With async tracing on and no span observer, tags set from Rust never reached the operation. The early return now applies only to SpanBuilder. The test that only ran the span code without a tracer now runs it with async tracing and no span observer, and checks the annotations.
It wraps C++ TraceContext::detachAsync(), for a user trace span that outlives the current call without being moved into the next await_io_and_run(), which would otherwise report its wait under the span's operation. A test checks that the wait then gets its own bridge.
Setting a tag twice is treated as a bug in C++, logged as an error in predictable mode, so the set_tag() docs now say to set each tag at most once rather than describing replacement, and the test no longer sets one twice. TagValue also converts from i32 and u32, so an integer literal needs no suffix.
The fallback_error tag of a fallback whose value cannot be serialized is empty, because kj::str() of the DOMException renders as empty. The test checks the span's other tags and that the error tag is present, without making the empty text expected behaviour.
There was a problem hiding this comment.
Superseded by the review of
ccb9247d.
This review, of aded3ac3
💬 Approved with comments — No correctness or safety issues; one INFO note: a doc comment is imprecise about when overwriting a tag is logged as an error.
Findings: 1 info
The decision is advisory: this review is a comment and never approves, blocks or requests changes on the pull request.
📝 Documentation: 💬 1 info
A doc comment is imprecise about when overwriting a tag is logged as an error.
Recommendation: 💬 Optional Nits
- [INFO]
SpanBuilder::set_tagdocs say a tag overwrite is logged as an error "in tests".trace.c++:2055-2057logs it at ERROR only in predictable mode and at WARNING otherwise; say "in predictable mode". (src/workerd/io/io-context/trace.rs:99)
View full review
src/workerd/io/io-context/trace.rs:98-99 says that setting a tag twice is "logged as a warning, and as an error in tests". trace.c++:2055-2057 logs it at ERROR only in predictable mode and at WARNING everywhere else, including tests that do not enable predictable mode. Saying "as an error in predictable mode" would be accurate.
✅ All other checks passed
🔒 Security: ✅ LGTM
No security concerns in the new FFI surface or the test changes.
📖 AGENTS.md: ✅ LGTM
No conflicts with the repository's guidelines.
📋 Code Quality: ✅ LGTM
- MemoryCache fix.
memory-cache.c++now setsfallback_success=trueonly after serialization and the expiration check succeed, so the catch handler no longer overwrites the tag. The new test covers the NaN-expiration and unserializable-value paths in predictable mode. - Literals.
wrapLiteral()inffi.c++requires the trailing NUL and uses a null disposer. Rust only passes&'static CStrbytes throughliteral(), so the lifetime requirement holds. - Tags.
SpanBuilderskips unobserved spans.TraceContextalways forwards its tags, so the async-trace operation is annotated as it is in C++, andio-context-test.c++covers this. - Thread affinity.
SpanBuilderandTraceContextwrap cxx opaque C++ types. The crate declares nounsafe impl Sendorunsafe impl Syncfor them, so neither type isSendorSync, as documented.
⚡ Performance: ✅ LGTM
Spans are created only when a Rust caller asks for one. String tags on unobserved SpanBuilder spans are not copied. The MemoryCache change only moves where one tag is set.
Reviewed commit: aded3ac3
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #8034 +/- ##
==========================================
+ Coverage 68.33% 68.40% +0.06%
==========================================
Files 726 727 +1
Lines 189924 190368 +444
Branches 29738 29791 +53
==========================================
+ Hits 129792 130222 +430
+ Misses 42705 42693 -12
- Partials 17427 17453 +26 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Superseded by the review of
eb56789c.
This review, of ccb9247d
💬 Approved with comments — No correctness, security or performance problems. One INFO note remains: a doc comment says when a tag overwrite is logged as an error, and it has that slightly wrong.
Findings: 1 info (1 is an inline comment on the diff)
The decision is advisory: this review is a comment and never approves, blocks or requests changes on the pull request.
📝 Documentation: 💬 1 info
One open INFO note: the set_tag doc comment says tag overwrites are logged as errors "in tests", but that happens only in predictable mode.
Recommendation: 💬 Optional Nits
- [INFO] The
SpanBuilder::set_tagdoc says a tag overwrite is logged "as an error in tests".trace.c++logs it at ERROR only in predictable mode and at WARNING otherwise, so say "in predictable mode". (src/workerd/io/io-context/trace.rs:99)
View full review
src/workerd/io/io-context/trace.rs:99says that setting a tag a second time is logged "as an error in tests".trace.c++logs it at ERROR only in predictable mode and at WARNING everywhere else, including tests that do not turn predictable mode on (io-context-testis one of them). "as an error in predictable mode" would be accurate.
✅ All other checks passed
🔒 Security: ✅ LGTM
No security concerns in the new FFI surface or the test changes.
📖 AGENTS.md: ✅ LGTM
No conflicts with the repository's guidelines.
📋 Code Quality: ✅ LGTM
memory-cache.c++now setsfallback_success=trueonly after serialization and the expiration check pass, so the catch path can no longer overwrite it. The new predictable-mode test covers both ways a result can be rejected.wrapLiteral()inffi.c++asserts the trailing NUL and uses a null disposer. Rust passes only&'static CStrbytes, so the literal's lifetime is guaranteed.SpanBuilderskips unobserved spans.TraceContextalways forwards its tags, which matches C++, and a test covers this.- In
test-fixture.c++, the recorder's members are initialized in the right order.
⚡ Performance: ✅ LGTM
Spans are created only when a Rust caller asks for one. String tags on unobserved SpanBuilder spans are not copied.
Reviewed commit: ccb9247d
There was a problem hiding this comment.
✅ Approved — The last open note, the wording of the set_tag doc comment, is fixed, and I found no correctness, security or performance problems.
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 concerns in the new bridge between Rust and C++ or in the test changes.
📖 AGENTS.md: ✅ LGTM
Nothing in this change conflicts with the repository's guidelines.
📋 Code Quality: ✅ LGTM
- In
memory-cache.c++,fallback_success=trueis now set only after serialization and the expiration check pass, so the catch path can no longer overwrite it. The new predictable-mode test covers both ways a result can be rejected. KjOwn<T>isSend/Synconly whenTis. The opaqueSpanBuilderandTraceContexttypes are neither, so the claim in the docs that the spans stay on theIoContext's thread holds.wrapLiteral()asserts the trailing NUL. Rust passes only&'static CStrbytes, so a span name or tag key that C++ does not copy lives longer than the span.
⚡ Performance: ✅ LGTM
- A span is created only when Rust code asks for one.
- A string tag on a
SpanBuilderspan is copied only when the span is observed.
📝 Documentation: ✅ LGTM
The SpanBuilder::set_tag doc comment now says a repeated tag is logged "as a warning, and as an error in predictable mode", which matches trace.c++.
Reviewed commit: eb56789c
`internal-build` is a required check, and GitHub counts a skipped required check as passing. Since the workflow started running on `edited` events, editing a PR's title or description created a skipped `internal-build` check run on the same commit as a build still in progress. On cloudflare#8034 such a run appeared while the build was running, the PR was merged, and the build then failed. Name the job `internal-build` only for runs that check the PR, and give it a different id so a skipped run can't fall back to that name. Filter stacked PRs with `branches: [main]` rather than a job condition, so they get no check instead of a skipped one. Fold the paths filter into the build job, so docs-only PRs report success and a failing filter fails the check rather than skipping it.
Gives Rust APIs the trace spans that porting
MemoryCache.read()toRust needs, so the Rust port records the same spans and tags as the C++
implementation:
IoContext::make_trace_span()andmake_user_trace_span()call C++makeTraceSpan()andmakeUserTraceSpan(), so spans are children ofthe span current in the async context frame.
SpanBuilderandTraceContexthold the C++ objects and end thespans when dropped. Like the C++ types, they stay on the
IoContext's thread: neither isSendorSync.SpanBuilder::set_tag(),SpanBuilder::new_child(),TraceContext::set_tag()andTraceContext::detach_async()mirrortheir 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::ConstStringwith a null disposer, as it does_kjcliterals. Atag 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()sharesawaitIo()'simplementation (#8030).
TraceContext::set_tag()annotates thatoperation 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(), whichthe Rust port has to reproduce tag for tag. The
memory_cache_fallbackspan was tagged
fallback_success=truebefore the fallback's resultwas 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, buttrace.c++treats overwritinga tag as a programming error and logs it at ERROR in predictable mode.
The tag is now set once the result is accepted.
RecordedSpansin test-fixture records the spans of the requests aTestFixtureruns; io-context-test and memory-cache-test share it. Anew test in
memory-cache-test.c++records the spans of two rejectedresults and checks their tags; without the fix it fails on the
overwrite's log.