[arrow] branch master updated: ARROW-8481: [Java] Provide an allocation manager based on Unsafe API

[em...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

emkornfield pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/master by this push:
     new 3aa3339  ARROW-8481: [Java] Provide an allocation manager based on Unsafe API
3aa3339 is described below

commit 3aa33392e7ca11c48c09a963d165275b26280157
Author: liyafan82 <fa...@foxmail.com>
AuthorDate: Thu May 14 21:24:48 2020 -0700

    ARROW-8481: [Java] Provide an allocation manager based on Unsafe API
    
    This is in response to the discussion in https://github.com/apache/arrow/pull/6323#issuecomment-614195070
    
    In this issue, we provide an allocation manager that is capable of allocation large (> 2GB) buffers. In addition, it does not depend on the netty library, which is aligning with the general trend of removing netty dependencies. In the future, we are going to make it the default allocation manager.
    
    Closes #6956 from liyafan82/fly_0416_unsf
    
    Lead-authored-by: liyafan82 <fa...@foxmail.com>
    Co-authored-by: emkornfield <em...@gmail.com>
    Signed-off-by: Micah Kornfield <em...@gmail.com>
---
 .../org/apache/arrow/memory/BaseAllocator.java     |  2 +-
 .../memory/DefaultAllocationManagerOption.java     | 98 ++++++++++++++++++++++
 .../arrow/memory/UnsafeAllocationManager.java      | 65 ++++++++++++++
 .../apache/arrow/memory/TestAllocationManager.java | 52 ++++++++++++
 .../arrow/memory/TestUnsafeAllocationManager.java  | 70 ++++++++++++++++
 5 files changed, 286 insertions(+), 1 deletion(-)

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 57ee58b..1b9b6b1 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 @@ public abstract class BaseAllocator extends Accountant implements BufferAllocato
      */
     @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 0000000..66f0ec6
--- /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 0000000..d450444
--- /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 0000000..3f68a6f
--- /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 0000000..af4f9aa
--- /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);
+    }
+  }
+}