Fix JSClosure leak

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
  }
 }
 
@@ -824,16 +830,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