Fix `setTimeout(0)`, improve test coverage slightly, reduce memory usage of timers

5 files changed, 221 insertions, 87 deletions

diff --git a/src/bun.js/api/bun.zig b/src/bun.js/api/bun.zig
index c173391cf..5c8bfbb0e 100644
--- a/src/bun.js/api/bun.zig
+++ b/src/bun.js/api/bun.zig
@@ -2201,10 +2201,24 @@ pub const TOML = struct {
 pub const Timer = struct {
     last_id: i32 = 1,
     warned: bool = false,
-    active: u32 = 0,
-    timeouts: TimeoutMap = TimeoutMap{},
 
-    const TimeoutMap = std.AutoArrayHashMapUnmanaged(i32, *Timeout);
+    /// Used by setTimeout()
+    timeout_map: TimeoutMap = TimeoutMap{},
+
+    /// Used by setInterval()
+    interval_map: TimeoutMap = TimeoutMap{},
+
+    /// TimeoutMap is map of i32 to nullable Timeout structs
+    /// i32 is exposed to JavaScript and can be used with clearTimeout, clearInterval, etc.
+    /// When Timeout is null, it means the tasks have been scheduled but not yet executed.
+    /// Timeouts are enqueued as a task to be run on the next tick of the task queue
+    /// The task queue runs after the event loop tasks have been run
+    /// Therefore, there is a race condition where you cancel the task after it has already been enqueued
+    /// In that case, it shouldn't run. It should be skipped.
+    pub const TimeoutMap = std.AutoArrayHashMapUnmanaged(
+        i32,
+        ?Timeout,
+    );
 
     pub fn getNextID() callconv(.C) i32 {
         VirtualMachine.vm.timer.last_id +%= 1;
@@ -2212,42 +2226,100 @@ pub const Timer = struct {
     }
 
     const uws = @import("uws");
-    pub const Timeout = struct {
+
+    // TODO: reference count to avoid multiple Strong references to the same
+    // object in setInterval
+    const CallbackJob = struct {
         id: i32 = 0,
+        task: JSC.AnyTask = undefined,
+        ref: JSC.Ref = JSC.Ref.init(),
+        globalThis: *JSC.JSGlobalObject,
         callback: JSC.Strong = .{},
-        interval: i32 = 0,
         repeat: bool = false,
+
+        pub const Task = JSC.AnyTask.New(CallbackJob, perform);
+
+        pub fn perform(this: *CallbackJob) void {
+            var vm = this.globalThis.bunVM();
+            var map: *TimeoutMap = if (this.repeat) &vm.timer.interval_map else &vm.timer.timeout_map;
+
+            defer {
+                this.callback.deinit();
+                this.ref.unref(this.globalThis.bunVM());
+                bun.default_allocator.destroy(this);
+            }
+
+            // This doesn't deinit the timer
+            // Timers are deinit'd separately
+            // We do need to handle when the timer is cancelled after the job has been enqueued
+            if (!this.repeat) {
+                if (map.fetchSwapRemove(this.id) == null) {
+                    // if the timeout was cancelled, don't run the callback
+                    return;
+                }
+            } else {
+                if (!map.contains(this.id)) {
+                    // if the interval was cancelled, don't run the callback
+                    return;
+                }
+            }
+
+            const callback = this.callback.get() orelse @panic("Expected CallbackJob to have a callback function");
+
+            const result = callback.call(this.globalThis, &.{});
+
+            if (result.isAnyError(this.globalThis)) {
+                vm.runErrorHandler(result, null);
+            }
+        }
+    };
+
+    pub const Timeout = struct {
+        callback: JSC.Strong = .{},
         globalThis: *JSC.JSGlobalObject,
         timer: *uws.Timer,
         poll_ref: JSC.PollRef = JSC.PollRef.init(),
 
-        pub var invalid_timer_ref: *Timeout = undefined;
+        // this is sized to be the same as one pointer
+        pub const ID = extern struct {
+            id: i32,
+
+            repeat: u32 = 0,
+        };
 
         pub fn run(timer: *uws.Timer) callconv(.C) void {
-            const id: i32 = timer.as(i32);
+            const timer_id: ID = timer.as(ID);
 
             // use the threadlocal despite being slow on macOS
             // to handle the timeout being cancelled after already enqueued
             var vm = JSC.VirtualMachine.vm;
 
-            var this_entry = vm.timer.timeouts.getEntry(
-                id,
-            ) orelse {
-                // this timer was cancelled after the event loop callback was queued
-                return;
-            };
+            const repeats = timer_id.repeat > 0;
 
-            var this: *Timeout = this_entry.value_ptr.*;
+            var map = if (repeats) &vm.timer.interval_map else &vm.timer.timeout_map;
 
-            std.debug.assert(this != invalid_timer_ref);
+            var this_: ?Timeout = map.get(
+                timer_id.id,
+            ) orelse return;
+            var this = this_ orelse
+                return;
 
             var cb: CallbackJob = .{
-                .callback = if (this.repeat)
-                    JSC.Strong.create(this.callback.get() orelse return, this.globalThis)
+                .callback = if (repeats)
+                    JSC.Strong.create(this.callback.get() orelse {
+                        // if the callback was freed, that's an error
+                        if (comptime Environment.allow_assert)
+                            unreachable;
+
+                        this.deinit();
+                        _ = map.swapRemove(timer_id.id);
+                        return;
+                    }, this.globalThis)
                 else
                     this.callback,
                 .globalThis = this.globalThis,
-                .id = this.id,
+                .id = timer_id.id,
+                .repeat = timer_id.repeat > 0,
             };
 
             var job = vm.allocator.create(CallbackJob) catch @panic(
@@ -2263,48 +2335,13 @@ pub const Timer = struct {
             // This allows us to:
             //  - free the memory before the job is run
             //  - reuse the JSC.Strong
-            if (!this.repeat) {
+            if (!repeats) {
                 this.callback = .{};
-                this_entry.value_ptr.* = invalid_timer_ref;
+                map.put(vm.allocator, timer_id.id, null) catch unreachable;
                 this.deinit();
             }
         }
 
-        // TODO: reference count to avoid multiple Strong references to the same
-        // object in setInterval
-        const CallbackJob = struct {
-            id: i32 = 0,
-            task: JSC.AnyTask = undefined,
-            ref: JSC.Ref = JSC.Ref.init(),
-            globalThis: *JSC.JSGlobalObject,
-            callback: JSC.Strong = .{},
-
-            pub const Task = JSC.AnyTask.New(CallbackJob, perform);
-
-            pub fn perform(this: *CallbackJob) void {
-                defer {
-                    this.callback.deinit();
-                    this.ref.unref(this.globalThis.bunVM());
-                    bun.default_allocator.destroy(this);
-                }
-                var vm = this.globalThis.bunVM();
-
-                if (!vm.timer.timeouts.contains(this.id)) {
-                    // we didn't find the timeout, so it was already cleared
-                    // that means this job shouldn't run.
-                    return;
-                }
-
-                const callback = this.callback.get() orelse @panic("Expected callback in timer");
-
-                const result = callback.call(this.globalThis, &.{});
-
-                if (result.isAnyError(this.globalThis)) {
-                    vm.runErrorHandler(result, null);
-                }
-            }
-        };
-
         pub fn deinit(this: *Timeout) void {
             if (comptime JSC.is_bindgen)
                 unreachable;
@@ -2313,7 +2350,6 @@ pub const Timer = struct {
             this.poll_ref.unref(vm);
             this.timer.deinit();
             this.callback.deinit();
-            vm.allocator.destroy(this);
         }
     };
 
@@ -2326,23 +2362,66 @@ pub const Timer = struct {
     ) !void {
         if (comptime is_bindgen) unreachable;
         var vm = globalThis.bunVM();
-        var timeout = try vm.allocator.create(Timeout);
-        timeout.* = Timeout{
-            .id = id,
+
+        // We don't deal with nesting levels directly
+        // but we do set the minimum timeout to be 1ms for repeating timers
+        const interval: i32 = @maximum(
+            countdown.toInt32(),
+            if (repeat) @as(i32, 1) else 0,
+        );
+
+        var map = if (repeat)
+            &vm.timer.interval_map
+        else
+            &vm.timer.timeout_map;
+
+        // setImmediate(foo)
+        // setTimeout(foo, 0)
+        if (interval == 0) {
+            var cb: CallbackJob = .{
+                .callback = JSC.Strong.create(callback, globalThis),
+                .globalThis = globalThis,
+                .id = id,
+                .repeat = false,
+            };
+
+            var job = vm.allocator.create(CallbackJob) catch @panic(
+                "Out of memory while allocating Timeout",
+            );
+
+            job.* = cb;
+            job.task = CallbackJob.Task.init(job);
+            job.ref.ref(vm);
+
+            vm.enqueueTask(JSC.Task.init(&job.task));
+            map.put(vm.allocator, id, null) catch unreachable;
+            return;
+        }
+
+        var timeout = Timeout{
             .callback = JSC.Strong.create(callback, globalThis),
-            .interval = countdown.toInt32(),
-            .repeat = repeat,
             .globalThis = globalThis,
-            .timer = uws.Timer.create(vm.uws_event_loop.?, false, timeout),
+            .timer = uws.Timer.create(
+                vm.uws_event_loop.?,
+                true,
+                Timeout.ID{
+                    .id = id,
+                    .repeat = @as(u32, @boolToInt(repeat)),
+                },
+            ),
         };
         timeout.timer.set(
-            id,
+            Timeout.ID{
+                .id = id,
+                .repeat = if (repeat) 1 else 0,
+            },
             Timeout.run,
-            timeout.interval,
-            @as(i32, @boolToInt(repeat)) * timeout.interval,
+            interval,
+            @as(i32, @boolToInt(repeat)) * interval,
         );
         timeout.poll_ref.ref(vm);
-        vm.timer.timeouts.put(vm.allocator, id, timeout) catch unreachable;
+
+        map.put(vm.allocator, id, timeout) catch unreachable;
     }
 
     pub fn setTimeout(
@@ -2374,17 +2453,22 @@ pub const Timer = struct {
         return JSValue.jsNumberWithType(i32, id);
     }
 
-    pub fn clearTimer(id: JSValue, _: *JSGlobalObject) void {
+    pub fn clearTimer(timer_id: JSValue, _: *JSGlobalObject, repeats: bool) void {
         if (comptime is_bindgen) unreachable;
 
-        var timer = VirtualMachine.vm.timer.timeouts.fetchSwapRemove(id.toInt32()) orelse return;
-        if (timer.value == Timeout.invalid_timer_ref) {
+        var map = if (repeats) &VirtualMachine.vm.timer.interval_map else &VirtualMachine.vm.timer.timeout_map;
+        const id: Timeout.ID = .{
+            .id = timer_id.toInt32(),
+            .repeat = @as(u32, @boolToInt(repeats)),
+        };
+        var timer = map.fetchSwapRemove(id.id) orelse return;
+        if (timer.value == null) {
             // this timer was scheduled to run but was cancelled before it was run
             // so long as the callback isn't already in progress, fetchSwapRemove will handle invalidating it
             return;
         }
 
-        timer.value.deinit();
+        timer.value.?.deinit();
     }
 
     pub fn clearTimeout(
@@ -2392,7 +2476,7 @@ pub const Timer = struct {
         id: JSValue,
     ) callconv(.C) JSValue {
         if (comptime is_bindgen) unreachable;
-        Timer.clearTimer(id, globalThis);
+        Timer.clearTimer(id, globalThis, false);
         return JSValue.jsUndefined();
     }
     pub fn clearInterval(
@@ -2400,7 +2484,7 @@ pub const Timer = struct {
         id: JSValue,
     ) callconv(.C) JSValue {
         if (comptime is_bindgen) unreachable;
-        Timer.clearTimer(id, globalThis);
+        Timer.clearTimer(id, globalThis, true);
         return JSValue.jsUndefined();
     }
 
diff --git a/src/bun.js/base.zig b/src/bun.js/base.zig
index 0770c77a4..23b00bd8e 100644
--- a/src/bun.js/base.zig
+++ b/src/bun.js/base.zig
@@ -3902,6 +3902,8 @@ pub const Ref = struct {
 pub const PollRef = struct {
     status: Status = .inactive,
 
+    const log = Output.scoped(.PollRef, true);
+
     const Status = enum { active, inactive, done };
 
     pub inline fn isActive(this: PollRef) bool {
@@ -3943,6 +3945,7 @@ pub const PollRef = struct {
         if (this.status != .active)
             return;
         this.status = .inactive;
+        log("unref", .{});
         vm.uws_event_loop.?.num_polls -= 1;
         vm.uws_event_loop.?.active -= 1;
     }
@@ -3951,6 +3954,7 @@ pub const PollRef = struct {
     pub fn ref(this: *PollRef, vm: *JSC.VirtualMachine) void {
         if (this.status != .inactive)
             return;
+        log("ref", .{});
         this.status = .active;
         vm.uws_event_loop.?.num_polls += 1;
         vm.uws_event_loop.?.active += 1;
diff --git a/src/bun.js/event_loop.zig b/src/bun.js/event_loop.zig
index a16c778a5..4c36f1c05 100644
--- a/src/bun.js/event_loop.zig
+++ b/src/bun.js/event_loop.zig
@@ -48,6 +48,8 @@ pub fn ConcurrentPromiseTask(comptime Context: type) type {
                 .globalThis = globalThis,
             };
             this.promise.protect();
+            this.ref.ref(this.event_loop.virtual_machine);
+
             return this;
         }
 
@@ -57,10 +59,13 @@ pub fn ConcurrentPromiseTask(comptime Context: type) type {
             this.onFinish();
         }
 
-        pub fn runFromJS(this: This) void {
+        pub fn runFromJS(this: *This) void {
             var promise_value = this.promise;
+            this.ref.unref(this.event_loop.virtual_machine);
+
             promise_value.ensureStillAlive();
             promise_value.unprotect();
+
             var promise = promise_value.asInternalPromise() orelse {
                 if (comptime @hasDecl(Context, "deinit")) {
                     @call(.{}, Context.deinit, .{this.ctx});
@@ -74,8 +79,6 @@ pub fn ConcurrentPromiseTask(comptime Context: type) type {
         }
 
         pub fn schedule(this: *This) void {
-            this.ref.ref(this.event_loop.virtual_machine);
-
             WorkPool.schedule(&this.task);
         }
 
@@ -84,7 +87,6 @@ pub fn ConcurrentPromiseTask(comptime Context: type) type {
         }
 
         pub fn deinit(this: *This) void {
-            this.ref.unref(this.event_loop.virtual_machine);
             this.allocator.destroy(this);
         }
     };
@@ -106,11 +108,12 @@ pub fn IOTask(comptime Context: type) type {
         pub fn createOnJSThread(allocator: std.mem.Allocator, globalThis: *JSGlobalObject, value: *Context) !*This {
             var this = try allocator.create(This);
             this.* = .{
-                .event_loop = VirtualMachine.vm.eventLoop(),
+                .event_loop = globalThis.bunVM().eventLoop(),
                 .ctx = value,
                 .allocator = allocator,
                 .globalThis = globalThis,
             };
+            this.ref.ref(this.event_loop.virtual_machine);
 
             return this;
         }
@@ -120,8 +123,9 @@ pub fn IOTask(comptime Context: type) type {
             Context.run(this.ctx, this);
         }
 
-        pub fn runFromJS(this: This) void {
+        pub fn runFromJS(this: *This) void {
             var ctx = this.ctx;
+            this.ref.unref(this.event_loop.virtual_machine);
             ctx.then(this.globalThis);
         }
 
@@ -376,7 +380,7 @@ pub const EventLoop = struct {
                     this.tick();
 
                     if (promise.status(this.global.vm()) == .Pending) {
-                        if (this.tickConcurrentWithCount() == 0) {
+                        if (this.virtual_machine.uws_event_loop.?.active > 0 or this.virtual_machine.uws_event_loop.?.num_polls > 0) {
                             this.virtual_machine.uws_event_loop.?.tick();
                         }
                     }
diff --git a/src/bun.js/javascript.zig b/src/bun.js/javascript.zig
index 8eeca5555..7f41ea717 100644
--- a/src/bun.js/javascript.zig
+++ b/src/bun.js/javascript.zig
@@ -295,14 +295,10 @@ pub export fn Bun__onDidAppendPlugin(jsc_vm: *VirtualMachine, globalObject: *JSG
     jsc_vm.bundler.linker.plugin_runner = &jsc_vm.plugin_runner.?;
 }
 
-// If you read JavascriptCore/API/JSVirtualMachine.mm - https://github.com/WebKit/WebKit/blob/acff93fb303baa670c055cb24c2bad08691a01a0/Source/JavaScriptCore/API/JSVirtualMachine.mm#L101
-// We can see that it's sort of like std.mem.Allocator but for JSGlobalContextRef, to support Automatic Reference Counting
-// Its unavailable on Linux
-
-// JavaScriptCore expects 1 VM per thread
-// However, there can be many JSGlobalObject
-// We currently assume a 1:1 correspondence between the two.
-// This is technically innacurate
+/// TODO: rename this to ScriptExecutionContext
+/// This is the shared global state for a single JS instance execution
+/// Today, Bun is one VM per thread, so the name "VirtualMachine" sort of makes sense
+/// However, that may change in the future
 pub const VirtualMachine = struct {
     global: *JSGlobalObject,
     allocator: std.mem.Allocator,
@@ -451,7 +447,13 @@ pub const VirtualMachine = struct {
 
     pub const MacroMap = std.AutoArrayHashMap(i32, js.JSObjectRef);
 
+    /// Threadlocals are slow on macOS
     pub threadlocal var vm_loaded = false;
+
+    /// Threadlocals are slow on macOS
+    /// Consider using `globalThis.bunVM()` instead.
+    /// There may be a time where we run multiple VMs in the same thread
+    /// At that point, this threadlocal will be a problem.
     pub threadlocal var vm: *VirtualMachine = undefined;
 
     pub fn enableMacroMode(this: *VirtualMachine) void {
diff --git a/test/bun.js/setTimeout.test.js b/test/bun.js/setTimeout.test.js
index 79039a568..a899867df 100644
--- a/test/bun.js/setTimeout.test.js
+++ b/test/bun.js/setTimeout.test.js
@@ -41,3 +41,43 @@ it("clearTimeout", async () => {
   });
   expect(called).toBe(false);
 });
+
+it("setTimeout(() => {}, 0)", async () => {
+  var called = false;
+  setTimeout(() => {
+    called = true;
+  }, 0);
+  await new Promise((resolve, reject) => {
+    setTimeout(() => {
+      resolve();
+    }, 10);
+  });
+  expect(called).toBe(true);
+  var ranFirst = -1;
+  setTimeout(() => {
+    if (ranFirst === -1) ranFirst = 1;
+  }, 1);
+  setTimeout(() => {
+    if (ranFirst === -1) ranFirst = 0;
+  }, 0);
+
+  await new Promise((resolve, reject) => {
+    setTimeout(() => {
+      resolve();
+    }, 10);
+  });
+  expect(ranFirst).toBe(0);
+
+  ranFirst = -1;
+
+  const id = setTimeout(() => {
+    ranFirst = 0;
+  }, 0);
+  clearTimeout(id);
+  await new Promise((resolve, reject) => {
+    setTimeout(() => {
+      resolve();
+    }, 10);
+  });
+  expect(ranFirst).toBe(-1);
+});