diff --git a/java/memory/src/main/java/org/apache/arrow/memory/BaseAllocator.java b/java/memory/src/main/java/org/apache/arrow/memory/BaseAllocator.java index 57ee58b0169..1b9b6b1e9fe 100644 --- a/java/memory/src/main/java/org/apache/arrow/memory/BaseAllocator.java +++ b/java/memory/src/main/java/org/apache/arrow/memory/BaseAllocator.java @@ -763,7 +763,7 @@ abstract static class Config { */ @Value.Default AllocationManager.Factory getAllocationManagerFactory() { - return NettyAllocationManager.FACTORY; + return DefaultAllocationManagerOption.DEFAULT_ALLOCATION_MANAGER_FACTORY; } /** diff --git a/java/memory/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java b/java/memory/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java new file mode 100644 index 00000000000..66f0ec638d1 --- /dev/null +++ b/java/memory/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java @@ -0,0 +1,98 @@ +/* + * 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; + +/** + * A class for choosing the default allocation manager. + */ +public class DefaultAllocationManagerOption { + + /** + * The environmental variable to set the default allocation manager type. + */ + public static final String ALLOCATION_MANAGER_TYPE_ENV_NAME = "ARROW_ALLOCATION_MANAGER_TYPE"; + + /** + * The system property to set the default allocation manager type. + */ + public static final String ALLOCATION_MANAGER_TYPE_PROPERTY_NAME = "arrow.allocation.manager.type"; + + static final org.slf4j.Logger LOGGER = org.slf4j.LoggerFactory.getLogger(DefaultAllocationManagerOption.class); + + /** + * The default allocation manager factory. + */ + public static final AllocationManager.Factory DEFAULT_ALLOCATION_MANAGER_FACTORY = + getDefaultAllocationManagerFactory(); + + /** + * The allocation manager type. + */ + public enum AllocationManagerType { + /** + * Netty based allocation manager. + */ + Netty, + + /** + * Unsafe based allocation manager. + */ + Unsafe, + + /** + * Unknown type. + */ + Unknown, + } + + static AllocationManagerType getDefaultAllocationManagerType() { + AllocationManagerType ret = AllocationManagerType.Unknown; + + try { + String envValue = System.getenv(ALLOCATION_MANAGER_TYPE_ENV_NAME); + ret = AllocationManagerType.valueOf(envValue); + } catch (IllegalArgumentException | NullPointerException e) { + // ignore the exception, and make the allocation manager type remain unchanged + } + + // system property takes precedence + try { + String propValue = System.getProperty(ALLOCATION_MANAGER_TYPE_PROPERTY_NAME); + ret = AllocationManagerType.valueOf(propValue); + } catch (IllegalArgumentException | NullPointerException e) { + // ignore the exception, and make the allocation manager type remain unchanged + } + return ret; + } + + static AllocationManager.Factory getDefaultAllocationManagerFactory() { + AllocationManagerType type = getDefaultAllocationManagerType(); + + switch (type) { + case Netty: + return NettyAllocationManager.FACTORY; + case Unsafe: + return UnsafeAllocationManager.FACTORY; + case Unknown: + LOGGER.info("allocation manager type not specified, using netty as the default type"); + return NettyAllocationManager.FACTORY; + default: + throw new IllegalStateException("Unknown allocation manager type: " + type); + } + } +} diff --git a/java/memory/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java b/java/memory/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java new file mode 100644 index 00000000000..d450444f811 --- /dev/null +++ b/java/memory/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java @@ -0,0 +1,65 @@ +/* + * 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 org.apache.arrow.memory.util.MemoryUtil; + +/** + * Allocation manager based on unsafe API. + */ +public final class UnsafeAllocationManager extends AllocationManager { + + public static final Factory FACTORY = new Factory(); + + private final long allocatedSize; + + private final long allocatedAddress; + + UnsafeAllocationManager(BaseAllocator accountingAllocator, long requestedSize) { + super(accountingAllocator); + allocatedAddress = MemoryUtil.UNSAFE.allocateMemory(requestedSize); + allocatedSize = requestedSize; + } + + @Override + public long getSize() { + return allocatedSize; + } + + @Override + protected long memoryAddress() { + return allocatedAddress; + } + + @Override + protected void release0() { + MemoryUtil.UNSAFE.freeMemory(allocatedAddress); + } + + /** + * Factory for creating {@link UnsafeAllocationManager}. + */ + public static class Factory implements AllocationManager.Factory { + private Factory() {} + + @Override + public AllocationManager create(BaseAllocator accountingAllocator, long size) { + return new UnsafeAllocationManager(accountingAllocator, size); + } + } +} diff --git a/java/memory/src/test/java/org/apache/arrow/memory/TestAllocationManager.java b/java/memory/src/test/java/org/apache/arrow/memory/TestAllocationManager.java new file mode 100644 index 00000000000..3f68a6f3737 --- /dev/null +++ b/java/memory/src/test/java/org/apache/arrow/memory/TestAllocationManager.java @@ -0,0 +1,52 @@ +/* + * 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.Assert.assertEquals; + +import org.junit.Test; + +/** + * Test cases for {@link AllocationManager}. + */ +public class TestAllocationManager { + + @Test + public void testAllocationManagerType() { + // test netty allocation manager type + System.setProperty( + DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME, "Netty"); + DefaultAllocationManagerOption.AllocationManagerType mgrType = + DefaultAllocationManagerOption.getDefaultAllocationManagerType(); + + assertEquals(DefaultAllocationManagerOption.AllocationManagerType.Netty, mgrType); + + // test unsafe allocation manager type + System.setProperty( + DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME, "Unsafe"); + mgrType = DefaultAllocationManagerOption.getDefaultAllocationManagerType(); + + assertEquals(DefaultAllocationManagerOption.AllocationManagerType.Unsafe, mgrType); + + // test unknown allocation manager type + System.clearProperty(DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME); + mgrType = DefaultAllocationManagerOption.getDefaultAllocationManagerType(); + + assertEquals(DefaultAllocationManagerOption.AllocationManagerType.Unknown, mgrType); + } +} diff --git a/java/memory/src/test/java/org/apache/arrow/memory/TestUnsafeAllocationManager.java b/java/memory/src/test/java/org/apache/arrow/memory/TestUnsafeAllocationManager.java new file mode 100644 index 00000000000..af4f9aa3cbb --- /dev/null +++ b/java/memory/src/test/java/org/apache/arrow/memory/TestUnsafeAllocationManager.java @@ -0,0 +1,70 @@ +/* + * 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.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; + +/** + * Test cases for {@link UnsafeAllocationManager}. + */ +public class TestUnsafeAllocationManager { + + private BaseAllocator createUnsafeAllocator() { + return new RootAllocator(BaseAllocator.configBuilder() + .allocationManagerFactory((accountingAllocator, requestedSize) -> + new UnsafeAllocationManager( + accountingAllocator, requestedSize)).build()); + } + + private void readWriteArrowBuf(ArrowBuf buffer) { + // write buffer + for (long i = 0; i < buffer.capacity() / 8; i++) { + buffer.setLong(i * 8, i); + } + + // read buffer + for (long i = 0; i < buffer.capacity() / 8; i++) { + long val = buffer.getLong(i * 8); + assertEquals(i, val); + } + } + + /** + * Test the memory allocation for {@link UnsafeAllocationManager}. + */ + @Test + public void testBufferAllocation() { + final long bufSize = 4096L; + try (BaseAllocator allocator = createUnsafeAllocator(); + ArrowBuf buffer = allocator.buffer(bufSize)) { + assertTrue(buffer.getReferenceManager() instanceof BufferLedger); + BufferLedger bufferLedger = (BufferLedger) buffer.getReferenceManager(); + + // make sure we are using unsafe allocation manager + AllocationManager allocMgr = bufferLedger.getAllocationManager(); + assertTrue(allocMgr instanceof UnsafeAllocationManager); + UnsafeAllocationManager unsafeMgr = (UnsafeAllocationManager) allocMgr; + + assertEquals(bufSize, unsafeMgr.getSize()); + readWriteArrowBuf(buffer); + } + } +}