replace read/write lock in JarResource to avoid virtual threads pinning

independent-projects/bootstrap/runner/src/main/java/io/quarkus/bootstrap/runner/JarResource.java

@@ -11,9 +11,8 @@
 import java.security.ProtectionDomain;
 import java.security.cert.Certificate;
 import java.util.Objects;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReadWriteLock;
-import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
 import java.util.zip.ZipEntry;
@@ -30,25 +29,27 @@ public class JarResource implements ClassLoadingResource {
     private final ManifestInfo manifestInfo;
     private final Path jarPath;
 
-    private final Lock readLock;
-    private final Lock writeLock;
-
     private volatile ProtectionDomain protectionDomain;
 
-    //Guarded by the read/write lock; open/close operations on the JarFile require the exclusive lock,
-    //while using an existing open reference can use the shared lock.
-    //If a lock is acquired, and as long as it's owned, we ensure that the zipFile reference
-    //points to an open JarFile instance, and read operations are valid.
-    //To close the jar, the exclusive lock must be owned, and reference will be set to null before releasing it.
-    //Likewise, opening a JarFile requires the exclusive lock.
+    // Guarded by a state machine + an atomic reader counter that emulate the behaviour of a read/write lock.
+    // To enable virtual threads compatibility and avoid pinning it is not possible to use an explicit read/write lock
+    // because the zipFile access may happen inside a native call (for example triggered by the RunnerClassLoader)
+    // and then it is necessary to avoid blocking on it.
     private volatile JarFile zipFile;
 
+    private AtomicReference<ZipFileLoadingState> zipFileLoadingState = new AtomicReference<>(ZipFileLoadingState.NOT_LOADED);
+    private AtomicInteger zipFileReadersCounter = new AtomicInteger(0);
+
+    enum ZipFileLoadingState {
+        NOT_LOADED, // The zipFile loading hasn't been started at all
+        LOADING, // Only one thread can start the zipFile loading process
+        LOADED, // The zipFile is fully loaded and can be used by multiple readers at the same time
+        CLOSING // Initiate the closing sequence waiting existing readers to terminate and disallowing new readers to enter
+    }
+
     public JarResource(ManifestInfo manifestInfo, Path jarPath) {
         this.manifestInfo = manifestInfo;
         this.jarPath = jarPath;
-        final ReadWriteLock readWriteLock = new ReentrantReadWriteLock();
-        this.readLock = readWriteLock.readLock();
-        this.writeLock = readWriteLock.writeLock();
     }
 
     @Override
@@ -92,7 +93,7 @@ public byte[] getResourceData(String resource) {
                 throw new RuntimeException("Failed to read zip entry " + resource, e);
             }
         } finally {
-            readLock.unlock();
+            zipFileReadersCounter.decrementAndGet();
         }
     }
 
@@ -137,7 +138,7 @@ public URL getResourceURL(String resource) {
                 throw new RuntimeException(e);
             }
         } finally {
-            readLock.unlock();
+            zipFileReadersCounter.decrementAndGet();
         }
     }
 
@@ -152,50 +153,35 @@ public ProtectionDomain getProtectionDomain() {
     }
 
     private JarFile readLockAcquireAndGetJarReference() {
-        while (true) {
-            readLock.lock();
-            final JarFile zipFileLocal = this.zipFile;
-            if (zipFileLocal != null) {
-                //Expected fast path: returns a reference to the open JarFile while owning the readLock
-                return zipFileLocal;
-            } else {
-                //This Lock implementation doesn't allow upgrading a readLock to a writeLock, so release it
-                //as we're going to need the WriteLock.
-                readLock.unlock();
-                //trigger the JarFile being (re)opened.
-                ensureJarFileIsOpen();
-                //Now since we no longer own any lock, we need to try again to obtain the readLock
-                //and check for the reference still being valid.
-                //This exposes us to a race with closing the just-opened JarFile;
-                //however this should be extremely rare, so we can trust we won't loop much;
-                //A counter doesn't seem necessary, as in fact we know that methods close()
-                //and resetInternalCaches() are invoked each at most once, which limits the amount
-                //of loops here in practice.
+        while (zipFileLoadingState.get() != ZipFileLoadingState.LOADED) {
+            if (zipFileLoadingState.compareAndSet(ZipFileLoadingState.NOT_LOADED, ZipFileLoadingState.LOADING)) {
+                return loadZipFile();
             }
+            Thread.yield();
         }
+        zipFileReadersCounter.incrementAndGet();
+        return this.zipFile;
     }
 
-    private void ensureJarFileIsOpen() {
-        writeLock.lock();
+    private JarFile loadZipFile() {
         try {
-            if (this.zipFile == null) {
-                try {
-                    this.zipFile = JarFiles.create(jarPath.toFile());
-                } catch (IOException e) {
-                    throw new RuntimeException("Failed to open " + jarPath, e);
-                }
-            }
-        } finally {
-            writeLock.unlock();
+            this.zipFile = JarFiles.create(jarPath.toFile());
+        } catch (IOException e) {
+            throw new RuntimeException("Failed to open " + jarPath, e);
         }
+        zipFileReadersCounter.incrementAndGet();
+        zipFileLoadingState.set(ZipFileLoadingState.LOADED);
+        return this.zipFile;
     }
 
     @Override
     public void close() {
-        writeLock.lock();
         try {
             final JarFile zipFileLocal = this.zipFile;
-            if (zipFileLocal != null) {
+            if (zipFileLocal != null && zipFileLoadingState.compareAndSet(ZipFileLoadingState.LOADED, ZipFileLoadingState.CLOSING)) {
+                while (zipFileReadersCounter.get() > 0) {
+                    Thread.yield();
+                }
                 try {
                     this.zipFile = null;
                     zipFileLocal.close();
@@ -204,7 +190,7 @@ public void close() {
                 }
             }
         } finally {
-            writeLock.unlock();
+            zipFileLoadingState.set(ZipFileLoadingState.NOT_LOADED);
         }
     }