Give Rust APIs the IoContext of the current request

src/rust/jsg/BUILD.bazel

@@ -7,7 +7,10 @@ load("//:build/wd_rust_crate.bzl", "wd_rust_crate")
 CPP_SHIMS = glob([
     "*.c++",
     "*.h",
-])
+]) + [
+    "//src/workerd/io/io-context:ffi.c++",
+    "//src/workerd/io/io-context:ffi.h",
+]
 
 # Fails when a shim adds a use of C++ JSG; see cpp-jsg-allowlist-test.py.
 py_test(

src/rust/jsg/jsg.h

@@ -18,8 +18,19 @@
 
 namespace workerd::rust::jsg {
 
-// Calls a Rust module callback. An empty result means it panicked and threw.
-inline v8::Local<v8::Value> callModuleCallback(ModuleCallback callback, ::workerd::jsg::Lock& js) {
+// The adapters below construct a `CallScope` around each call of a Rust module callback. A
+// module's exports must not depend on the request that happens to import it first, but either
+// registry can create them during a request: the legacy registry when a request first resolves
+// the module (for example, by reading the `process` global), and the new one when a request
+// resolves it directly. Callers pass `SuppressIoContextScope`, so that the callback sees no
+// current IoContext, as ES module evaluation in the new registry already does. jsg cannot depend
+// on the IoContext, so the scope is a template parameter, with no default so that every caller
+// has to choose it.
+
+// Calls a Rust module callback in a `CallScope`. An empty result means it panicked and threw.
+template <typename CallScope>
+v8::Local<v8::Value> callModuleCallback(ModuleCallback callback, ::workerd::jsg::Lock& js) {
+  CallScope scope;
   v8::Local<v8::Value> value;
   if (!maybe_local_from_ffi<v8::Value>(callback(js.v8Isolate)).ToLocal(&value)) {
     throw ::workerd::jsg::JsExceptionThrown();
@@ -29,7 +40,7 @@ inline v8::Local<v8::Value> callModuleCallback(ModuleCallback callback, ::worker
 
 // Adapter for registering Rust-implemented modules into the OLD module registry
 // (jsg::ModuleRegistryImpl). Translates Rust ModuleCallback into ModuleInfo.
-template <typename Registry>
+template <typename Registry, typename CallScope>
 struct RustModuleRegistry: public ::workerd::rust::jsg::ModuleRegistry {
   virtual ~RustModuleRegistry() = default;
   RustModuleRegistry(Registry& registry): registry(registry) {}
@@ -56,7 +67,7 @@ struct RustModuleRegistry: public ::workerd::rust::jsg::ModuleRegistry {
             ::workerd::jsg::Lock& js, ::workerd::jsg::ModuleRegistry::ResolveMethod,
             kj::Maybe<const kj::Path&>&) mutable
         -> kj::Maybe<::workerd::jsg::ModuleRegistry::ModuleInfo> {
-      auto value = callModuleCallback(callback, js);
+      v8::Local<v8::Value> value = callModuleCallback<CallScope>(callback, js);
       KJ_DASSERT(value->IsObject());
 
       using ModuleInfo = ::workerd::jsg::ModuleRegistry::ModuleInfo;
@@ -77,8 +88,9 @@ struct RustModuleRegistry: public ::workerd::rust::jsg::ModuleRegistry {
 // callModuleCallback() unwraps before the synthetic module is registered.
 //
 // Usage:
-//   RustBuiltinModuleAdapter adapter(builder);
+//   RustBuiltinModuleAdapter<SuppressIoContextScope> adapter(builder);
 //   ::workerd::rust::api::register_nodejs_modules(adapter);
+template <typename CallScope>
 struct RustBuiltinModuleAdapter: public ::workerd::rust::jsg::ModuleRegistry {
   ::workerd::jsg::modules::ModuleBundle::BuiltinBuilder& builder;
 
@@ -112,7 +124,7 @@ struct RustBuiltinModuleAdapter: public ::workerd::rust::jsg::ModuleRegistry {
           [callback = kj::mv(moduleCallback)](::workerd::jsg::Lock& js, const ::workerd::jsg::Url&,
               const ::workerd::jsg::modules::Module::ModuleNamespace& ns,
               const ::workerd::jsg::CompilationObserver&) mutable {
-        auto value = callModuleCallback(callback, js);
+        v8::Local<v8::Value> value = callModuleCallback<CallScope>(callback, js);
         KJ_DASSERT(value->IsObject());
         ns.setDefault(js, ::workerd::jsg::JsValue(value));
         return true;

src/rust/jsg/modules.rs

@@ -20,9 +20,10 @@ pub use crate::v8::ffi::ModuleType;
 ///
 /// `callback` produces the module's exports when JavaScript imports it. It runs in the module
 /// registry's current scope, so it receives a [`Lock`] for that scope and returns a handle
-/// branded with the same scope; an error it returns is thrown to the importer. It must capture
-/// nothing: C++ calls it through a plain function pointer, so there is nowhere to keep captured
-/// state (see [`UnitType`]).
+/// branded with the same scope; an error it returns is thrown to the importer. It sees no
+/// request, even when a request imports the module: the C++ registration runs it with no current
+/// `IoContext`. It must capture nothing: C++ calls it through a plain function pointer, so there
+/// is nowhere to keep captured state (see [`UnitType`]).
 pub fn add_builtin<F>(
     registry: Pin<&mut ffi::ModuleRegistry>,
     specifier: &str,

src/workerd/api/node/node.h

@@ -13,6 +13,7 @@
 #include <workerd/api/node/timers.h>
 #include <workerd/api/node/util.h>
 #include <workerd/io/compatibility-date.capnp.h>
+#include <workerd/io/io-context.h>
 #include <workerd/jsg/jsg.h>
 #include <workerd/jsg/modules-new.h>
 #include <workerd/jsg/url.h>
@@ -234,7 +235,7 @@ void registerNodeJsCompatModules(
     }
   }
 
-  ::workerd::rust::jsg::RustModuleRegistry r(registry);
+  ::workerd::rust::jsg::RustModuleRegistry<Registry, SuppressIoContextScope> r(registry);
   ::workerd::rust::api::register_nodejs_modules(r);
 }
 
@@ -267,7 +268,7 @@ kj::Own<jsg::modules::ModuleBundle> getInternalNodeJsCompatModuleBundle(
   // by the same registration functions run in
   // getExternalNodeJsCompatModuleBundle().
   {
-    ::workerd::rust::jsg::RustBuiltinModuleAdapter adapter(builder);
+    ::workerd::rust::jsg::RustBuiltinModuleAdapter<SuppressIoContextScope> adapter(builder);
     ::workerd::rust::api::register_nodejs_modules(adapter);
   }
 
@@ -386,7 +387,7 @@ kj::Own<jsg::modules::ModuleBundle> getExternalNodeJsCompatModuleBundle(
   // legacy path registers Rust modules unconditionally with their declared
   // type).
   {
-    ::workerd::rust::jsg::RustBuiltinModuleAdapter adapter(builder);
+    ::workerd::rust::jsg::RustBuiltinModuleAdapter<SuppressIoContextScope> adapter(builder);
     ::workerd::rust::api::register_nodejs_modules(adapter);
   }
   return builder.finish();

src/workerd/io/AGENTS.md

@@ -27,6 +27,7 @@ I/O lifecycle, per-request context, worker/isolate management, actor storage, co
 | ------------------------------- | --------------------------------------------------------------------------------- |
 | Promise bridging KJ↔JS         | `io-context.h` — `awaitIo()`, `awaitJs()`                                         |
 | Request lifecycle / subrequests | `io-context.{h,c++}`, `worker-entrypoint.{h,c++}`                                 |
+| `IoContext` from Rust           | `io-context/` — the `io-context` crate; identifies an `IoContext` by its `Id`     |
 | Actor storage ops               | `actor-cache.h` (`ActorCacheOps`), `actor-sqlite.h`, `actor-storage.capnp`        |
 | DO gate semantics               | `io-gate.{h,c++}` — `InputGate::CriticalSection`, `OutputGate::lockWhile()`       |
 | Worker/isolate creation         | `worker.{h,c++}`, `worker-modules.{h,c++}`                                        |

src/workerd/io/io-context.c++

@@ -1840,18 +1840,11 @@ void IoContext::runImpl(Runnable& runnable,
   });
 }
 
-static constexpr auto kAsyncIoErrorMessage =
-    "Disallowed operation called within global scope. Asynchronous I/O "
-    "(ex: fetch() or connect()), setting a timeout, and generating random "
-    "values are not allowed within global scope. To fix this error, perform this "
-    "operation within a handler. "
-    "https://developers.cloudflare.com/workers/runtime-apis/handlers/";
-
 IoContext& IoContext::current() {
   if (!hasCurrent()) {
     v8::Isolate* isolate = v8::Isolate::TryGetCurrent();
     KJ_REQUIRE(isolate != nullptr, "there is no current request on this thread");
-    isolate->ThrowError(jsg::v8StrIntern(isolate, kAsyncIoErrorMessage));
+    isolate->ThrowError(jsg::v8StrIntern(isolate, NO_REQUEST_MESSAGE));
     throw jsg::JsExceptionThrown();
   } else {
     return *threadLocalRequest;

src/workerd/io/io-context.h

@@ -554,6 +554,14 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler
   // Throws an exception if there is no current context (see hasCurrent() below).
   static IoContext& current();
 
+  // The message of the error current() throws when there is no current context.
+  static constexpr kj::StringPtr NO_REQUEST_MESSAGE =
+      "Disallowed operation called within global scope. Asynchronous I/O "
+      "(ex: fetch() or connect()), setting a timeout, and generating random "
+      "values are not allowed within global scope. To fix this error, perform this "
+      "operation within a handler. "
+      "https://developers.cloudflare.com/workers/runtime-apis/handlers/"_kj;
+
   // Like current(), but returns kj::none if there is no current context.
   static kj::Maybe<IoContext&> tryCurrent();
 

src/workerd/io/io-context/BUILD.bazel

@@ -0,0 +1,61 @@
+load("//:build/kj_test.bzl", "kj_test")
+load("//:build/wd_cc_library.bzl", "wd_cc_library")
+load("//:build/wd_rust_crate.bzl", "wd_rust_crate")
+
+# The IoContext of the current request, for APIs implemented in Rust. A package of its own
+# because //src/workerd/io depends on jsg.
+wd_rust_crate(
+    name = "io-context",
+    srcs = [
+        "ffi.rs",
+        "lib.rs",
+        "lib-test.rs",
+    ],
+    cxx_bridge_deps = ["//src/workerd/io"],
+    cxx_bridge_hdrs = ["ffi.h"],
+    cxx_bridge_src = "ffi.rs",
+    link_deps = [":ffi"],
+    test_deps = ["//src/rust/jsg-test"],
+    visibility = ["//visibility:public"],
+    deps = ["//src/rust/jsg"],
+)
+
+wd_cc_library(
+    name = "ffi",
+    srcs = ["ffi.c++"],
+    hdrs = ["ffi.h"],
+    deps = ["//src/workerd/io"],
+)
+
+# Checked by //src/rust/jsg:cpp-jsg-allowlist-test.
+exports_files([
+    "ffi.c++",
+    "ffi.h",
+])
+
+# The Rust half of io-context-test.
+wd_rust_crate(
+    name = "io-context-test-bridge",
+    testonly = True,
+    srcs = ["test_bridge.rs"],
+    crate_root = "test_bridge.rs",
+    cxx_bridge_deps = ["//src/rust/jsg"],
+    cxx_bridge_hdrs = [],
+    cxx_bridge_src = "test_bridge.rs",
+    deps = [
+        ":io-context",
+        "//src/rust/jsg",
+    ],
+)
+
+# io-context inside a request of a real worker and in a Durable Object, and the IoContext that
+# Rust builtin modules see when a request resolves them.
+kj_test(
+    src = "io-context-test.c++",
+    deps = [
+        ":ffi",
+        ":io-context-test-bridge",
+        "//src/rust/jsg:bridge",
+        "//src/workerd/tests:test-fixture",
+    ],
+)

src/workerd/io/io-context/ffi.c++

@@ -0,0 +1,81 @@
+// Copyright (c) 2026 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+#include "ffi.h"
+
+#include <capnp/message.h>
+#include <capnp/serialize.h>
+
+#include <bit>
+#include <type_traits>
+
+namespace workerd::io_context {
+
+namespace {
+
+// IoContext::Id holds only the process-unique uint64_t that Rust stores.
+static_assert(sizeof(IoContext::Id) == sizeof(uint64_t));
+static_assert(std::is_trivially_copyable_v<IoContext::Id>);
+
+// The current IoContext, which must be the one `id` identifies.
+IoContext& currentWithId(uint64_t id) {
+  auto& context = KJ_REQUIRE_NONNULL(IoContext::tryCurrent(), "there is no current request");
+  KJ_REQUIRE(context.getId() == std::bit_cast<IoContext::Id>(id),
+      "the IoContext is not the one of the current request");
+  return context;
+}
+
+}  // namespace
+
+::rust::Str io_context_no_request_message() {
+  return IoContext::NO_REQUEST_MESSAGE.cStr();
+}
+
+kj::Maybe<uint64_t> io_context_try_current_id() {
+  return IoContext::tryGetCurrentId().map(
+      [](IoContext::Id id) { return std::bit_cast<uint64_t>(id); });
+}
+
+int64_t io_context_now(uint64_t id) {
+  return (currentWithId(id).now() - kj::UNIX_EPOCH) / kj::NANOSECONDS;
+}
+
+void io_context_write_logfwdr(uint64_t id, uint32_t channel, ::rust::Slice<const uint8_t> message) {
+  auto& context = currentWithId(id);
+
+  // Rust does not align the bytes to words, so they are copied.
+  auto bytes = kj::arrayPtr(message.data(), message.size());
+  KJ_REQUIRE(bytes.size() > 0, "the serialized message is empty");
+  KJ_REQUIRE(bytes.size() % sizeof(capnp::word) == 0,
+      "the serialized message is not a whole number of words");
+  auto words = kj::heapArray<capnp::word>(bytes.size() / sizeof(capnp::word));
+  words.asBytes().copyFrom(bytes);
+  capnp::FlatArrayMessageReader reader(words);
+  KJ_REQUIRE(reader.getEnd() == words.end(), "bytes follow the serialized message");
+
+  KJ_IF_SOME(outputLocks, context.waitForOutputLocksIfNecessary()) {
+    // In a Durable Object, the message must not leave before the storage writes that precede it
+    // are confirmed, as in AnalyticsEngine::writeDataPoint(). The message is copied now, so that
+    // an invalid one fails this call rather than the write.
+    auto copy = kj::heap<capnp::MallocMessageBuilder>();
+    copy->getRoot<capnp::AnyPointer>().set(reader.getRoot<capnp::AnyPointer>());
+    context.addWaitUntil(
+        outputLocks.then([weakRef = context.getWeakRef(), channel, copy = kj::mv(copy)]() mutable {
+      // The IoContext owns its waitUntil tasks, so it is alive while one runs.
+      auto& context = KJ_ASSERT_NONNULL(weakRef->tryGet());
+      // IoChannelFactory::writeLogfwdr() needs the current IoContext, so the write runs in it.
+      return context.run([channel, copy = kj::mv(copy)](Worker::Lock&, IoContext& context) mutable {
+        context.writeLogfwdr(channel, [&](capnp::AnyPointer::Builder builder) {
+          builder.set(copy->getRoot<capnp::AnyPointer>().asReader());
+        });
+      });
+    }));
+  } else {
+    context.writeLogfwdr(channel, [&](capnp::AnyPointer::Builder builder) {
+      builder.set(reader.getRoot<capnp::AnyPointer>());
+    });
+  }
+}
+
+}  // namespace workerd::io_context

src/workerd/io/io-context/ffi.h

@@ -0,0 +1,32 @@
+// Copyright (c) 2026 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+#pragma once
+
+#include <workerd/io/io-context.h>
+
+#include <rust/cxx.h>
+
+#include <cstdint>
+
+// The IoContext calls `ffi.rs` makes. Rust identifies an IoContext by its IoContext::Id, as a
+// uint64_t, and holds no reference to it. Each call below that takes an `id` uses the current
+// IoContext, and fails unless `id` identifies it: ids are never reused, so a stale id cannot
+// reach another request's IoContext, or one that has been destroyed.
+namespace workerd::io_context {
+
+// The message of the error IoContext::current() throws outside a request.
+::rust::Str io_context_no_request_message();
+
+// IoContext::tryGetCurrentId(), as a uint64_t.
+kj::Maybe<uint64_t> io_context_try_current_id();
+
+// IoContext::now(), in nanoseconds since the Unix epoch.
+int64_t io_context_now(uint64_t id);
+
+// IoContext::writeLogfwdr(), for a serialized Cap'n Proto message. In a Durable Object, the
+// write waits for the output gate, as AnalyticsEngine::writeDataPoint() does.
+void io_context_write_logfwdr(uint64_t id, uint32_t channel, ::rust::Slice<const uint8_t> message);
+
+}  // namespace workerd::io_context

src/workerd/io/io-context/ffi.rs

@@ -0,0 +1,19 @@
+// Copyright (c) 2026 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+//! The bridge to the C++ `IoContext` calls in `ffi.h`.
+
+#![expect(unsafe_code, reason = "the cxx bridge expands to unsafe FFI glue")]
+
+#[cxx::bridge(namespace = "workerd::io_context")]
+pub mod bridge {
+    unsafe extern "C++" {
+        include!("workerd/io/io-context/ffi.h");
+
+        fn io_context_no_request_message() -> &'static str;
+        fn io_context_try_current_id() -> KjMaybe<u64>;
+        fn io_context_now(id: u64) -> Result<i64>;
+        fn io_context_write_logfwdr(id: u64, channel: u32, message: &[u8]) -> Result<()>;
+    }
+}