Merge pull request #393 from swiftwasm/yt/jsclosure-fix

Fix wrong deallocation management for JSClosure with FinalizationRegistry

Author: kateinoigakukun

Makefile

@@ -13,7 +13,7 @@ unittest:
 	    -Xlinker --global-base=524288 \
 	    -Xlinker -z \
 	    -Xlinker stack-size=524288 \
-	    js test --prelude ./Tests/prelude.mjs
+	    js test --prelude ./Tests/prelude.mjs -Xnode --expose-gc
 
 .PHONY: regenerate_swiftpm_resources
 regenerate_swiftpm_resources:

Plugins/PackageToJS/Templates/runtime.mjs

@@ -443,7 +443,7 @@ class SwiftRuntime {
         });
         const alreadyReleased = this.exports.swjs_call_host_function(host_func_id, argv, argc, callback_func_ref);
         if (alreadyReleased) {
-            throw new Error(`The JSClosure has been already released by Swift side. The closure is created at ${file}:${line}`);
+            throw new Error(`The JSClosure has been already released by Swift side. The closure is created at ${file}:${line} @${host_func_id}`);
         }
         this.exports.swjs_cleanup_host_function_call(argv);
         return output;
@@ -635,7 +635,13 @@ class SwiftRuntime {
                 const fileString = this.memory.getObject(file);
                 const func = (...args) => this.callHostFunction(host_func_id, line, fileString, args);
                 const func_ref = this.memory.retain(func);
-                (_a = this.closureDeallocator) === null || _a === void 0 ? void 0 : _a.track(func, func_ref);
+                (_a = this.closureDeallocator) === null || _a === void 0 ? void 0 : _a.track(func, host_func_id);
+                return func_ref;
+            },
+            swjs_create_oneshot_function: (host_func_id, line, file) => {
+                const fileString = this.memory.getObject(file);
+                const func = (...args) => this.callHostFunction(host_func_id, line, fileString, args);
+                const func_ref = this.memory.retain(func);
                 return func_ref;
             },
             swjs_create_typed_array: (constructor_ref, elementsPtr, length) => {

Runtime/src/index.ts

@@ -224,7 +224,7 @@ export class SwiftRuntime {
         );
         if (alreadyReleased) {
             throw new Error(
-                `The JSClosure has been already released by Swift side. The closure is created at ${file}:${line}`
+                `The JSClosure has been already released by Swift side. The closure is created at ${file}:${line} @${host_func_id}`
             );
         }
         this.exports.swjs_cleanup_host_function_call(argv);
@@ -553,7 +553,19 @@ export class SwiftRuntime {
                 const func = (...args: any[]) =>
                     this.callHostFunction(host_func_id, line, fileString, args);
                 const func_ref = this.memory.retain(func);
-                this.closureDeallocator?.track(func, func_ref);
+                this.closureDeallocator?.track(func, host_func_id);
+                return func_ref;
+            },
+
+            swjs_create_oneshot_function: (
+                host_func_id: number,
+                line: number,
+                file: ref
+            ) => {
+                const fileString = this.memory.getObject(file) as string;
+                const func = (...args: any[]) =>
+                    this.callHostFunction(host_func_id, line, fileString, args);
+                const func_ref = this.memory.retain(func);
                 return func_ref;
             },
 

Sources/JavaScriptKit/FundamentalObjects/JSClosure.swift

@@ -25,7 +25,7 @@ public class JSOneshotClosure: JSObject, JSClosureProtocol {
         // 2. Create a new JavaScript function which calls the given Swift function.
         hostFuncRef = JavaScriptHostFuncRef(bitPattern: ObjectIdentifier(self))
         _id = withExtendedLifetime(JSString(file)) { file in
-            swjs_create_function(hostFuncRef, line, file.asInternalJSRef())
+            swjs_create_oneshot_function(hostFuncRef, line, file.asInternalJSRef())
         }
 
         // 3. Retain the given body in static storage by `funcRef`.

Sources/_CJavaScriptKit/include/_CJavaScriptKit.h

@@ -278,6 +278,16 @@ IMPORT_JS_FUNCTION(swjs_create_function, JavaScriptObjectRef, (const JavaScriptH
                                                                unsigned int line,
                                                                JavaScriptObjectRef file))
 
+/// Creates a oneshot JavaScript thunk function that calls Swift side closure.
+///
+/// @param host_func_id The target Swift side function called by the created thunk function.
+/// @param line The line where the function is created. Will be used for diagnostics
+/// @param file The file name where the function is created. Will be used for diagnostics
+/// @returns A reference to the newly-created JavaScript thunk function
+IMPORT_JS_FUNCTION(swjs_create_oneshot_function, JavaScriptObjectRef, (const JavaScriptHostFuncRef host_func_id,
+                                                                       unsigned int line,
+                                                                       JavaScriptObjectRef file))
+
 /// Instantiates a new `TypedArray` object with given elements
 /// This is used to provide an efficient way to create `TypedArray`.
 ///

Tests/JavaScriptKitTests/JSClosureTests.swift

@@ -1,4 +1,4 @@
-import JavaScriptKit
+@_spi(JSObject_id) import JavaScriptKit
 import XCTest
 
 class JSClosureTests: XCTestCase {
@@ -85,4 +85,69 @@ class JSClosureTests: XCTestCase {
         hostFunc2.release()
         #endif
     }
+
+    func testRegressionTestForMisDeallocation() async throws {
+        // Use Node.js's `--expose-gc` flag to enable manual garbage collection.
+        guard let gc = JSObject.global.gc.function else {
+            throw XCTSkip("Missing --expose-gc flag")
+        }
+
+        // Step 1: Create many JSClosure instances
+        let obj = JSObject()
+        var closurePointers: Set<UInt32> = []
+        let numberOfSourceClosures = 10_000
+
+        do {
+            var closures: [JSClosure] = []
+            for i in 0..<numberOfSourceClosures {
+                let closure = JSClosure { _ in .undefined }
+                obj["c\(i)"] = closure.jsValue
+                closures.append(closure)
+                // Store
+                closurePointers.insert(UInt32(UInt(bitPattern: Unmanaged.passUnretained(closure).toOpaque())))
+
+                // To avoid all JSClosures having a common address diffs, randomly allocate a new object.
+                if Bool.random() {
+                    _ = JSObject()
+                }
+            }
+        }
+
+        // Step 2: Create many JSObject to make JSObject.id close to Swift heap object address
+        let minClosurePointer = closurePointers.min() ?? 0
+        let maxClosurePointer = closurePointers.max() ?? 0
+        while true {
+            let obj = JSObject()
+            if minClosurePointer == obj.id {
+                break
+            }
+        }
+
+        // Step 3: Create JSClosure instances and find the one with JSClosure.id == &closurePointers[x]
+        do {
+            while true {
+                let c = JSClosure { _ in .undefined }
+                if closurePointers.contains(c.id) || c.id > maxClosurePointer {
+                    break
+                }
+                // To avoid all JSClosures having a common JSObject.id diffs, randomly allocate a new JS object.
+                if Bool.random() {
+                    _ = JSObject()
+                }
+            }
+        }
+
+        // Step 4: Trigger garbage collection to call the finalizer of the conflicting JSClosure instance
+        for _ in 0..<100 {
+            gc()
+            // Tick the event loop to allow the garbage collector to run finalizers
+            // registered by FinalizationRegistry.
+            try await Task.sleep(for: .milliseconds(0))
+        }
+
+        // Step 5: Verify that the JSClosure instances are still alive and can be called
+        for i in 0..<numberOfSourceClosures {
+            _ = obj["c\(i)"].function!()
+        }
+    }
 }