Support concurrent use of a shared JavaVirtualMachine

Sources/Java2Swift/JavaToSwift.swift

@@ -22,9 +22,6 @@ import JavaKitVM
 import SwiftSyntax
 import SwiftSyntaxBuilder
 
-/// Global instance of the Java virtual machine that we keep alive forever.
-var javaVirtualMachine: JavaVirtualMachine? = nil
-
 /// Command-line utility to drive the export of Java classes into Swift types.
 @main
 struct JavaToSwift: ParsableCommand {
@@ -68,10 +65,8 @@ struct JavaToSwift: ParsableCommand {
       vmOptions.append(contentsOf: classpath)
     }
 
-    let jvm = try JavaVirtualMachine(vmOptions: vmOptions)
-    javaVirtualMachine = jvm
-
-    try run(environment: jvm.environment)
+    let jvm = try JavaVirtualMachine.shared(vmOptions: vmOptions)
+    try run(environment: jvm.environment())
   }
 
   mutating func run(environment: JNIEnvironment) throws {

Sources/JavaKitVM/JavaVirtualMachine.swift

@@ -17,11 +17,20 @@ 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
+  /// Whether to destroy the JVM on deinit.
+  private let destroyOnDeinit: Bool
+
+  /// Adopt an existing JVM pointer.
+  private init(adoptingJVM jvm: JavaVMPointer) {
+    self.jvm = jvm
+    self.destroyOnDeinit = false
+  }
 
   /// Initialize a new Java virtual machine instance.
   ///
@@ -33,15 +42,15 @@ public final class JavaVirtualMachine: @unchecked Sendable {
   ///     be prefixed by the class-path argument described above.
   ///   - ignoreUnrecognized: Whether the JVM should ignore any VM options it
   ///     does not recognize.
-  public init(
+  private init(
     classPath: [String] = [],
     vmOptions: [String] = [],
     ignoreUnrecognized: Bool = true
   ) throws {
     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.
@@ -74,25 +83,191 @@ public final class JavaVirtualMachine: @unchecked Sendable {
     vmArgs.options = optionsBuffer.baseAddress
     vmArgs.nOptions = jint(optionsBuffer.count)
 
-    // Create the JVM.
-    if JNI_CreateJavaVM(&jvm, &environment, &vmArgs) != JNI_OK {
-      throw VMError.failedToCreateVM
+    // Create the JVM instance.
+    if let createError = VMError(fromJNIError: JNI_CreateJavaVM(&jvm, &environment, &vmArgs)) {
+      throw createError
     }
 
     self.jvm = jvm!
-    self.environment = environment!.assumingMemoryBound(to: JNIEnv?.self)
+    self.destroyOnDeinit = true
   }
 
   deinit {
-    // Destroy the JVM.
-    if jvm.pointee!.pointee.DestroyJavaVM(jvm) != JNI_OK {
-      fatalError("Failed to destroy the JVM.")
+    if destroyOnDeinit {
+      // Destroy the JVM.
+      if let resultError = VMError(fromJNIError: jvm.pointee!.pointee.DestroyJavaVM(jvm)) {
+        fatalError("Failed to destroy the JVM: \(resultError)")
+      }
     }
   }
 }
 
+// MARK: Java thread management.
 extension JavaVirtualMachine {
+  /// 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 we failed to attach, report that.
+    if let attachError = VMError(fromJNIError: attachResult) {
+      throw attachError
+    }
+
+    return environment!.assumingMemoryBound(to: JNIEnv?.self)
+  }
+
+  /// Detach the current thread from the Java Virtual Machine. All Java
+  /// threads waiting for this thread to die are notified.
+  public func detachCurrentThread() throws {
+    if let resultError = VMError(fromJNIError: jvm.pointee!.pointee.DetachCurrentThread(jvm)) {
+      throw resultError
+    }
+  }
+}
+
+// MARK: Shared Java Virtual Machine management.
+extension JavaVirtualMachine {
+  /// The globally shared JavaVirtualMachine instance, behind a lock.
+  ///
+  /// TODO: If the use of the lock itself ends up being slow, we could
+  /// use an atomic here instead because our access pattern is fairly
+  /// simple.
+  private static let sharedJVM: LockedState<JavaVirtualMachine?> = .init(initialState: nil)
+
+  /// Access the shared Java Virtual Machine instance.
+  ///
+  /// If there is no shared Java Virtual Machine, create one with the given
+  /// arguments. Note that this function makes no attempt to try to augment
+  /// an existing virtual machine instance with the options given, so it is
+  /// up to clients to ensure that consistent arguments are provided to all
+  /// calls.
+  ///
+  /// - Parameters:
+  ///   - classPath: The directories, JAR files, and ZIP files in which the JVM
+  ///     should look to find classes. This maps to the VM option
+  ///     `-Djava.class.path=`.
+  ///   - vmOptions: Options that should be passed along to the JVM, which will
+  ///     be prefixed by the class-path argument described above.
+  ///   - ignoreUnrecognized: Whether the JVM should ignore any VM options it
+  ///     does not recognize.
+  public static func shared(
+    classPath: [String] = [],
+    vmOptions: [String] = [],
+    ignoreUnrecognized: Bool = true
+  ) throws -> JavaVirtualMachine {
+    try sharedJVM.withLock { (sharedJVMPointer: inout JavaVirtualMachine?) in
+      // If we already have a JavaVirtualMachine instance, return it.
+      if let existingInstance = sharedJVMPointer {
+        return existingInstance
+      }
+
+      while true {
+        var wasExistingVM: Bool = false
+        while true {
+          // Query the JVM itself to determine whether there is a JVM
+          // instance that we don't yet know about.
+          var jvm: UnsafeMutablePointer<JavaVM?>? = nil
+          var numJVMs: jsize = 0
+          if JNI_GetCreatedJavaVMs(&jvm, 1, &numJVMs) == JNI_OK, numJVMs >= 1 {
+            // Adopt this JVM into a new instance of the JavaVirtualMachine
+            // wrapper.
+            let javaVirtualMachine = JavaVirtualMachine(adoptingJVM: jvm!)
+            sharedJVMPointer = javaVirtualMachine
+            return javaVirtualMachine
+          }
+
+          precondition(
+            !wasExistingVM,
+            "JVM reports that an instance of the JVM was already created, but we didn't see it."
+          )
+
+          // Create a new instance of the JVM.
+          let javaVirtualMachine: JavaVirtualMachine
+          do {
+            javaVirtualMachine = try JavaVirtualMachine(
+              classPath: classPath,
+              vmOptions: vmOptions,
+              ignoreUnrecognized: ignoreUnrecognized
+            )
+          } catch VMError.existingVM {
+            // We raced with code outside of this JavaVirtualMachine instance
+            // that created a VM while we were trying to do the same. Go
+            // through the loop again to pick up the underlying JVM pointer.
+            wasExistingVM = true
+            continue
+          }
+
+          sharedJVMPointer = javaVirtualMachine
+          return javaVirtualMachine
+        }
+      }
+    }
+  }
+
+  /// "Forget" the shared JavaVirtualMachine instance.
+  ///
+  /// This will allow the shared JavaVirtualMachine instance to be deallocated.
+  public static func forgetShared() {
+    sharedJVM.withLock { sharedJVMPointer in
+      sharedJVMPointer = nil
+    }
+  }
+}
+
+extension JavaVirtualMachine {
+  /// Describes the kinds of errors that can occur when interacting with JNI.
   enum VMError: Error {
-    case failedToCreateVM
+    /// There is already a Java Virtual Machine.
+    case existingVM
+
+    /// JNI version mismatch error.
+    case jniVersion
+
+    /// Thread is detached from the VM.
+    case threadDetached
+
+    /// Out of memory.
+    case outOfMemory
+
+    /// Invalid arguments.
+    case invalidArguments
+
+    /// Unknown JNI error.
+    case unknown(jint)
+
+    init?(fromJNIError error: jint) {
+      switch error {
+      case JNI_OK: return nil
+      case JNI_EDETACHED: self = .threadDetached
+      case JNI_EVERSION: self = .jniVersion
+      case JNI_ENOMEM: self = .outOfMemory
+      case JNI_EEXIST: self = .existingVM
+      case JNI_EINVAL: self = .invalidArguments
+      default: self = .unknown(error)
+      }
+    }
   }
 }

Sources/_Subprocess/LockedState.swift → Sources/JavaKitVM/LockedState.swift

@@ -2,11 +2,13 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2022 Apple Inc. and the Swift project authors
-// Licensed under Apache License v2.0 with Runtime Library Exception
+// Copyright (c) 2024 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
 //
-// See https://swift.org/LICENSE.txt for license information
-// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
 //
 //===----------------------------------------------------------------------===//
 

Tests/JavaKitTests/BasicRuntimeTests.swift

@@ -17,19 +17,23 @@ import JavaKitNetwork
 import JavaKitVM
 import XCTest // NOTE: Workaround for https://github.com/swiftlang/swift-java/issues/43
 
-@MainActor
-let jvm = try! JavaVirtualMachine(vmOptions: [])
+/// Handy reference to the JVM abstraction.
+var jvm: JavaVirtualMachine {
+  get throws {
+    try .shared()
+  }
+}
 
-@MainActor
 class BasicRuntimeTests: XCTestCase {
   func testJavaObjectManagement() async throws {
     if isLinux {
       throw XCTSkip("Attempts to refcount a null pointer on Linux")
     }
 
+    let environment = try 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 +45,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 +60,10 @@ class BasicRuntimeTests: XCTestCase {
       throw XCTSkip("Attempts to refcount a null pointer on Linux")
     }
 
+    let environment = try 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 +74,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 jvm.environment()
+
+    let urlConnectionClass = try JavaClass<URLConnection>(in: environment)
     XCTAssert(urlConnectionClass.getDefaultAllowUserInteraction() == false)
   }
 
   func testClassInstanceLookup() async throws {
+    let environment = try jvm.environment()
+
     do {
-      _ = try JavaClass<Nonexistent>(in: jvm.environment)
+      _ = try JavaClass<Nonexistent>(in: environment)
     } catch {
       XCTAssertEqual(String(describing: error), "org/swift/javakit/Nonexistent")
     }