Reduce locking in HostClassLoader

engine/runtime-language-epb/src/main/java/org/enso/interpreter/epb/HostClassLoader.java

@@ -17,8 +17,10 @@
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.net.URLClassLoader;
-import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
 import org.graalvm.polyglot.Context;
 
 /**
@@ -29,7 +31,8 @@
 @ExportLibrary(InteropLibrary.class)
 final class HostClassLoader extends URLClassLoader implements AutoCloseable, TruffleObject {
 
-  private final Map<String, Class<?>> loadedClasses = new ConcurrentHashMap<>();
+  private final ConcurrentHashMap<String, Future<Class<?>>> loadedClasses =
+      new ConcurrentHashMap<>();
   private static final Logger logger = System.getLogger(HostClassLoader.class.getName());
   // Classes from "org.graalvm" packages are loaded either by a class loader for the boot
   // module layer, or by a specific class loader, depending on how enso is run. For example,
@@ -62,43 +65,53 @@ public Class<?> loadClass(String name) throws ClassNotFoundException {
 
   @Override
   @CompilerDirectives.TruffleBoundary
+  @SuppressWarnings("unchecked")
   protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
     logger.log(Logger.Level.TRACE, "Loading class {0}", name);
-    var l = loadedClasses.get(name);
-    if (l != null) {
-      logger.log(Logger.Level.TRACE, "Class {0} found in cache", name);
-      return l;
+    var placeholder = new CompletableFuture[1];
+    var pendingClass =
+        loadedClasses.computeIfAbsent(
+            name,
+            t -> {
+              logger.log(Logger.Level.TRACE, "Class {0} found in cache", name);
+              var f = new CompletableFuture<Class<?>>();
+              placeholder[0] = f;
+              return f;
+            });
+    if (placeholder[0] != null) {
+      placeholder[0].complete(loadClassUnsafe(name, resolve));
     }
-    synchronized (this) {
-      l = loadedClasses.get(name);
-      if (l != null) {
-        logger.log(Logger.Level.TRACE, "Class {0} found in cache", name);
-        return l;
-      }
-      if (!isRuntimeModInBootLayer && name.startsWith("org.graalvm")) {
-        return polyglotClassLoader.loadClass(name);
-      }
-      if (name.startsWith("org.slf4j")) {
-        // Delegating to system class loader ensures that log classes are not loaded again
-        // and do not require special setup. In other words, it is using log configuration that
-        // has been setup by the runner that started the process. See #11641.
-        return polyglotClassLoader.loadClass(name);
-      }
-      try {
-        l = findClass(name);
-        if (resolve) {
-          l.getMethods();
-        }
-        logger.log(Logger.Level.TRACE, "Class {0} found, putting in cache", name);
-        loadedClasses.put(name, l);
-        return l;
-      } catch (ClassNotFoundException ex) {
-        logger.log(Logger.Level.TRACE, "Class {0} not found, delegating to super", name);
-        return super.loadClass(name, resolve);
-      } catch (Throwable e) {
-        logger.log(Logger.Level.TRACE, "Failure while loading a class: " + e.getMessage(), e);
-        throw e;
+    try {
+      return pendingClass.get();
+    } catch (InterruptedException | ExecutionException e) {
+      throw new ClassNotFoundException("Unable to find class " + name, e);
+    }
+  }
+
+  /** Find a class with a given name without giving any Thread-safety guarantees. */
+  private Class<?> loadClassUnsafe(String name, boolean resolve) throws ClassNotFoundException {
+    if (!isRuntimeModInBootLayer && name.startsWith("org.graalvm")) {
+      return polyglotClassLoader.loadClass(name);
+    }
+    if (name.startsWith("org.slf4j")) {
+      // Delegating to system class loader ensures that log classes are not loaded again
+      // and do not require special setup. In other words, it is using log configuration that
+      // has been setup by the runner that started the process. See #11641.
+      return polyglotClassLoader.loadClass(name);
+    }
+    try {
+      var l = findClass(name);
+      if (resolve) {
+        l.getMethods();
       }
+      logger.log(Logger.Level.TRACE, "Class {0} found, putting in cache", name);
+      return l;
+    } catch (ClassNotFoundException ex) {
+      logger.log(Logger.Level.TRACE, "Class {0} not found, delegating to super", name);
+      return super.loadClass(name, resolve);
+    } catch (Throwable e) {
+      logger.log(Logger.Level.TRACE, "Failure while loading a class: " + e.getMessage(), e);
+      throw e;
     }
   }