Lazily attach threads to the JVM when querying the JNI environment

Instead of having the JavaVirtualMachine instance keep track of the JNI
environment for the thread in which it was initialized, provide a JNI
environment query function that attaches the current thread to the JVM
if needed. This lets us query the JNI environment for multiple threads,
should we want to.

We no longer need to force all of the JavaKit runtime tests to the main
actor.

Author: DougGregor

Sources/Java2Swift/JavaToSwift.swift

@@ -71,7 +71,7 @@ struct JavaToSwift: ParsableCommand {
     let jvm = try JavaVirtualMachine(vmOptions: vmOptions)
     javaVirtualMachine = jvm
 
-    try run(environment: jvm.environment)
+    try run(environment: jvm.environment())
   }
 
   mutating func run(environment: JNIEnvironment) throws {

Sources/JavaKitVM/JavaVirtualMachine.swift

@@ -17,12 +17,12 @@ import JavaKit
 typealias JavaVMPointer = UnsafeMutablePointer<JavaVM?>
 
 public final class JavaVirtualMachine: @unchecked Sendable {
+  /// The JNI version that we depend on.
+  static let jniVersion = JNI_VERSION_1_6
+
   /// The Java virtual machine instance.
   private let jvm: JavaVMPointer
 
-  /// The JNI environment for the JVM.
-  public let environment: JNIEnvironment
-
   /// Initialize a new Java virtual machine instance.
   ///
   /// - Parameters:
@@ -41,7 +41,7 @@ public final class JavaVirtualMachine: @unchecked Sendable {
     var jvm: JavaVMPointer? = nil
     var environment: UnsafeMutableRawPointer? = nil
     var vmArgs = JavaVMInitArgs()
-    vmArgs.version = JNI_VERSION_1_6
+    vmArgs.version = JavaVirtualMachine.jniVersion
     vmArgs.ignoreUnrecognized = jboolean(ignoreUnrecognized ? JNI_TRUE : JNI_FALSE)
 
     // Construct the complete list of VM options.
@@ -80,7 +80,6 @@ public final class JavaVirtualMachine: @unchecked Sendable {
     }
 
     self.jvm = jvm!
-    self.environment = environment!.assumingMemoryBound(to: JNIEnv?.self)
   }
 
   deinit {
@@ -89,10 +88,55 @@ public final class JavaVirtualMachine: @unchecked Sendable {
       fatalError("Failed to destroy the JVM.")
     }
   }
+
+  /// Produce the JNI environment for the active thread, attaching this
+  /// thread to the JVM if it isn't already.
+  ///
+  /// - Parameter
+  ///   - asDaemon: Whether this thread should be treated as a daemon
+  ///     thread in the Java Virtual Machine.
+  public func environment(asDaemon: Bool = false) throws -> JNIEnvironment {
+    // Check whether this thread is already attached. If so, return the
+    // corresponding environment.
+    var environment: UnsafeMutableRawPointer? = nil
+    let getEnvResult = jvm.pointee!.pointee.GetEnv(
+      jvm,
+      &environment,
+      JavaVirtualMachine.jniVersion
+    )
+    if getEnvResult == JNI_OK, let environment {
+      return environment.assumingMemoryBound(to: JNIEnv?.self)
+    }
+
+    // Attach the current thread to the JVM.
+    let attachResult: jint
+    if asDaemon {
+      attachResult = jvm.pointee!.pointee.AttachCurrentThreadAsDaemon(jvm, &environment, nil)
+    } else {
+      attachResult = jvm.pointee!.pointee.AttachCurrentThread(jvm, &environment, nil)
+    }
+
+    if attachResult == JNI_OK, let environment {
+      return environment.assumingMemoryBound(to: JNIEnv?.self)
+    }
+
+    throw VMError.failedToAttachThread
+  }
+
+  /// Detach the current thread from the Java Virtual Machine. All Java
+  /// threads waiting for this thread to die are notified.
+  public func detachCurrentThread() throws {
+    let result = jvm.pointee!.pointee.DetachCurrentThread(jvm)
+    if result != JNI_OK {
+      throw VMError.failedToDetachThread
+    }
+  }
 }
 
 extension JavaVirtualMachine {
   enum VMError: Error {
     case failedToCreateVM
+    case failedToAttachThread
+    case failedToDetachThread
   }
 }

Tests/JavaKitTests/BasicRuntimeTests.swift

@@ -20,16 +20,16 @@ import XCTest // NOTE: Workaround for https://github.com/swiftlang/swift-java/is
 @MainActor
 let jvm = try! JavaVirtualMachine(vmOptions: [])
 
-@MainActor
 class BasicRuntimeTests: XCTestCase {
   func testJavaObjectManagement() async throws {
     if isLinux {
       throw XCTSkip("Attempts to refcount a null pointer on Linux")
     }
 
+    let environment = try await jvm.environment()
     let sneakyJavaThis: jobject
     do {
-      let object = JavaObject(environment: jvm.environment)
+      let object = JavaObject(environment: environment)
       XCTAssert(object.toString().starts(with: "java.lang.Object"))
 
       // Make sure this object was promoted to a global reference.
@@ -41,10 +41,10 @@ class BasicRuntimeTests: XCTestCase {
 
     // The reference should now be invalid, because we've deleted the
     // global reference.
-    XCTAssertEqual(jvm.environment.pointee?.pointee.GetObjectRefType(jvm.environment, sneakyJavaThis), JNIInvalidRefType)
+    XCTAssertEqual(environment.pointee?.pointee.GetObjectRefType(environment, sneakyJavaThis), JNIInvalidRefType)
 
     // 'super' and 'as' don't require allocating a new holder.
-    let url = try URL("http://swift.org", environment: jvm.environment)
+    let url = try URL("http://swift.org", environment: environment)
     let superURL = url.super
     XCTAssert(url.javaHolder === superURL.javaHolder)
     let urlAgain = superURL.as(URL.self)!
@@ -56,8 +56,10 @@ class BasicRuntimeTests: XCTestCase {
       throw XCTSkip("Attempts to refcount a null pointer on Linux")
     }
 
+    let environment = try await jvm.environment()
+
     do {
-      _ = try URL("bad url", environment: jvm.environment)
+      _ = try URL("bad url", environment: environment)
     } catch {
       XCTAssert(String(describing: error) == "no protocol: bad url")
     }
@@ -68,13 +70,17 @@ class BasicRuntimeTests: XCTestCase {
       throw XCTSkip("Attempts to refcount a null pointer on Linux")
     }
 
-    let urlConnectionClass = try JavaClass<URLConnection>(in: jvm.environment)
+    let environment = try await jvm.environment()
+
+    let urlConnectionClass = try JavaClass<URLConnection>(in: environment)
     XCTAssert(urlConnectionClass.getDefaultAllowUserInteraction() == false)
   }
 
   func testClassInstanceLookup() async throws {
+    let environment = try await jvm.environment()
+
     do {
-      _ = try JavaClass<Nonexistent>(in: jvm.environment)
+      _ = try JavaClass<Nonexistent>(in: environment)
     } catch {
       XCTAssertEqual(String(describing: error), "org/swift/javakit/Nonexistent")
     }