GH-1469: Infer the allocation manager from the factory found on the classpath

memory/memory-core/src/main/java/org/apache/arrow/memory/CheckAllocator.java

@@ -18,9 +18,13 @@
 
 import java.io.IOException;
 import java.net.URL;
+import java.util.EnumMap;
 import java.util.Enumeration;
 import java.util.LinkedHashSet;
+import java.util.Map;
 import java.util.Set;
+import org.apache.arrow.util.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -29,85 +33,90 @@
  */
 final class CheckAllocator {
   private static final Logger logger = LoggerFactory.getLogger(CheckAllocator.class);
-  // unique package names needed by JPMS module naming
-  private static final String ALLOCATOR_PATH_CORE =
-      "org/apache/arrow/memory/DefaultAllocationManagerFactory.class";
-  private static final String ALLOCATOR_PATH_UNSAFE =
-      "org/apache/arrow/memory/unsafe/DefaultAllocationManagerFactory.class";
-  private static final String ALLOCATOR_PATH_NETTY =
-      "org/apache/arrow/memory/netty/DefaultAllocationManagerFactory.class";
 
-  private CheckAllocator() {}
+  /**
+   * The DefaultAllocationManagerFactory implementations to look for, in order of precedence. Unsafe
+   * and Netty match {@link DefaultAllocationManagerOption.AllocationManagerType}, so the logged
+   * name can be used to set the type explicitly.
+   */
+  private enum Allocator {
+    // only provided by the memory-core tests
+    Core("org.apache.arrow.memory.DefaultAllocationManagerFactory"),
+    // unique package names needed by JPMS module naming
+    Unsafe("org.apache.arrow.memory.unsafe.DefaultAllocationManagerFactory"),
+    Netty("org.apache.arrow.memory.netty.DefaultAllocationManagerFactory");
 
-  static String check() {
-    Set<URL> urls = scanClasspath();
-    URL rootAllocator = assertOnlyOne(urls);
-    reportResult(rootAllocator);
-    if (rootAllocator.getPath().contains("memory-core")
-        || rootAllocator.getPath().contains("/org/apache/arrow/memory/core/")) {
-      return "org.apache.arrow.memory.DefaultAllocationManagerFactory";
-    } else if (rootAllocator.getPath().contains("memory-unsafe")
-        || rootAllocator.getPath().contains("/org/apache/arrow/memory/unsafe/")) {
-      return "org.apache.arrow.memory.unsafe.DefaultAllocationManagerFactory";
-    } else if (rootAllocator.getPath().contains("memory-netty")
-        || rootAllocator.getPath().contains("/org/apache/arrow/memory/netty/")) {
-      return "org.apache.arrow.memory.netty.DefaultAllocationManagerFactory";
-    } else {
-      throw new IllegalStateException(
-          "Unknown allocation manager type to infer. Current: " + rootAllocator.getPath());
-    }
-  }
+    private final String factoryClassName;
+    private final String resourcePath;
 
-  @SuppressWarnings("URLEqualsHashCode")
-  private static Set<URL> scanClasspath() {
-    // LinkedHashSet appropriate here because it preserves insertion order
-    // during iteration
-    Set<URL> allocatorPathSet = new LinkedHashSet<>();
-    try {
-      ClassLoader allocatorClassLoader = CheckAllocator.class.getClassLoader();
-      Enumeration<URL> paths;
-      if (allocatorClassLoader == null) {
-        paths = ClassLoader.getSystemResources(ALLOCATOR_PATH_CORE);
-        if (!paths.hasMoreElements()) {
-          paths = ClassLoader.getSystemResources(ALLOCATOR_PATH_UNSAFE);
-        }
-        if (!paths.hasMoreElements()) {
-          paths = ClassLoader.getSystemResources(ALLOCATOR_PATH_NETTY);
-        }
-      } else {
-        paths = allocatorClassLoader.getResources(ALLOCATOR_PATH_CORE);
-        if (!paths.hasMoreElements()) {
-          paths = allocatorClassLoader.getResources(ALLOCATOR_PATH_UNSAFE);
-        }
-        if (!paths.hasMoreElements()) {
-          paths = allocatorClassLoader.getResources(ALLOCATOR_PATH_NETTY);
-        }
-      }
-      while (paths.hasMoreElements()) {
-        URL path = paths.nextElement();
-        allocatorPathSet.add(path);
-      }
-    } catch (IOException ioe) {
-      logger.error("Error getting resources from path", ioe);
+    Allocator(String factoryClassName) {
+      this.factoryClassName = factoryClassName;
+      this.resourcePath = factoryClassName.replace('.', '/') + ".class";
     }
-    return allocatorPathSet;
   }
 
-  private static void reportResult(URL rootAllocator) {
-    String path = rootAllocator.getPath();
-    String subPath = path.substring(path.indexOf("memory"));
-    logger.info("Using DefaultAllocationManager at {}", subPath);
+  private CheckAllocator() {}
+
+  static String check() {
+    return check(CheckAllocator.class.getClassLoader());
   }
 
-  private static URL assertOnlyOne(Set<URL> urls) {
-    if (urls.size() > 1) {
-      logger.warn("More than one DefaultAllocationManager on classpath. Choosing first found");
-    }
-    if (urls.isEmpty()) {
+  @VisibleForTesting
+  static String check(@Nullable ClassLoader classLoader) {
+    Map<Allocator, Set<URL>> found = scanClasspath(classLoader);
+    if (found.isEmpty()) {
       throw new RuntimeException(
           "No DefaultAllocationManager found on classpath. Can't allocate Arrow buffers."
               + " Please consider adding arrow-memory-netty or arrow-memory-unsafe as a dependency.");
     }
-    return urls.iterator().next();
+    Map.Entry<Allocator, Set<URL>> first = found.entrySet().iterator().next();
+    Allocator allocator = first.getKey();
+    Set<URL> urls = first.getValue();
+    if (found.size() > 1) {
+      logger.warn(
+          "Found allocation manager types {} on the classpath, using {}. Set the {} system property"
+              + " or the {} environment variable to choose one explicitly",
+          found.keySet(),
+          allocator,
+          DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME,
+          DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_ENV_NAME);
+    }
+    if (urls.size() > 1) {
+      logger.warn(
+          "More than one {} DefaultAllocationManager on classpath {}. Choosing first found",
+          allocator,
+          urls);
+    }
+    logger.info(
+        "Allocation manager type not specified, inferred {} from the classpath at {}",
+        allocator,
+        urls.iterator().next());
+    return allocator.factoryClassName;
+  }
+
+  @SuppressWarnings("URLEqualsHashCode")
+  private static Map<Allocator, Set<URL>> scanClasspath(@Nullable ClassLoader classLoader) {
+    // EnumMap iterates in declaration order, which is the order of precedence
+    Map<Allocator, Set<URL>> found = new EnumMap<>(Allocator.class);
+    for (Allocator allocator : Allocator.values()) {
+      try {
+        Enumeration<URL> paths =
+            classLoader == null
+                ? ClassLoader.getSystemResources(allocator.resourcePath)
+                : classLoader.getResources(allocator.resourcePath);
+        // LinkedHashSet appropriate here because it preserves insertion order
+        // during iteration
+        Set<URL> urls = new LinkedHashSet<>();
+        while (paths.hasMoreElements()) {
+          urls.add(paths.nextElement());
+        }
+        if (!urls.isEmpty()) {
+          found.put(allocator, urls);
+        }
+      } catch (IOException ioe) {
+        logger.error("Error getting resources from path", ioe);
+      }
+    }
+    return found;
   }
 }

memory/memory-core/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java

@@ -88,11 +88,7 @@ static AllocationManager.Factory getDefaultAllocationManagerFactory() {
         DEFAULT_ALLOCATION_MANAGER_FACTORY = getUnsafeFactory();
         break;
       case Unknown:
-        String allocationManagerClass = CheckAllocator.check();
-        DEFAULT_ALLOCATION_MANAGER_FACTORY = getFactory(allocationManagerClass);
-        LOGGER.info(
-            "allocation manager type not specified, inferred {} from the classpath",
-            allocationManagerClass);
+        DEFAULT_ALLOCATION_MANAGER_FACTORY = getFactory(CheckAllocator.check());
         break;
       default:
         throw new IllegalStateException("Unknown allocation manager type: " + type);

memory/memory-core/src/test/java/org/apache/arrow/memory/TestCheckAllocator.java

@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.memory;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.io.IOException;
+import java.net.URI;
+import java.net.URL;
+import java.util.Collections;
+import java.util.Enumeration;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+/** Test cases for {@link CheckAllocator}. */
+public class TestCheckAllocator {
+
+  private static final String CORE = "org.apache.arrow.memory.DefaultAllocationManagerFactory";
+  private static final String UNSAFE =
+      "org.apache.arrow.memory.unsafe.DefaultAllocationManagerFactory";
+  private static final String NETTY =
+      "org.apache.arrow.memory.netty.DefaultAllocationManagerFactory";
+
+  @Test
+  public void testCheckClasspath() {
+    // memory-core tests provide their own DefaultAllocationManagerFactory
+    assertEquals(CORE, CheckAllocator.check());
+  }
+
+  @Test
+  public void testCheckIgnoresLocation() {
+    ClassLoader classLoader =
+        classLoaderWith(Map.of(NETTY, "file:/work/arrow-memory-core/memory-netty/classes/"));
+    assertEquals(NETTY, CheckAllocator.check(classLoader));
+  }
+
+  @Test
+  public void testCheckUnsafeBeforeNetty() {
+    ClassLoader classLoader =
+        classLoaderWith(
+            Map.of(
+                NETTY, "file:/lib/arrow-memory-netty/", UNSAFE, "file:/lib/arrow-memory-unsafe/"));
+    assertEquals(UNSAFE, CheckAllocator.check(classLoader));
+  }
+
+  @Test
+  public void testCheckNoAllocator() {
+    assertThrows(RuntimeException.class, () -> CheckAllocator.check(classLoaderWith(Map.of())));
+  }
+
+  /** Returns a class loader that finds each factory class only at the given location. */
+  private static ClassLoader classLoaderWith(Map<String, String> locations) {
+    return new ClassLoader(null) {
+      @Override
+      public Enumeration<URL> getResources(String name) throws IOException {
+        for (Map.Entry<String, String> location : locations.entrySet()) {
+          if (name.equals(location.getKey().replace('.', '/') + ".class")) {
+            return Collections.enumeration(List.of(URI.create(location.getValue() + name).toURL()));
+          }
+        }
+        return Collections.emptyEnumeration();
+      }
+    };
+  }
+}