diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/CheckAllocator.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/CheckAllocator.java index f595858ebf..d7d874894b 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/CheckAllocator.java +++ b/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 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 scanClasspath() { - // LinkedHashSet appropriate here because it preserves insertion order - // during iteration - Set allocatorPathSet = new LinkedHashSet<>(); - try { - ClassLoader allocatorClassLoader = CheckAllocator.class.getClassLoader(); - Enumeration 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 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> 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> first = found.entrySet().iterator().next(); + Allocator allocator = first.getKey(); + Set 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> scanClasspath(@Nullable ClassLoader classLoader) { + // EnumMap iterates in declaration order, which is the order of precedence + Map> found = new EnumMap<>(Allocator.class); + for (Allocator allocator : Allocator.values()) { + try { + Enumeration paths = + classLoader == null + ? ClassLoader.getSystemResources(allocator.resourcePath) + : classLoader.getResources(allocator.resourcePath); + // LinkedHashSet appropriate here because it preserves insertion order + // during iteration + Set 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; } } diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java index 2ecf6722b5..0ad80d6fa5 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java +++ b/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); diff --git a/memory/memory-core/src/test/java/org/apache/arrow/memory/TestCheckAllocator.java b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestCheckAllocator.java new file mode 100644 index 0000000000..07d5534ae0 --- /dev/null +++ b/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 locations) { + return new ClassLoader(null) { + @Override + public Enumeration getResources(String name) throws IOException { + for (Map.Entry location : locations.entrySet()) { + if (name.equals(location.getKey().replace('.', '/') + ".class")) { + return Collections.enumeration(List.of(URI.create(location.getValue() + name).toURL())); + } + } + return Collections.emptyEnumeration(); + } + }; + } +}