refactor(js): unify env lifecycle refcounting

src/js.zig

@@ -62,6 +62,7 @@ test {
     // which would force-link free functions (throwError) against C symbols
     // unavailable in the native test runner.
     _ = @import("js/context.zig");
+    _ = @import("js/lifecycle.zig");
     _ = @import("js/io.zig");
     _ = @import("js/number.zig");
     _ = @import("js/string.zig");

src/js/export_module.zig

@@ -2,6 +2,7 @@ const std = @import("std");
 const napi = @import("../napi.zig");
 const context = @import("context.zig");
 const io_context = @import("io.zig");
+const lifecycle = @import("lifecycle.zig");
 const wrap_function = @import("wrap_function.zig");
 const wrap_class = @import("wrap_class.zig");
 const class_meta = @import("class_meta.zig");
@@ -32,9 +33,11 @@ const class_runtime = @import("class_runtime.zig");
 ///   have been processed and *before* the module's `init` hook (if present).
 ///   `exports` is the JavaScript object that will hold the module's exports.
 ///
-/// The DSL internally manages an atomic refcount for module instances across
+/// The DSL internally manages a refcount for module instances across
 /// different N-API environments, and uses the same env lifecycle to retain a
-/// shared `js.io()` handle for the addon.
+/// shared `js.io()` handle for the addon. Lifecycle hooks are serialized
+/// across environments; if export registration fails after a successful
+/// `init`, `cleanup` runs before the error propagates.
 ///
 /// Usage Examples:
 /// ```zig
@@ -62,10 +65,20 @@ pub fn exportModule(comptime Module: type, comptime options: anytype) void {
     const has_init = @hasField(@TypeOf(options), "init");
     const has_cleanup = @hasField(@TypeOf(options), "cleanup");
     const has_register = @hasField(@TypeOf(options), "register");
-    const has_lifecycle = has_init or has_cleanup;
 
     const State = struct {
-        var env_refcount: std.atomic.Value(u32) = std.atomic.Value(u32).init(0);
+        const Lifecycle = lifecycle.SharedResource(void, .{
+            .init = lifecycleInit,
+            .cleanup = lifecycleCleanup,
+        });
+
+        fn lifecycleInit(_: *void, prev_refcount: u32) !void {
+            if (has_init) try options.init(prev_refcount);
+        }
+
+        fn lifecycleCleanup(_: *void, new_refcount: u32) void {
+            if (has_cleanup) options.cleanup(new_refcount);
+        }
 
         // addEnvCleanupHook requires a non-null *Data pointer.
         const CleanupData = struct {
@@ -74,13 +87,7 @@ pub fn exportModule(comptime Module: type, comptime options: anytype) void {
         var cleanup_data: CleanupData = .{};
 
         fn cleanupHook(_: *CleanupData) void {
-            if (has_lifecycle) {
-                const prev = env_refcount.fetchSub(1, .acq_rel);
-                const new_refcount = prev - 1;
-                if (has_cleanup) {
-                    options.cleanup(new_refcount);
-                }
-            }
+            Lifecycle.release();
             io_context.release();
         }
     };
@@ -96,17 +103,12 @@ pub fn exportModule(comptime Module: type, comptime options: anytype) void {
                 io_context.release();
             };
 
-            var prev_refcount: u32 = 0;
-            if (has_lifecycle) {
-                prev_refcount = State.env_refcount.fetchAdd(1, .monotonic);
-                errdefer if (!cleanup_hook_registered) {
-                    _ = State.env_refcount.fetchSub(1, .acq_rel);
-                };
-
-                if (has_init) {
-                    try options.init(prev_refcount);
-                }
-            }
+            // On init failure the refcount is untouched; on later registration
+            // failure the rollback release runs `cleanup` with the right count.
+            _ = try State.Lifecycle.retain();
+            errdefer if (!cleanup_hook_registered) {
+                State.Lifecycle.release();
+            };
 
             _ = try registerDecls(Module, env, module, 0);
 

src/js/io.zig

@@ -1,45 +1,37 @@
 const std = @import("std");
+const lifecycle = @import("lifecycle.zig");
 
 const gpa: std.mem.Allocator = std.heap.page_allocator;
 
-const State = struct {
-    var mutex: std.Io.Mutex = .init;
-    var instance: std.Io.Threaded = undefined;
-    var initialized: bool = false;
-    var refcount: u32 = 0;
+const hooks = struct {
+    fn init(instance: *std.Io.Threaded, prev_refcount: u32) !void {
+        if (prev_refcount == 0) instance.* = std.Io.Threaded.init(gpa, .{});
+    }
+    fn cleanup(instance: *std.Io.Threaded, new_refcount: u32) void {
+        if (new_refcount == 0) instance.deinit();
+    }
 };
 
+const SharedIo = lifecycle.SharedResource(std.Io.Threaded, .{
+    .init = hooks.init,
+    .cleanup = hooks.cleanup,
+});
+
 /// Retains the shared DSL `std.Io` instance for an active N-API environment.
 ///
 /// Called internally from `js.exportModule(...)` on module registration.
 /// The underlying `std.Io.Threaded` is initialized lazily on the first retain
 /// and torn down after the last matching `release()`.
 pub fn retain() void {
-    std.Io.Threaded.mutexLock(&State.mutex);
-    defer std.Io.Threaded.mutexUnlock(&State.mutex);
-
-    if (State.refcount == 0) {
-        State.instance = std.Io.Threaded.init(gpa, .{});
-        State.initialized = true;
-    }
-    State.refcount += 1;
+    _ = SharedIo.retain() catch unreachable;
 }
 
 /// Releases one active N-API environment's hold on the shared DSL `std.Io`.
 ///
 /// Called internally from the env cleanup hook installed by
 /// `js.exportModule(...)`.
 pub fn release() void {
-    std.Io.Threaded.mutexLock(&State.mutex);
-    defer std.Io.Threaded.mutexUnlock(&State.mutex);
-
-    std.debug.assert(State.refcount > 0);
-    State.refcount -= 1;
-
-    if (State.refcount == 0 and State.initialized) {
-        State.instance.deinit();
-        State.initialized = false;
-    }
+    SharedIo.release();
 }
 
 /// Returns the shared `std.Io` handle managed by the JS DSL.
@@ -52,32 +44,24 @@ pub fn release() void {
 /// SAFETY: `js.io()` panics if called before the addon has been registered in a
 /// JS environment or after the last environment has been cleaned up.
 pub fn io() std.Io {
-    std.Io.Threaded.mutexLock(&State.mutex);
-    defer std.Io.Threaded.mutexUnlock(&State.mutex);
-
-    if (!State.initialized) {
+    const instance = SharedIo.get() orelse
         @panic("js.io() called before DSL module registration or after env cleanup");
-    }
-    return State.instance.io();
+    return instance.io();
 }
 
 test "shared io retains and releases across env registrations" {
-    try std.testing.expect(!State.initialized);
-    try std.testing.expectEqual(@as(u32, 0), State.refcount);
+    try std.testing.expect(SharedIo.get() == null);
 
     retain();
-    try std.testing.expect(State.initialized);
-    try std.testing.expectEqual(@as(u32, 1), State.refcount);
-    _ = io();
+    const a = io();
 
     retain();
-    try std.testing.expectEqual(@as(u32, 2), State.refcount);
+    const b = io();
+    try std.testing.expectEqual(a.userdata, b.userdata);
 
     release();
-    try std.testing.expect(State.initialized);
-    try std.testing.expectEqual(@as(u32, 1), State.refcount);
+    _ = io();
 
     release();
-    try std.testing.expect(!State.initialized);
-    try std.testing.expectEqual(@as(u32, 0), State.refcount);
+    try std.testing.expect(SharedIo.get() == null);
 }

src/js/lifecycle.zig

@@ -0,0 +1,172 @@
+const std = @import("std");
+
+/// Refcounted shared state with lifecycle hooks serialized under a mutex.
+///
+/// Optional `hooks`:
+///
+/// - `.init = fn (instance: *T, prev_refcount: u32) !void`: called under the
+///   lock before the refcount is incremented (`prev_refcount` is 0 for the
+///   first holder). On error the refcount is left untouched.
+/// - `.cleanup = fn (instance: *T, new_refcount: u32) void`: called under the
+///   lock after the refcount is decremented (`new_refcount` is 0 for the last
+///   holder).
+///
+/// Each instantiation owns its own refcount, mutex, and `T` instance.
+pub fn SharedResource(comptime T: type, comptime hooks: anytype) type {
+    const has_init = @hasField(@TypeOf(hooks), "init");
+    const has_cleanup = @hasField(@TypeOf(hooks), "cleanup");
+
+    return struct {
+        var mutex: std.Io.Mutex = .init;
+        var refcount: u32 = 0;
+        var instance: T = undefined;
+
+        /// Returns the refcount before this retain (0 for the first holder).
+        pub fn retain() !u32 {
+            std.Io.Threaded.mutexLock(&mutex);
+            defer std.Io.Threaded.mutexUnlock(&mutex);
+
+            const prev = refcount;
+            if (has_init) try hooks.init(&instance, prev);
+            refcount += 1;
+            return prev;
+        }
+
+        pub fn release() void {
+            std.Io.Threaded.mutexLock(&mutex);
+            defer std.Io.Threaded.mutexUnlock(&mutex);
+
+            std.debug.assert(refcount > 0);
+            refcount -= 1;
+            if (has_cleanup) hooks.cleanup(&instance, refcount);
+        }
+
+        /// Returns the shared instance, or null when no holder is active.
+        ///
+        /// The pointer is only valid while the caller holds a retain; without
+        /// one, a concurrent last `release()` may tear the instance down.
+        pub fn get() ?*T {
+            std.Io.Threaded.mutexLock(&mutex);
+            defer std.Io.Threaded.mutexUnlock(&mutex);
+
+            if (refcount == 0) return null;
+            return &instance;
+        }
+    };
+}
+
+test "SharedResource runs init before increment and cleanup after decrement" {
+    const Hooks = struct {
+        var inits: [2]u32 = undefined;
+        var init_count: usize = 0;
+        var cleanups: [2]u32 = undefined;
+        var cleanup_count: usize = 0;
+
+        fn init(_: *void, prev_refcount: u32) !void {
+            inits[init_count] = prev_refcount;
+            init_count += 1;
+        }
+        fn cleanup(_: *void, new_refcount: u32) void {
+            cleanups[cleanup_count] = new_refcount;
+            cleanup_count += 1;
+        }
+    };
+    const S = SharedResource(void, .{ .init = Hooks.init, .cleanup = Hooks.cleanup });
+
+    try std.testing.expectEqual(@as(u32, 0), try S.retain());
+    try std.testing.expectEqual(@as(u32, 1), try S.retain());
+    S.release();
+    S.release();
+
+    try std.testing.expectEqualSlices(u32, &.{ 0, 1 }, Hooks.inits[0..2]);
+    try std.testing.expectEqualSlices(u32, &.{ 1, 0 }, Hooks.cleanups[0..2]);
+}
+
+test "SharedResource leaves refcount untouched and skips cleanup when init fails" {
+    const Hooks = struct {
+        var fail: bool = true;
+        var cleanup_count: u32 = 0;
+
+        fn init(_: *void, _: u32) !void {
+            if (fail) return error.InitFailed;
+        }
+        fn cleanup(_: *void, _: u32) void {
+            cleanup_count += 1;
+        }
+    };
+    const S = SharedResource(void, .{ .init = Hooks.init, .cleanup = Hooks.cleanup });
+
+    try std.testing.expectError(error.InitFailed, S.retain());
+    try std.testing.expect(S.get() == null);
+    try std.testing.expectEqual(@as(u32, 0), Hooks.cleanup_count);
+
+    Hooks.fail = false;
+    try std.testing.expectEqual(@as(u32, 0), try S.retain());
+    S.release();
+    try std.testing.expectEqual(@as(u32, 1), Hooks.cleanup_count);
+}
+
+test "SharedResource exposes the instance only while retained" {
+    const Hooks = struct {
+        fn init(instance: *u32, prev_refcount: u32) !void {
+            if (prev_refcount == 0) instance.* = 42;
+        }
+    };
+    const S = SharedResource(u32, .{ .init = Hooks.init });
+
+    try std.testing.expect(S.get() == null);
+    _ = try S.retain();
+    try std.testing.expectEqual(@as(u32, 42), S.get().?.*);
+    S.release();
+    try std.testing.expect(S.get() == null);
+}
+
+test "SharedResource serializes hooks across threads" {
+    const thread_count = 4;
+    const iterations = 100;
+
+    const Hooks = struct {
+        var busy: std.atomic.Value(bool) = std.atomic.Value(bool).init(false);
+        var overlapped: std.atomic.Value(bool) = std.atomic.Value(bool).init(false);
+        // Non-atomic on purpose: serialized hooks must not race on these.
+        var init_count: u32 = 0;
+        var cleanup_count: u32 = 0;
+
+        fn enter() void {
+            if (busy.swap(true, .acquire)) overlapped.store(true, .seq_cst);
+        }
+        fn exit() void {
+            busy.store(false, .release);
+        }
+        fn init(_: *void, _: u32) !void {
+            enter();
+            defer exit();
+            init_count += 1;
+        }
+        fn cleanup(_: *void, _: u32) void {
+            enter();
+            defer exit();
+            cleanup_count += 1;
+        }
+    };
+    const S = SharedResource(void, .{ .init = Hooks.init, .cleanup = Hooks.cleanup });
+
+    const worker = struct {
+        fn run() void {
+            for (0..iterations) |_| {
+                _ = S.retain() catch unreachable;
+                std.debug.assert(S.get() != null);
+                S.release();
+            }
+        }
+    };
+
+    var threads: [thread_count]std.Thread = undefined;
+    for (&threads) |*t| t.* = try std.Thread.spawn(.{}, worker.run, .{});
+    for (threads) |t| t.join();
+
+    try std.testing.expect(!Hooks.overlapped.load(.seq_cst));
+    try std.testing.expectEqual(@as(u32, thread_count * iterations), Hooks.init_count);
+    try std.testing.expectEqual(@as(u32, thread_count * iterations), Hooks.cleanup_count);
+    try std.testing.expect(S.get() == null);
+}