From 2b0abac3da9f4f98080cad9d1ee83c918c8de7c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Fri, 13 Nov 2026 05:29:10 +0100 Subject: [PATCH 1/2] GH-1469: Infer the allocation manager from the factory found on the classpath CheckAllocator now looks up each known DefaultAllocationManagerFactory and returns the one it found, instead of guessing from the location name. It warns when several allocation manager types are on the classpath, and logs the inferred type (Netty/Unsafe) with its location on a single line. Closes #1469. --- .../apache/arrow/memory/CheckAllocator.java | 145 ++++++++++-------- .../DefaultAllocationManagerOption.java | 6 +- .../arrow/memory/TestCheckAllocator.java | 80 ++++++++++ 3 files changed, 158 insertions(+), 73 deletions(-) create mode 100644 memory/memory-core/src/test/java/org/apache/arrow/memory/TestCheckAllocator.java 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..34e7cf2e1b --- /dev/null +++ b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestCheckAllocator.java @@ -0,0 +1,80 @@ +/* + * 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(); + } + }; + } +} From 2f5ee59ab74b088b230913de7472bf52ec087925 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Fri, 13 Nov 2026 07:19:47 +0100 Subject: [PATCH 2/2] GH-1469: Fix Spotless formatting in TestCheckAllocator --- .../test/java/org/apache/arrow/memory/TestCheckAllocator.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 index 34e7cf2e1b..07d5534ae0 100644 --- 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 @@ -54,7 +54,8 @@ public void testCheckIgnoresLocation() { public void testCheckUnsafeBeforeNetty() { ClassLoader classLoader = classLoaderWith( - Map.of(NETTY, "file:/lib/arrow-memory-netty/", UNSAFE, "file:/lib/arrow-memory-unsafe/")); + Map.of( + NETTY, "file:/lib/arrow-memory-netty/", UNSAFE, "file:/lib/arrow-memory-unsafe/")); assertEquals(UNSAFE, CheckAllocator.check(classLoader)); }