Cleanup how we handle local references in JNI interactions

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+NativeTranslation.swift

@@ -1293,6 +1293,7 @@ extension JNISwift2JavaGenerator {
           """
           let class$ = environment.interface.GetObjectClass(environment, \(placeholder))
           let methodID$ = environment.interface.GetMethodID(environment, class$, "apply", "\(methodSignature.mangledName)")!
+          environment.interface.DeleteLocalRef(environment, class$)
           let arguments$: [jvalue] = [\(arguments.joined(separator: ", "))]
           """
         )

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+SwiftThunkPrinting.swift

@@ -171,30 +171,41 @@ extension JNISwift2JavaGenerator {
     }
 
     printer.printBraceBlock(function.swiftDecl.signatureString) { printer in
-      var upcallArguments = zip(
-        function.originalFunctionSignature.parameters,
-        function.parameterConversions
-      ).map { param, conversion in
-        // Wrap-java does not extract parameter names, so no labels
-        conversion.render(&printer, param.parameterName!)
-      }
+      let resultType = function.originalFunctionSignature.result.type
+      let returnStmt = !resultType.isVoid ? "return " : ""
+      // If the protocol function is non-throwing, we have no option but to force try.
+      // The error thrown by `withLocalFrame` is an OOM error anyway.
+      let withLocalFrameTryKeyword = function.originalFunctionSignature.isThrowing ? "try" : "try!"
+
+      // Push a local JNI frame so refs created during this upcall are freed on exit.
+      // When called from a Swift async context (e.g. cooperative thread pool) there is
+      // no enclosing JNI frame, so refs would otherwise accumulate indefinitely. When
+      // called from a Java-initiated native call there is already a frame, but pushing
+      // a sub-frame still frees refs earlier and prevents overflow within a single call.
+      let paramCount = function.originalFunctionSignature.parameters.count
+      let estimatedRefCount = paramCount * 2 + 4
+      printer.print("let environment$ = try! JavaVirtualMachine.shared().environment()")
+      printer.printBraceBlock("\(returnStmt)\(withLocalFrameTryKeyword) environment$.withLocalFrame(capacity: \(estimatedRefCount))") { printer in
+        var upcallArguments = zip(
+          function.originalFunctionSignature.parameters,
+          function.parameterConversions
+        ).map { param, conversion in
+          // Wrap-java does not extract parameter names, so no labels
+          conversion.render(&printer, param.parameterName!)
+        }
 
-      // If the underlying translated method requires
-      // a SwiftArena, we pass in the global arena
-      if translatedDecl.translatedFunctionSignature.requiresSwiftArena {
-        upcallArguments.append("JavaSwiftArena.defaultAutoArena")
-      }
+        // If the underlying translated method requires
+        // a SwiftArena, we pass in the global arena
+        if translatedDecl.translatedFunctionSignature.requiresSwiftArena {
+          upcallArguments.append("JavaSwiftArena.defaultAutoArena")
+        }
 
-      let tryClause = function.originalFunctionSignature.isThrowing ? "try " : ""
-      let javaUpcall =
-        "\(tryClause)\(wrapper.javaInterfaceVariableName).\(function.swiftFunctionName)(\(upcallArguments.joined(separator: ", ")))"
+        let tryClause = function.originalFunctionSignature.isThrowing ? "try " : ""
+        let javaUpcall =
+          "\(tryClause)\(wrapper.javaInterfaceVariableName).\(function.swiftFunctionName)(\(upcallArguments.joined(separator: ", ")))"
 
-      let resultType = function.originalFunctionSignature.result.type
-      let result = function.resultConversion.render(&printer, javaUpcall)
-      if resultType.isVoid {
-        printer.print(result)
-      } else {
-        printer.print("return \(result)")
+        let result = function.resultConversion.render(&printer, javaUpcall)
+        printer.print("\(returnStmt)\(result)")
       }
     }
   }

Sources/JavaKit/Helpers/_JNIMethodIDCache.swift

@@ -46,6 +46,8 @@ public final class _JNIMethodIDCache: Sendable {
         fatalError("Method \(method.signature) with signature \(method.signature) not found in class \(className)")
       }
     }
+    // clazz is still needed by GetMethodID above; delete the local ref only after reduce completes.
+    environment.interface.DeleteLocalRef(environment, clazz)
   }
 
   public subscript(_ method: Method) -> jmethodID? {

Sources/SwiftJava/Exceptions/ExceptionHandling.swift

@@ -54,6 +54,7 @@ extension JNIEnvironment {
       }
       return
     }
+    defer { self.interface.DeleteLocalRef(self, exceptionClass) }
 
     _ = interface.ThrowNew(self, exceptionClass, javaException.message ?? "")
   }

Sources/SwiftJava/JavaObject+MethodCalls.swift

@@ -108,6 +108,7 @@ extension AnyJavaObject {
     let thisClass = try environment.translatingJNIExceptions {
       environment.interface.GetObjectClass(environment, javaThis)
     }!
+    defer { environment.interface.DeleteLocalRef(environment, thisClass) }
 
     return try environment.translatingJNIExceptions {
       try Self.javaMethodLookup(
@@ -131,6 +132,7 @@ extension AnyJavaObject {
     let thisClass = try environment.translatingJNIExceptions {
       environment.interface.GetObjectClass(environment, javaThis)
     }!
+    defer { environment.interface.DeleteLocalRef(environment, thisClass) }
 
     return try environment.translatingJNIExceptions {
       try Self.javaMethodLookup(
@@ -269,7 +271,6 @@ extension AnyJavaObject {
         in: environment
       )
 
-      // Retrieve the constructor, then map the arguments and call it.
       let jniArgs = getJValues(repeat each arguments, in: environment)
       return try environment.translatingJNIExceptions {
         environment.interface.NewObjectA!(environment, thisClass, methodID, jniArgs)
@@ -285,6 +286,7 @@ extension AnyJavaObject {
 
     // Retrieve the Java class instance from the object.
     let thisClass = environment.interface.GetObjectClass(environment, this)!
+    defer { environment.interface.DeleteLocalRef(environment, thisClass) }
 
     return environment.interface.GetFieldID(environment, thisClass, fieldName, FieldType.jniMangling)
   }
@@ -337,7 +339,6 @@ extension JavaClass {
       )
     }!
 
-    // Retrieve the method that performs this call, then
     let jniMethod = Result.jniStaticMethodCall(in: environment)
     let jniArgs = getJValues(repeat each arguments, in: environment)
     let jniResult = try environment.translatingJNIExceptions {
@@ -371,7 +372,6 @@ extension JavaClass {
       )
     }!
 
-    // Retrieve the method that performs this call, then
     let jniMethod = environment.interface.CallStaticVoidMethodA
     let jniArgs = getJValues(repeat each arguments, in: environment)
     try environment.translatingJNIExceptions {

Sources/SwiftJava/JavaObjectHolder.swift

@@ -26,6 +26,12 @@ public final class JavaObjectHolder {
   public init(object: jobject, environment: JNIEnvironment) {
     self.object = environment.interface.NewGlobalRef(environment, object)
     self.environment = environment
+
+    // If we are taking over a local ref, let's delete that.
+    let refType = environment.interface.GetObjectRefType(environment, object)
+    if refType == JNILocalRefType {
+      environment.interface.DeleteLocalRef(environment, object)
+    }
   }
 
   /// Forget this Java object, meaning that it is no longer used from anywhere

Sources/SwiftJavaRuntimeSupport/_JNIMethodIDCache.swift

@@ -63,6 +63,7 @@ public final class _JNIMethodIDCache: Sendable {
     let clazz: jobject
     if let jniClass = environment.interface.FindClass(environment, className) {
       clazz = environment.interface.NewGlobalRef(environment, jniClass)!
+      environment.interface.DeleteLocalRef(environment, jniClass)
       self.javaObjectHolder = nil
     } else {
       // Clear any ClassNotFound exceptions from FindClass

Tests/JExtractSwiftTests/JNI/JNIClosureTests.swift

@@ -70,6 +70,7 @@ struct JNIClosureTests {
           SwiftModule.emptyClosure(closure: {
             let class$ = environment.interface.GetObjectClass(environment, closure)
             let methodID$ = environment.interface.GetMethodID(environment, class$, "apply", "()V")!
+            environment.interface.DeleteLocalRef(environment, class$)
             let arguments$: [jvalue] = []
             environment.interface.CallVoidMethodA(environment, closure, methodID$, arguments$)
           }
@@ -127,6 +128,7 @@ struct JNIClosureTests {
           SwiftModule.closureWithArgumentsAndReturn(closure: { _0, _1 in
             let class$ = environment.interface.GetObjectClass(environment, closure)
             let methodID$ = environment.interface.GetMethodID(environment, class$, "apply", "(JZ)J")!
+            environment.interface.DeleteLocalRef(environment, class$)
             let arguments$: [jvalue] = [_0.getJValue(in: environment), _1.getJValue(in: environment)]
             return Int64(fromJNI: environment.interface.CallLongMethodA(environment, closure, methodID$, arguments$), in: environment)
           }

Tests/JExtractSwiftTests/JNI/JNIProtocolTests.swift

@@ -315,18 +315,24 @@ struct JNIProtocolTests {
         """
         extension SwiftJavaSomeProtocolWrapper {
           public func method() {
-            _javaSomeProtocolInterface.method()
+            let environment$ = try! JavaVirtualMachine.shared().environment()
+            try! environment$.withLocalFrame(capacity: 4) {
+              _javaSomeProtocolInterface.method()
+            }
           }
           public func withObject(c: SomeClass) -> SomeClass {
-            let cClass = try! JavaClass<JavaSomeClass>(environment: JavaVirtualMachine.shared().environment())
-            let cPointer = UnsafeMutablePointer<SomeClass>.allocate(capacity: 1)
-            cPointer.initialize(to: c)
-            guard let unwrapped$ = _javaSomeProtocolInterface.withObject(cClass.wrapMemoryAddressUnsafe(Int64(Int(bitPattern: cPointer))), JavaSwiftArena.defaultAutoArena) else {
-              fatalError("Upcall to withObject unexpectedly returned nil")
+            let environment$ = try! JavaVirtualMachine.shared().environment()
+            return try! environment$.withLocalFrame(capacity: 6) {
+              let cClass = try! JavaClass<JavaSomeClass>(environment: JavaVirtualMachine.shared().environment())
+              let cPointer = UnsafeMutablePointer<SomeClass>.allocate(capacity: 1)
+              cPointer.initialize(to: c)
+              guard let unwrapped$ = _javaSomeProtocolInterface.withObject(cClass.wrapMemoryAddressUnsafe(Int64(Int(bitPattern: cPointer))), JavaSwiftArena.defaultAutoArena) else {
+                fatalError("Upcall to withObject unexpectedly returned nil")
+              }
+              let result$MemoryAddress$ = unwrapped$.as(JavaJNISwiftInstance.self)!.memoryAddress()
+              let result$Pointer = UnsafeMutablePointer<SomeClass>(bitPattern: Int(result$MemoryAddress$))!
+              return result$Pointer.pointee
             }
-            let result$MemoryAddress$ = unwrapped$.as(JavaJNISwiftInstance.self)!.memoryAddress()
-            let result$Pointer = UnsafeMutablePointer<SomeClass>(bitPattern: Int(result$MemoryAddress$))!
-            return result$Pointer.pointee
           }
         }
         """,