Reduce memory usage in long-running processes (#14885)

oven-sh · Oct 29, 2024 · b5a7313 · b5a7313
1 parent d5f9978
commit b5a7313
Show file tree

Hide file tree

Showing 16 changed files with 120 additions and 60 deletions.
diff --git a/packages/bun-usockets/src/eventing/epoll_kqueue.c b/packages/bun-usockets/src/eventing/epoll_kqueue.c
@@ -230,6 +230,8 @@ void us_loop_run(struct us_loop_t *loop) {
     }
 }
 
+extern int Bun__JSC_onBeforeWait(void*);
+extern void Bun__JSC_onAfterWait(void*);
 
 void us_loop_run_bun_tick(struct us_loop_t *loop, const struct timespec* timeout) {
     if (loop->num_polls == 0)
@@ -246,6 +248,11 @@ void us_loop_run_bun_tick(struct us_loop_t *loop, const struct timespec* timeout
     /* Emit pre callback */
     us_internal_loop_pre(loop);
 
+    int needs_after_wait = 0;
+    if (loop->data.jsc_vm) {
+        needs_after_wait = Bun__JSC_onBeforeWait(loop->data.jsc_vm);
+    }
+
     /* Fetch ready polls */
 #ifdef LIBUS_USE_EPOLL
 
@@ -256,6 +263,10 @@ void us_loop_run_bun_tick(struct us_loop_t *loop, const struct timespec* timeout
     } while (IS_EINTR(loop->num_ready_polls));
 #endif
 
+    if (needs_after_wait) {
+        Bun__JSC_onAfterWait(loop->data.jsc_vm);
+    }
+
     /* Iterate ready polls, dispatching them by type */
     for (loop->current_ready_poll = 0; loop->current_ready_poll < loop->num_ready_polls; loop->current_ready_poll++) {
         struct us_poll_t *poll = GET_READY_POLL(loop, loop->current_ready_poll);

diff --git a/packages/bun-usockets/src/internal/loop_data.h b/packages/bun-usockets/src/internal/loop_data.h
@@ -44,6 +44,7 @@ struct us_internal_loop_data_t {
     char parent_tag;
     /* We do not care if this flips or not, it doesn't matter */
     size_t iteration_nr;
+    void* jsc_vm;
 };
 
 #endif // LOOP_DATA_H
diff --git a/packages/bun-usockets/src/loop.c b/packages/bun-usockets/src/loop.c
@@ -48,6 +48,7 @@ void us_internal_loop_data_init(struct us_loop_t *loop, void (*wakeup_cb)(struct
     loop->data.parent_tag = 0;
 
     loop->data.closed_context_head = 0;
+    loop->data.jsc_vm = 0;
 
     loop->data.wakeup_async = us_internal_create_async(loop, 1, 0);
     us_internal_async_set(loop->data.wakeup_async, (void (*)(struct us_internal_async *)) wakeup_cb);

diff --git a/src/async/posix_event_loop.zig b/src/async/posix_event_loop.zig
@@ -55,11 +55,11 @@ pub const KeepAlive = struct {
         this.status = .inactive;
 
         if (comptime @TypeOf(event_loop_ctx_) == JSC.EventLoopHandle) {
-            event_loop_ctx_.loop().subActive(1);
+            event_loop_ctx_.loop().unref();
             return;
         }
         const event_loop_ctx = JSC.AbstractVM(event_loop_ctx_);
-        event_loop_ctx.platformEventLoop().subActive(1);
+        event_loop_ctx.platformEventLoop().unref();
     }
 
     /// From another thread, Prevent a poll from keeping the process alive.

diff --git a/src/bun.js/bindings/BunJSCEventLoop.cpp b/src/bun.js/bindings/BunJSCEventLoop.cpp
@@ -0,0 +1,18 @@
+#include "root.h"
+
+#include <JavaScriptCore/VM.h>
+#include <JavaScriptCore/Heap.h>
+
+extern "C" int Bun__JSC_onBeforeWait(JSC::VM* vm)
+{
+    if (vm->heap.hasAccess()) {
+        vm->heap.releaseAccess();
+        return 1;
+    }
+    return 0;
+}
+
+extern "C" void Bun__JSC_onAfterWait(JSC::VM* vm)
+{
+    vm->heap.acquireAccess();
+}
diff --git a/src/bun.js/bindings/JSCUSocketsLoopIntegration.cpp b/src/bun.js/bindings/JSCUSocketsLoopIntegration.cpp
diff --git a/src/bun.js/bindings/bindings.zig b/src/bun.js/bindings/bindings.zig
@@ -6161,6 +6161,22 @@ pub const VM = extern struct {
         LargeHeap = 1,
     };
 
+    extern fn Bun__JSC_onBeforeWait(vm: *VM) i32;
+    extern fn Bun__JSC_onAfterWait(vm: *VM) void;
+    pub const ReleaseHeapAccess = struct {
+        vm: *VM,
+        needs_to_release: bool,
+        pub fn acquire(this: *const ReleaseHeapAccess) void {
+            if (this.needs_to_release) {
+                Bun__JSC_onAfterWait(this.vm);
+            }
+        }
+    };
+
+    pub fn releaseHeapAccess(vm: *VM) ReleaseHeapAccess {
+        return .{ .vm = vm, .needs_to_release = Bun__JSC_onBeforeWait(vm) != 0 };
+    }
+
     pub fn create(heap_type: HeapType) *VM {
         return cppFn("create", .{@intFromEnum(heap_type)});
     }

diff --git a/src/bun.js/event_loop.zig b/src/bun.js/event_loop.zig
@@ -545,6 +545,7 @@ pub const GarbageCollectionController = struct {
         const actual = uws.Loop.get();
         this.gc_timer = uws.Timer.createFallthrough(actual, this);
         this.gc_repeating_timer = uws.Timer.createFallthrough(actual, this);
+        actual.internal_loop_data.jsc_vm = vm.jsc;
 
         if (comptime Environment.isDebug) {
             if (bun.getenvZ("BUN_TRACK_LAST_FN_NAME") != null) {

diff --git a/src/bun.js/module_loader.zig b/src/bun.js/module_loader.zig
@@ -1646,6 +1646,12 @@ pub const ModuleLoader = struct {
                 var parse_result: ParseResult = switch (disable_transpilying or
                     (loader == .json and !path.isJSONCFile())) {
                     inline else => |return_file_only| brk: {
+                        const heap_access = if (!disable_transpilying)
+                            jsc_vm.jsc.releaseHeapAccess()
+                        else
+                            JSC.VM.ReleaseHeapAccess{ .vm = jsc_vm.jsc, .needs_to_release = false };
+                        defer heap_access.acquire();
+
                         break :brk jsc_vm.bundler.parseMaybeReturnFileOnly(
                             parse_options,
                             null,

diff --git a/src/bun.js/node/node_zlib_binding.zig b/src/bun.js/node/node_zlib_binding.zig
@@ -114,35 +114,25 @@ pub fn CompressionStream(comptime T: type) type {
             //
 
             const vm = globalThis.bunVM();
-            var task = AsyncJob.new(.{
-                .binding = this,
-            });
+            this.task = .{ .callback = &AsyncJob.runTask };
             this.poll_ref.ref(vm);
-            JSC.WorkPool.schedule(&task.task);
+            JSC.WorkPool.schedule(&this.task);
 
             return .undefined;
         }
 
         const AsyncJob = struct {
-            task: JSC.WorkPoolTask = .{ .callback = &runTask },
-            binding: *T,
-
-            pub usingnamespace bun.New(@This());
-
-            pub fn runTask(this: *JSC.WorkPoolTask) void {
-                var job: *AsyncJob = @fieldParentPtr("task", this);
-                job.run();
-                job.destroy();
+            pub fn runTask(task: *JSC.WorkPoolTask) void {
+                const this: *T = @fieldParentPtr("task", task);
+                AsyncJob.run(this);
             }
 
-            pub fn run(job: *AsyncJob) void {
-                const this = job.binding;
+            pub fn run(this: *T) void {
                 const globalThis: *JSC.JSGlobalObject = this.globalThis;
                 const vm = globalThis.bunVMConcurrently();
 
                 this.stream.doWork();
 
-                this.poll_ref.refConcurrently(vm);
                 vm.enqueueTaskConcurrent(JSC.ConcurrentTask.create(JSC.Task.init(this)));
             }
         };
@@ -294,6 +284,29 @@ pub fn CompressionStream(comptime T: type) type {
 
 pub const NativeZlib = JSC.Codegen.JSNativeZlib.getConstructor;
 
+const CountedKeepAlive = struct {
+    keep_alive: bun.Async.KeepAlive = .{},
+    ref_count: u32 = 0,
+
+    pub fn ref(this: *@This(), vm: *JSC.VirtualMachine) void {
+        if (this.ref_count == 0) {
+            this.keep_alive.ref(vm);
+        }
+        this.ref_count += 1;
+    }
+
+    pub fn unref(this: *@This(), vm: *JSC.VirtualMachine) void {
+        this.ref_count -= 1;
+        if (this.ref_count == 0) {
+            this.keep_alive.unref(vm);
+        }
+    }
+
+    pub fn deinit(this: *@This()) void {
+        this.keep_alive.disable();
+    }
+};
+
 pub const SNativeZlib = struct {
     pub usingnamespace bun.NewRefCounted(@This(), deinit);
     pub usingnamespace JSC.Codegen.JSNativeZlib;
@@ -306,11 +319,12 @@ pub const SNativeZlib = struct {
     write_result: ?[*]u32 = null,
     write_callback: JSC.Strong = .{},
     onerror_value: JSC.Strong = .{},
-    poll_ref: bun.Async.KeepAlive = .{},
+    poll_ref: CountedKeepAlive = .{},
     this_value: JSC.Strong = .{},
     write_in_progress: bool = false,
     pending_close: bool = false,
     closed: bool = false,
+    task: JSC.WorkPoolTask = .{ .callback = undefined },
 
     pub fn constructor(globalThis: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) ?*@This() {
         const arguments = callframe.argumentsUndef(4).ptr;
@@ -392,6 +406,7 @@ pub const SNativeZlib = struct {
     pub fn deinit(this: *@This()) void {
         this.write_callback.deinit();
         this.onerror_value.deinit();
+        this.poll_ref.deinit();
         this.destroy();
     }
 };
@@ -662,11 +677,14 @@ pub const SNativeBrotli = struct {
     write_result: ?[*]u32 = null,
     write_callback: JSC.Strong = .{},
     onerror_value: JSC.Strong = .{},
-    poll_ref: bun.Async.KeepAlive = .{},
+    poll_ref: CountedKeepAlive = .{},
     this_value: JSC.Strong = .{},
     write_in_progress: bool = false,
     pending_close: bool = false,
     closed: bool = false,
+    task: JSC.WorkPoolTask = .{
+        .callback = undefined,
+    },
 
     pub fn constructor(globalThis: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) ?*@This() {
         const arguments = callframe.argumentsUndef(1).ptr;
@@ -752,6 +770,7 @@ pub const SNativeBrotli = struct {
     pub fn deinit(this: *@This()) void {
         this.write_callback.deinit();
         this.onerror_value.deinit();
+        this.poll_ref.deinit();
         this.destroy();
     }
 };

diff --git a/src/bun.js/web_worker.zig b/src/bun.js/web_worker.zig
@@ -459,7 +459,9 @@ pub const WebWorker = struct {
         var exit_code: i32 = 0;
         var globalObject: ?*JSC.JSGlobalObject = null;
         var vm_to_deinit: ?*JSC.VirtualMachine = null;
+        var loop: ?*bun.uws.Loop = null;
         if (this.vm) |vm| {
+            loop = vm.uwsLoop();
             this.vm = null;
             vm.is_shutting_down = true;
             vm.onExit();
@@ -470,6 +472,9 @@ pub const WebWorker = struct {
         var arena = this.arena;
 
         WebWorker__dispatchExit(globalObject, cpp_worker, exit_code);
+        if (loop) |loop_| {
+            loop_.internal_loop_data.jsc_vm = null;
+        }
         bun.uws.onThreadExit();
         this.deinit();
 

diff --git a/src/bun.zig b/src/bun.zig
@@ -3146,8 +3146,8 @@ pub fn NewRefCounted(comptime T: type, comptime deinit_fn: ?fn (self: *T) void)
             const ptr = bun.new(T, t);
 
             if (Environment.enable_logs) {
-                if (ptr.ref_count != 1) {
-                    Output.panic("Expected ref_count to be 1, got {d}", .{ptr.ref_count});
+                if (ptr.ref_count == 0) {
+                    Output.panic("Expected ref_count to be > 0, got {d}", .{ptr.ref_count});
                 }
             }
 

diff --git a/src/deps/uws.zig b/src/deps/uws.zig
@@ -55,6 +55,7 @@ pub const InternalLoopData = extern struct {
     parent_ptr: ?*anyopaque,
     parent_tag: c_char,
     iteration_nr: usize,
+    jsc_vm: ?*JSC.VM,
 
     pub fn recvSlice(this: *InternalLoopData) []u8 {
         return this.recv_buf[0..LIBUS_RECV_BUFFER_LENGTH];

diff --git a/test/bundler/expectBundled.ts b/test/bundler/expectBundled.ts
@@ -556,6 +556,12 @@ function expectBundled(
     let root = path.join(
       tempDirectory,
       id
+        .replaceAll("\\", "/")
+        .replaceAll(":", "-")
+        .replaceAll(" ", "-")
+        .replaceAll("\r\n", "-")
+        .replaceAll("\n", "-")
+        .replaceAll(".", "-")
         .split("/")
         .map(a => filenamify(a))
         .join("/"),

diff --git a/test/integration/next-pages/test/dev-server.test.ts b/test/integration/next-pages/test/dev-server.test.ts
@@ -79,7 +79,7 @@ async function getDevServerURL() {
   readStream()
     .catch(e => reject(e))
     .finally(() => {
-      dev_server.unref?.();
+      dev_server?.unref?.();
     });
   await promise;
   return baseUrl;

diff --git a/test/js/bun/css/doesnt_crash.test.ts b/test/js/bun/css/doesnt_crash.test.ts
@@ -15,15 +15,18 @@ describe("doesnt_crash", async () => {
   files = readdirSync(files_dir).map(file => path.join(files_dir, file));
   console.log("Tempdir", temp_dir);
 
-  files.map(file => {
-    const outfile1 = path.join(temp_dir, file).replaceAll("\\", "/");
-    const outfile2 = path.join(temp_dir, "lmao1-" + file).replaceAll("\\", "/");
-    const outfile3 = path.join(temp_dir, "lmao2-" + file).replaceAll("\\", "/");
-    const outfile4 = path.join(temp_dir, "lmao3-" + file).replaceAll("\\", "/");
+  files.map(absolute => {
+    absolute = absolute.replaceAll("\\", "/");
+    const file = path.basename(absolute);
+    const outfile1 = path.join(temp_dir, "file-1" + file).replaceAll("\\", "/");
+    const outfile2 = path.join(temp_dir, "file-2" + file).replaceAll("\\", "/");
+    const outfile3 = path.join(temp_dir, "file-3" + file).replaceAll("\\", "/");
+    const outfile4 = path.join(temp_dir, "file-4" + file).replaceAll("\\", "/");
+
     test(file, async () => {
       {
         const { stdout, stderr, exitCode } =
-          await Bun.$`${bunExe()} build --experimental-css ${file} --outfile=${outfile1}`.quiet().env(bunEnv);
+          await Bun.$`${bunExe()} build --experimental-css ${absolute} --outfile=${outfile1}`.quiet().env(bunEnv);
         expect(exitCode).toBe(0);
         expect(stdout.toString()).not.toContain("error");
         expect(stderr.toString()).toBeEmpty();
@@ -39,7 +42,9 @@ describe("doesnt_crash", async () => {
     test(`(minify) ${file}`, async () => {
       {
         const { stdout, stderr, exitCode } =
-          await Bun.$`${bunExe()} build --experimental-css ${file} --minify --outfile=${outfile3}`.quiet().env(bunEnv);
+          await Bun.$`${bunExe()} build --experimental-css ${absolute} --minify --outfile=${outfile3}`
+            .quiet()
+            .env(bunEnv);
         expect(exitCode).toBe(0);
         expect(stdout.toString()).not.toContain("error");
         expect(stderr.toString()).toBeEmpty();