Fix JSClosure leak (#240)

* Fix retain cycle in JSClosure by weak reference to the closure holder

```swift
let c1 = JSClosure { _ in .undefined }
consume c1
```

did not release the `JSClosure` itself and it leaked the underlying
JavaScript closure too because `JSClosure` -> JS closure thunk ->
Closure registry entry -> `JSClosure` reference cycle was not broken
when using FinalizationRegistry. (Without FR, it was broken by manual
`release` call.)

Note that weakening the reference does not violates the contract that
function reference should be unique because holding a weak reference does
deinit but not deallocate the object, so ObjectIdentifier is not reused
until the weak reference in the registry is removed.

* Fix the test suite for violation of the closure lifetime contract

The test suite was not properly releasing the closures but they have not
been revealed as a problem because those closures were leaked
conservatively.

* Report where the closure is created when it violates the lifetime contract

This additional information will help developers to find the root cause

* Flatten two-level object reference in ClosureEntry

Author: kateinoigakukun

IntegrationTests/TestSuites/Sources/ConcurrencyTests/main.swift

@@ -123,12 +123,16 @@ func entrypoint() async throws {
             try expectEqual(result, .number(3))
         }
         try expectGTE(diff, 200)
+#if JAVASCRIPTKIT_WITHOUT_WEAKREFS
+        delayObject.closure = nil
+        delayClosure.release()
+#endif
     }
 
     try await asyncTest("Async JSPromise: then") {
         let promise = JSPromise { resolve in
             _ = JSObject.global.setTimeout!(
-                JSClosure { _  in
+                JSOneshotClosure { _  in
                     resolve(.success(JSValue.number(3)))
                     return .undefined
                 }.jsValue,
@@ -149,7 +153,7 @@ func entrypoint() async throws {
     try await asyncTest("Async JSPromise: then(success:failure:)") {
         let promise = JSPromise { resolve in
             _ = JSObject.global.setTimeout!(
-                JSClosure { _ in
+                JSOneshotClosure { _ in
                     resolve(.failure(JSError(message: "test").jsValue))
                     return .undefined
                 }.jsValue,
@@ -168,7 +172,7 @@ func entrypoint() async throws {
     try await asyncTest("Async JSPromise: catch") {
         let promise = JSPromise { resolve in
             _ = JSObject.global.setTimeout!(
-                JSClosure { _ in
+                JSOneshotClosure { _ in
                     resolve(.failure(JSError(message: "test").jsValue))
                     return .undefined
                 }.jsValue,

IntegrationTests/TestSuites/Sources/PrimaryTests/UnitTestUtils.swift

@@ -111,7 +111,7 @@ func expectThrow<T>(_ body: @autoclosure () throws -> T, file: StaticString = #f
 }
 
 func wrapUnsafeThrowableFunction(_ body: @escaping () -> Void, file: StaticString = #file, line: UInt = #line, column: UInt = #column) throws -> Error {
-    JSObject.global.callThrowingClosure.function!(JSClosure { _ in 
+    JSObject.global.callThrowingClosure.function!(JSOneshotClosure { _ in
             body() 
             return .undefined
     })

IntegrationTests/TestSuites/Sources/PrimaryTests/main.swift

@@ -254,12 +254,18 @@ try test("Closure Lifetime") {
     do {
         let c1 = JSClosure { _ in .number(4) }
         try expectEqual(c1(), .number(4))
+#if JAVASCRIPTKIT_WITHOUT_WEAKREFS
+        c1.release()
+#endif
     }
 
     do {
         let c1 = JSClosure { _ in fatalError("Crash while closure evaluation") }
         let error = try expectThrow(try evalClosure.throws(c1)) as! JSValue
         try expectEqual(error.description, "RuntimeError: unreachable")
+#if JAVASCRIPTKIT_WITHOUT_WEAKREFS
+        c1.release()
+#endif
     }
 }
 
@@ -866,16 +872,24 @@ try test("Symbols") {
     // }.prop
     // Object.defineProperty(hasInstanceClass, Symbol.hasInstance, { value: () => true })
     let hasInstanceObject = JSObject.global.Object.function!.new()
-    hasInstanceObject.prop = JSClosure { _ in .undefined }.jsValue
+    let hasInstanceObjectClosure = JSClosure { _ in .undefined }
+    hasInstanceObject.prop = hasInstanceObjectClosure.jsValue
     let hasInstanceClass = hasInstanceObject.prop.function!
     let propertyDescriptor = JSObject.global.Object.function!.new()
-    propertyDescriptor.value = JSClosure { _ in .boolean(true) }.jsValue
+    let propertyDescriptorClosure = JSClosure { _ in .boolean(true) }
+    propertyDescriptor.value = propertyDescriptorClosure.jsValue
     _ = JSObject.global.Object.function!.defineProperty!(
         hasInstanceClass, JSSymbol.hasInstance,
         propertyDescriptor
     )
     try expectEqual(hasInstanceClass[JSSymbol.hasInstance].function!().boolean, true)
     try expectEqual(JSObject.global.Object.isInstanceOf(hasInstanceClass), true)
+#if JAVASCRIPTKIT_WITHOUT_WEAKREFS
+    hasInstanceObject.prop = .undefined
+    propertyDescriptor.value = .undefined
+    hasInstanceObjectClosure.release()
+    propertyDescriptorClosure.release()
+#endif
 }
 
 struct AnimalStruct: Decodable {

Sources/JavaScriptKit/FundamentalObjects/JSClosure.swift

@@ -19,17 +19,15 @@ public class JSOneshotClosure: JSObject, JSClosureProtocol {
         // 1. Fill `id` as zero at first to access `self` to get `ObjectIdentifier`.
         super.init(id: 0)
 
-        // 2. Create a new JavaScript function which calls the given Swift function.
-        hostFuncRef = JavaScriptHostFuncRef(bitPattern: ObjectIdentifier(self))
-        id = withExtendedLifetime(JSString(file)) { file in
-          _create_function(hostFuncRef, line, file.asInternalJSRef())
-        }
-
-        // 3. Retain the given body in static storage by `funcRef`.
-        JSClosure.sharedClosures[hostFuncRef] = (self, {
+        // 2. Retain the given body in static storage
+        // Leak the self object globally and release once it's called
+        hostFuncRef = JSClosure.sharedClosures.register(ObjectIdentifier(self), object: .strong(self), body: {
             defer { self.release() }
             return body($0)
         })
+        id = withExtendedLifetime(JSString(file)) { file in
+          _create_function(hostFuncRef, line, file.asInternalJSRef())
+        }
     }
 
     #if compiler(>=5.5)
@@ -42,7 +40,7 @@ public class JSOneshotClosure: JSObject, JSClosureProtocol {
     /// Release this function resource.
     /// After calling `release`, calling this function from JavaScript will fail.
     public func release() {
-        JSClosure.sharedClosures[hostFuncRef] = nil
+        JSClosure.sharedClosures.unregister(hostFuncRef)
     }
 }
 
@@ -62,13 +60,13 @@ public class JSOneshotClosure: JSObject, JSClosureProtocol {
 /// ```
 ///
 public class JSClosure: JSFunction, JSClosureProtocol {
-
-    // Note: Retain the closure object itself also to avoid funcRef conflicts
-    fileprivate static var sharedClosures: [JavaScriptHostFuncRef: (object: JSObject, body: ([JSValue]) -> JSValue)] = [:]
+    fileprivate static var sharedClosures = SharedJSClosureRegistry()
 
     private var hostFuncRef: JavaScriptHostFuncRef = 0
 
     #if JAVASCRIPTKIT_WITHOUT_WEAKREFS
+    private let file: String
+    private let line: UInt32
     private var isReleased: Bool = false
     #endif
 
@@ -82,30 +80,35 @@ public class JSClosure: JSFunction, JSClosureProtocol {
     }
 
     public init(_ body: @escaping ([JSValue]) -> JSValue, file: String = #fileID, line: UInt32 = #line) {
+        #if JAVASCRIPTKIT_WITHOUT_WEAKREFS
+        self.file = file
+        self.line = line
+        #endif
         // 1. Fill `id` as zero at first to access `self` to get `ObjectIdentifier`.
         super.init(id: 0)
 
-        // 2. Create a new JavaScript function which calls the given Swift function.
-        hostFuncRef = JavaScriptHostFuncRef(bitPattern: ObjectIdentifier(self))
+        // 2. Retain the given body in static storage
+        hostFuncRef = Self.sharedClosures.register(
+            ObjectIdentifier(self), object: .weak(self), body: body
+        )
+
         id = withExtendedLifetime(JSString(file)) { file in
           _create_function(hostFuncRef, line, file.asInternalJSRef())
         }
 
-        // 3. Retain the given body in static storage by `funcRef`.
-        Self.sharedClosures[hostFuncRef] = (self, body)
     }
 
     #if compiler(>=5.5)
     @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
-    public static func async(_ body: @escaping ([JSValue]) async throws -> JSValue) -> JSClosure {
-        JSClosure(makeAsyncClosure(body))
+    public static func async(_ body: @escaping ([JSValue]) async throws -> JSValue, file: String = #fileID, line: UInt32 = #line) -> JSClosure {
+        JSClosure(makeAsyncClosure(body), file: file, line: line)
     }
     #endif
 
     #if JAVASCRIPTKIT_WITHOUT_WEAKREFS
     deinit {
         guard isReleased else {
-            fatalError("release() must be called on JSClosure objects manually before they are deallocated")
+            fatalError("release() must be called on JSClosure object (\(file):\(line)) manually before they are deallocated")
         }
     }
     #endif
@@ -133,6 +136,60 @@ private func makeAsyncClosure(_ body: @escaping ([JSValue]) async throws -> JSVa
 }
 #endif
 
+/// Registry for Swift closures that are referenced from JavaScript.
+private struct SharedJSClosureRegistry {
+    struct ClosureEntry {
+        // Note: Retain the closure object itself also to avoid funcRef conflicts.
+        var object: AnyObjectReference
+        var body: ([JSValue]) -> JSValue
+
+        init(object: AnyObjectReference, body: @escaping ([JSValue]) -> JSValue) {
+            self.object = object
+            self.body = body
+        }
+    }
+    enum AnyObjectReference {
+        case strong(AnyObject)
+        case weak(WeakObject)
+
+        static func `weak`(_ object: AnyObject) -> AnyObjectReference {
+            .weak(SharedJSClosureRegistry.WeakObject(underlying: object))
+        }
+    }
+    struct WeakObject {
+        weak var underlying: AnyObject?
+        init(underlying: AnyObject) {
+            self.underlying = underlying
+        }
+    }
+    private var closures: [JavaScriptHostFuncRef: ClosureEntry] = [:]
+
+    /// Register a Swift closure to be called from JavaScript.
+    /// - Parameters:
+    ///   - hint: A hint to identify the closure.
+    ///   - object: The object should be retained until the closure is released from JavaScript.
+    ///   - body: The closure to be called from JavaScript.
+    /// - Returns: An unique identifier for the registered closure.
+    mutating func register(
+        _ hint: ObjectIdentifier,
+        object: AnyObjectReference, body: @escaping ([JSValue]) -> JSValue
+    ) -> JavaScriptHostFuncRef {
+        let ref = JavaScriptHostFuncRef(bitPattern: hint)
+        closures[ref] = ClosureEntry(object: object, body: body)
+        return ref
+    }
+
+    /// Unregister a Swift closure from the registry.
+    mutating func unregister(_ ref: JavaScriptHostFuncRef) {
+        closures[ref] = nil
+    }
+
+    /// Get the Swift closure from the registry.
+    subscript(_ ref: JavaScriptHostFuncRef) -> (([JSValue]) -> JSValue)? {
+        closures[ref]?.body
+    }
+}
+
 // MARK: - `JSClosure` mechanism note
 //
 // 1. Create a thunk in the JavaScript world, which has a reference
@@ -174,7 +231,7 @@ func _call_host_function_impl(
     _ argv: UnsafePointer<RawJSValue>, _ argc: Int32,
     _ callbackFuncRef: JavaScriptObjectRef
 ) -> Bool {
-    guard let (_, hostFunc) = JSClosure.sharedClosures[hostFuncRef] else {
+    guard let hostFunc = JSClosure.sharedClosures[hostFuncRef] else {
         return true
     }
     let arguments = UnsafeBufferPointer(start: argv, count: Int(argc)).map(\.jsValue)
@@ -195,7 +252,7 @@ func _call_host_function_impl(
 extension JSClosure {
     public func release() {
         isReleased = true
-        Self.sharedClosures[hostFuncRef] = nil
+        Self.sharedClosures.unregister(hostFuncRef)
     }
 }
 
@@ -213,6 +270,6 @@ extension JSClosure {
 
 @_cdecl("_free_host_function_impl")
 func _free_host_function_impl(_ hostFuncRef: JavaScriptHostFuncRef) {
-    JSClosure.sharedClosures[hostFuncRef] = nil
+    JSClosure.sharedClosures.unregister(hostFuncRef)
 }
 #endif