Merge pull request #2506 from cloudflare/jsnell/implement-setimmediat…

…e-for-node-compat-v2

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

@@ -915,4 +915,55 @@ bool Navigator::sendBeacon(jsg::Lock& js, kj::String url,
   return false;
 }
 
+// ======================================================================================
+
+Immediate::Immediate(IoContext& context, TimeoutId timeoutId)
+    : contextRef(context.getWeakRef()),
+      timeoutId(timeoutId) {}
+
+void Immediate::dispose() {
+  contextRef->runIfAlive([&](IoContext& context) {
+    context.clearTimeoutImpl(timeoutId);
+  });
+}
+
+jsg::Ref<Immediate> ServiceWorkerGlobalScope::setImmediate(
+    jsg::Lock& js,
+    jsg::Function<void(jsg::Arguments<jsg::Value>)> function,
+    jsg::Arguments<jsg::Value> args) {
+
+  // This is an approximation of the Node.js setImmediate global API.
+  // We implement it in terms of setting a 0 ms timeout. This is not
+  // how Node.js does it so there will be some edge cases where the
+  // timing of the callback will differ relative to the equivalent
+  // operations in Node.js. For the vast majority of cases, users
+  // really shouldn't be able to tell a difference. It would likely
+  // only be somewhat pathological edge cases that could be affected
+  // by the differences. Unfortunately, changing this later to match
+  // Node.js would likely be a breaking change for some users that
+  // would require a compat flag... but that's ok for now?
+
+  auto& context = IoContext::current();
+  auto fn = [function=kj::mv(function),
+             args=kj::mv(args),
+             context=jsg::AsyncContextFrame::currentRef(js)](jsg::Lock& js) mutable {
+    jsg::AsyncContextFrame::Scope scope(js, context);
+    function(js, kj::mv(args));
+  };
+  auto timeoutId = context.setTimeoutImpl(
+      timeoutIdGenerator,
+      /* repeats = */ false,
+      [function = kj::mv(fn)](jsg::Lock& js) mutable {
+    function(js);
+  }, 0);
+  return jsg::alloc<Immediate>(context, timeoutId);
+}
+
+void ServiceWorkerGlobalScope::clearImmediate(
+    kj::Maybe<jsg::Ref<Immediate>> maybeImmediate) {
+  KJ_IF_SOME(immediate, maybeImmediate) {
+    immediate->dispose();
+  }
+}
+
 } // namespace workerd::api

src/workerd/api/global-scope.h

@@ -312,6 +312,40 @@ struct ExportedHandler {
   }
 };
 
+// An approximation of Node.js setImmediate `Immediate` object.
+// This is used only when the `nodejs_compat_v2` compatibility flag is enabled.
+class Immediate final: public jsg::Object {
+public:
+  Immediate(IoContext& context, TimeoutId timeoutId);
+
+  // In Node.js, the "ref" mechanism refers to whether or not an i/o object
+  // will keep the libuv event loop alive (and therefore keep the process alive).
+  // We do not implement a similar mechanism in workerd. These are here only to
+  // satisfy the API contract for the `Immediate` object but are never expected
+  // to actually do anything.
+  bool hasRef() { return false; }
+  void ref() { /* non-op */ }
+  void unref() { /* non-op */ }
+
+  void dispose();
+
+  JSG_RESOURCE_TYPE(Immediate) {
+    JSG_METHOD(ref);
+    JSG_METHOD(unref);
+    JSG_METHOD(hasRef);
+    JSG_DISPOSE(dispose);
+  }
+
+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 dipose 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;
+  TimeoutId timeoutId;
+};
+
 // Global object API exposed to JavaScript.
 class ServiceWorkerGlobalScope: public WorkerGlobalScope {
 public:
@@ -460,6 +494,10 @@ class ServiceWorkerGlobalScope: public WorkerGlobalScope {
   // properties.
   jsg::JsValue getBuffer(jsg::Lock& js);
   jsg::JsValue getProcess(jsg::Lock& js);
+  jsg::Ref<Immediate> setImmediate(jsg::Lock& js,
+                                   jsg::Function<void(jsg::Arguments<jsg::Value>)> function,
+                                   jsg::Arguments<jsg::Value> args);
+  void clearImmediate(kj::Maybe<jsg::Ref<Immediate>> immediate);
 
   JSG_RESOURCE_TYPE(ServiceWorkerGlobalScope, CompatibilityFlags::Reader flags) {
     JSG_INHERIT(WorkerGlobalScope);
@@ -540,6 +578,8 @@ class ServiceWorkerGlobalScope: public WorkerGlobalScope {
       JSG_LAZY_INSTANCE_PROPERTY(Buffer, getBuffer);
       JSG_LAZY_INSTANCE_PROPERTY(process, getProcess);
       JSG_LAZY_INSTANCE_PROPERTY(global, getSelf);
+      JSG_METHOD(setImmediate);
+      JSG_METHOD(clearImmediate);
     }
 
     JSG_NESTED_TYPE(CompressionStream);
@@ -762,6 +802,7 @@ class ServiceWorkerGlobalScope: public WorkerGlobalScope {
   api::PromiseRejectionEvent,                            \
   api::Navigator,                                        \
   api::Performance,                                      \
-  api::AlarmInvocationInfo
+  api::AlarmInvocationInfo,                              \
+  api::Immediate
 // The list of global-scope.h types that are added to worker.c++'s JSG_DECLARE_ISOLATE_TYPE
 }  // namespace workerd::api

src/workerd/api/node/BUILD.bazel

@@ -1,5 +1,6 @@
 load("//:build/kj_test.bzl", "kj_test")
 load("//:build/wd_cc_library.bzl", "wd_cc_library")
+load("//:build/wd_test.bzl", "wd_test")
 
 wd_cc_library(
     name = "node",
@@ -21,3 +22,11 @@ kj_test(
     src = "buffer-test.c++",
     deps = ["//src/workerd/tests:test-fixture"],
 )
+
+[wd_test(
+    src = f,
+    args = ["--experimental"],
+    data = [f.removesuffix(".wd-test") + ".js"],
+) for f in glob(
+    ["**/*.wd-test"],
+)]

src/workerd/api/node/tests/node-compat-v2-test.js

@@ -7,6 +7,8 @@
 // the 'node:' prefix.
 import { default as assert } from 'node:assert';
 import { default as assert2 } from 'assert';
+import { AsyncLocalStorage } from 'async_hooks';
+
 const assert3 = (await import('node:assert')).default;
 const assert4 = (await import('assert')).default;
 
@@ -86,3 +88,27 @@ export const nodeJsBufferExports = {
     assert.strictEqual(Blob, globalThis.Blob);
   }
 };
+
+export const nodeJsSetImmediate = {
+  async test() {
+    const als = new AsyncLocalStorage();
+    const { promise, resolve } = Promise.withResolvers();
+    als.run('abc', () => setImmediate((a) => {
+      assert.strictEqual(als.getStore(), 'abc');
+      resolve(a);
+    }, 1));
+    assert.strictEqual(await promise, 1);
+
+    const i = setImmediate(() => {
+      throw new Error('should not have fired');
+    });
+    i[Symbol.dispose]();  // Calls clear immediate
+    i[Symbol.dispose]();  // Should be a no-op
+
+    const i2 = setImmediate(() => {
+      throw new Error('should not have fired');
+    });
+    clearImmediate(i2);
+    clearImmediate(i2); // clearing twice works fine
+  }
+};