make Immediate class thread-safe (#5207)

Author: anonrig

src/workerd/api/global-scope.c++

@@ -997,11 +997,14 @@ bool Navigator::sendBeacon(jsg::Lock& js, kj::String url, jsg::Optional<Body::In
 // ======================================================================================
 
 Immediate::Immediate(IoContext& context, TimeoutId timeoutId)
-    : contextRef(context.getWeakRef()),
+    : ioContext(context.addObject(context)),
       timeoutId(timeoutId) {}
 
 void Immediate::dispose() {
-  contextRef->runIfAlive([&](IoContext& context) { context.clearTimeoutImpl(timeoutId); });
+  // IoPtr will throw if the IoContext is no longer valid, which is fine - accessing the Immediate
+  // from the wrong IoContext should throw an error, just as accessing the Immediate when it's
+  // owning IoContext is gone.
+  ioContext->clearTimeoutImpl(timeoutId);
 }
 
 jsg::Ref<Immediate> ServiceWorkerGlobalScope::setImmediate(jsg::Lock& js,

src/workerd/api/global-scope.h

@@ -414,12 +414,10 @@ class Immediate final: public jsg::Object {
   }
 
  private:
-  // On the off chance user code holds onto to the Ref<Immediate> longer than
-  // the IoContext remains alive, let's maintain just a weak reference to the
-  // IoContext here to avoid problems. This reference is used only for handling
-  // the dispose operation, so it should be perfectly fine for it to be weak
-  // and a non-op after the IoContext is gone.
-  kj::Own<IoContext::WeakRef> contextRef;
+  // Note: We cannot use IoContext::WeakRef here because it's not thread-safe (it's only intended
+  // to be held from KJ I/O objects, but this is a JSG object which can be accessed by V8's GC
+  // on different threads). Instead, we use IoPtr<IoContext> which is safe to hold from JSG objects.
+  IoPtr<IoContext> ioContext;
   TimeoutId timeoutId;
 };
 

src/workerd/api/node/tests/timers-nodejs-test.js

@@ -496,3 +496,27 @@ export const timersIntervalPromise = {
     }
   },
 };
+
+export const testImmediateFromOtherContext = {
+  async test(_, env) {
+    // Set up the global setImmediate...
+    await env.SUBREQUEST.fetch('http://example.com');
+
+    // Try to access it from a different IoContext.
+    await env.SUBREQUEST.fetch('http://example.com');
+  },
+};
+
+export default {
+  fetch() {
+    if (globalThis.IMMEDIATE_TEST === undefined) {
+      globalThis.IMMEDIATE_TEST = setImmediate(() => {});
+      globalThis.IMMEDIATE_TEST[Symbol.dispose]();
+    } else {
+      throws(() => globalThis.IMMEDIATE_TEST[Symbol.dispose](), {
+        message: /perform I\/O on behalf of a different request/,
+      });
+    }
+    return new Response();
+  },
+};

src/workerd/api/node/tests/timers-nodejs-test.wd-test

@@ -7,8 +7,11 @@ const unitTests :Workerd.Config = (
         modules = [
           (name = "worker", esModule = embed "timers-nodejs-test.js")
         ],
-        compatibilityDate = "2025-01-09",
-        compatibilityFlags = ["nodejs_compat", "no_nodejs_compat_v2"],
+        compatibilityDate = "2025-10-04",
+        compatibilityFlags = ["nodejs_compat"],
+        bindings = [
+          ( name = "SUBREQUEST", service = "nodejs-timers-test" ),
+        ],
       )
     ),
   ],