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

chaokunyang pushed a commit to branch releases-0.12
in repository https://gitbox.apache.org/repos/asf/fory.git

commit 4ed0c872edd9b457d885164efa3563d7cb793878
Author: adri <[email protected]>
AuthorDate: Thu Aug 14 12:57:35 2025 +0200

    feat(memory): add customizable MemoryAllocator interface  (#2467)
    
    ## What does this PR do?
    
    This PR introduces a new `MemoryAllocator` interface that allows
    customisation of memory allocation strategies in `MemoryBuffer`. This
    enables users to implement custom allocation policies.
    
    ## Related issues
    - Closes https://github.com/apache/fory/issues/2459
    - https://github.com/apache/fory/pull/2457
    - Closes https://github.com/apache/fory/issues/2350.
    
    ## Does this PR introduce any user-facing change?
    
    <!--
    If any user-facing interface changes, please [open an
    issue](https://github.com/apache/fory/issues/new/choose) describing the
    need to do so and update the document if necessary.
    -->
    
    - [ ] Does this PR introduce any public API change?
    - [ ] Does this PR introduce any binary protocol compatibility change?
    
    ---------
    
    Co-authored-by: Shawn Yang <[email protected]>
---
 docs/guide/java_serialization_guide.md             |  78 +++++++++++
 .../org/apache/fory/memory/MemoryAllocator.java    |  41 ++++++
 .../java/org/apache/fory/memory/MemoryBuffer.java  |  72 ++++++++--
 .../apache/fory/memory/MemoryAllocatorTest.java    | 153 +++++++++++++++++++++
 4 files changed, 331 insertions(+), 13 deletions(-)

diff --git a/docs/guide/java_serialization_guide.md 
b/docs/guide/java_serialization_guide.md
index a4f7da2b8..7f1f7a2af 100644
--- a/docs/guide/java_serialization_guide.md
+++ b/docs/guide/java_serialization_guide.md
@@ -1089,6 +1089,84 @@ Note that when implementing custom map or collection 
serializers:
 
 Besides registering serializes, one can also implement 
`java.io.Externalizable` for a class to customize serialization logic, such 
type will be serialized by fory `ExternalizableSerializer`.
 
+### Memory Allocation Customization
+
+Fory provides a `MemoryAllocator` interface that allows you to customize how 
memory buffers are allocated and grown during serialization operations. This 
can be useful for performance optimization, memory pooling, or debugging memory 
usage.
+
+#### MemoryAllocator Interface
+
+The `MemoryAllocator` interface defines two key methods:
+
+```java
+public interface MemoryAllocator {
+  /**
+   * Allocates a new MemoryBuffer with the specified initial capacity.
+   */
+  MemoryBuffer allocate(int initialCapacity);
+
+  /**
+   * Grows an existing buffer to accommodate the new capacity.
+   * The implementation must grow the buffer in-place by modifying
+   * the existing buffer instance.
+   */
+  MemoryBuffer grow(MemoryBuffer buffer, int newCapacity);
+}
+```
+
+#### Using Custom Memory Allocators
+
+You can set a global memory allocator that will be used by all `MemoryBuffer` 
instances:
+
+```java
+// Create a custom allocator
+MemoryAllocator customAllocator = new MemoryAllocator() {
+  @Override
+  public MemoryBuffer allocate(int initialCapacity) {
+    // Add extra capacity for debugging or pooling
+    return MemoryBuffer.fromByteArray(new byte[initialCapacity + 100]);
+  }
+
+  @Override
+  public MemoryBuffer grow(MemoryBuffer buffer, int newCapacity) {
+    if (newCapacity <= buffer.size()) {
+      return buffer;
+    }
+
+    // Custom growth strategy - add 100% extra capacity
+    int newSize = (int) (newCapacity * 2);
+    byte[] data = new byte[newSize];
+    buffer.copyToUnsafe(0, data, Platform.BYTE_ARRAY_OFFSET, buffer.size());
+    buffer.initHeapBuffer(data, 0, data.length);
+    return buffer;
+  }
+};
+
+// Set the custom allocator globally
+MemoryBuffer.setGlobalAllocator(customAllocator);
+
+// All subsequent MemoryBuffer allocations will use your custom allocator
+Fory fory = Fory.builder().withLanguage(Language.JAVA).build();
+byte[] bytes = fory.serialize(someObject); // Uses custom allocator
+```
+
+#### Default Memory Allocator Behavior
+
+The default allocator uses the following growth strategy:
+
+- For buffers smaller than `BUFFER_GROW_STEP_THRESHOLD` (100MB): multiply 
capacity by 2
+- For larger buffers: multiply capacity by 1.5 (capped at `Integer.MAX_VALUE - 
8`)
+
+This provides a balance between avoiding frequent reallocations and preventing 
excessive memory usage.
+
+#### Use Cases
+
+Custom memory allocators are useful for:
+
+- **Memory Pooling**: Reuse allocated buffers to reduce GC pressure
+- **Performance Tuning**: Use different growth strategies based on your 
workload
+- **Debugging**: Add logging or tracking to monitor memory usage
+- **Off-heap Memory**: Integrate with off-heap memory management systems
+
 ### Security & Class Registration
 
 `ForyBuilder#requireClassRegistration` can be used to disable class 
registration, this will allow to deserialize objects
diff --git 
a/java/fory-core/src/main/java/org/apache/fory/memory/MemoryAllocator.java 
b/java/fory-core/src/main/java/org/apache/fory/memory/MemoryAllocator.java
new file mode 100644
index 000000000..ebbc6098c
--- /dev/null
+++ b/java/fory-core/src/main/java/org/apache/fory/memory/MemoryAllocator.java
@@ -0,0 +1,41 @@
+/*
+ * 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.fory.memory;
+
+/** Interface for customizing memory allocation strategies in MemoryBuffer. */
+public interface MemoryAllocator {
+  /**
+   * Allocates a new MemoryBuffer with the specified initial capacity.
+   *
+   * @param initialCapacity the initial capacity for the buffer
+   * @return a new MemoryBuffer instance
+   */
+  MemoryBuffer allocate(int initialCapacity);
+
+  /**
+   * Grows an existing buffer to accommodate the new capacity. The 
implementation must grow the
+   * buffer in-place by modifying the existing buffer instance.
+   *
+   * @param buffer the existing buffer to grow
+   * @param newCapacity the required new capacity
+   * @return the same MemoryBuffer instance with at least the new capacity
+   */
+  MemoryBuffer grow(MemoryBuffer buffer, int newCapacity);
+}
diff --git 
a/java/fory-core/src/main/java/org/apache/fory/memory/MemoryBuffer.java 
b/java/fory-core/src/main/java/org/apache/fory/memory/MemoryBuffer.java
index fe7704ca9..97f1fa0d4 100644
--- a/java/fory-core/src/main/java/org/apache/fory/memory/MemoryBuffer.java
+++ b/java/fory-core/src/main/java/org/apache/fory/memory/MemoryBuffer.java
@@ -64,6 +64,9 @@ public final class MemoryBuffer {
   private static final Unsafe UNSAFE = Platform.UNSAFE;
   private static final boolean LITTLE_ENDIAN = (ByteOrder.nativeOrder() == 
ByteOrder.LITTLE_ENDIAN);
 
+  // Global allocator instance that can be customized
+  private static volatile MemoryAllocator globalAllocator = new 
DefaultMemoryAllocator();
+
   // If the data in on the heap, `heapMemory` will be non-null, and its' the 
object relative to
   // which we access the memory.
   // If we have this buffer, we must never void this reference, or the memory 
buffer will point
@@ -1233,27 +1236,17 @@ public final class MemoryBuffer {
   public void grow(int neededSize) {
     int length = writerIndex + neededSize;
     if (length > size) {
-      growBuffer(length);
+      globalAllocator.grow(this, length);
     }
   }
 
   /** For off-heap buffer, this will make a heap buffer internally. */
   public void ensure(int length) {
     if (length > size) {
-      growBuffer(length);
+      globalAllocator.grow(this, length);
     }
   }
 
-  private void growBuffer(int length) {
-    int newSize =
-        length < BUFFER_GROW_STEP_THRESHOLD
-            ? length << 2
-            : (int) Math.min(length * 1.5d, Integer.MAX_VALUE - 8);
-    byte[] data = new byte[newSize];
-    copyToUnsafe(0, data, Platform.BYTE_ARRAY_OFFSET, size());
-    initHeapBuffer(data, 0, data.length);
-  }
-
   // -------------------------------------------------------------------------
   //                          Read Methods
   // -------------------------------------------------------------------------
@@ -2607,6 +2600,59 @@ public final class MemoryBuffer {
         + '}';
   }
 
+  // ------------------------------------------------------------------------
+  // Memory Allocator Support
+  // ------------------------------------------------------------------------
+
+  /** Default memory allocator that uses the original heap-based allocation 
strategy. */
+  private static final class DefaultMemoryAllocator implements MemoryAllocator 
{
+    @Override
+    public MemoryBuffer allocate(int initialSize) {
+      return fromByteArray(new byte[initialSize]);
+    }
+
+    @Override
+    public MemoryBuffer grow(MemoryBuffer buffer, int newCapacity) {
+      if (newCapacity <= buffer.size()) {
+        return buffer;
+      }
+
+      int newSize =
+          newCapacity < BUFFER_GROW_STEP_THRESHOLD
+              ? newCapacity << 1
+              : (int) Math.min(newCapacity * 1.5d, Integer.MAX_VALUE - 8);
+
+      byte[] data = new byte[newSize];
+      buffer.copyToUnsafe(0, data, Platform.BYTE_ARRAY_OFFSET, buffer.size());
+      buffer.initHeapBuffer(data, 0, data.length);
+
+      return buffer;
+    }
+  }
+
+  /**
+   * Sets the global memory allocator. This affects all new MemoryBuffer 
allocations and growth
+   * operations.
+   *
+   * @param allocator the new global allocator to use
+   * @throws NullPointerException if allocator is null
+   */
+  public static void setGlobalAllocator(MemoryAllocator allocator) {
+    if (allocator == null) {
+      throw new NullPointerException("Memory allocator cannot be null");
+    }
+    globalAllocator = allocator;
+  }
+
+  /**
+   * Gets the current global memory allocator.
+   *
+   * @return the current global allocator
+   */
+  public static MemoryAllocator getGlobalAllocator() {
+    return globalAllocator;
+  }
+
   /** Point this buffer to a new byte array. */
   public void pointTo(byte[] buffer, int offset, int length) {
     initHeapBuffer(buffer, offset, length);
@@ -2663,6 +2709,6 @@ public final class MemoryBuffer {
    * enough.
    */
   public static MemoryBuffer newHeapBuffer(int initialSize) {
-    return fromByteArray(new byte[initialSize]);
+    return globalAllocator.allocate(initialSize);
   }
 }
diff --git 
a/java/fory-core/src/test/java/org/apache/fory/memory/MemoryAllocatorTest.java 
b/java/fory-core/src/test/java/org/apache/fory/memory/MemoryAllocatorTest.java
new file mode 100644
index 000000000..68c3e613d
--- /dev/null
+++ 
b/java/fory-core/src/test/java/org/apache/fory/memory/MemoryAllocatorTest.java
@@ -0,0 +1,153 @@
+/*
+ * 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.fory.memory;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertTrue;
+
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+public class MemoryAllocatorTest {
+
+  private MemoryAllocator originalAllocator;
+
+  @BeforeMethod
+  public void setUp() {
+    // Save the original allocator before each test
+    originalAllocator = MemoryBuffer.getGlobalAllocator();
+  }
+
+  @AfterMethod
+  public void tearDown() {
+    // Restore the original allocator after each test
+    MemoryBuffer.setGlobalAllocator(originalAllocator);
+  }
+
+  @Test
+  public void testDefaultMemoryAllocator() {
+    MemoryAllocator defaultAllocator = MemoryBuffer.getGlobalAllocator();
+
+    MemoryBuffer buffer = defaultAllocator.allocate(100);
+    assertEquals(buffer.size(), 100);
+    assertFalse(buffer.isOffHeap());
+
+    // Test growth below BUFFER_GROW_STEP_THRESHOLD (should multiply by 2)
+    defaultAllocator.grow(buffer, 200);
+    assertEquals(buffer.size(), 200 << 1);
+
+    // Test growth above BUFFER_GROW_STEP_THRESHOLD
+    buffer = defaultAllocator.allocate(100);
+    int largeCapacity = MemoryBuffer.BUFFER_GROW_STEP_THRESHOLD + 1000;
+    defaultAllocator.grow(buffer, largeCapacity);
+    int expectedSize = (int) Math.min(largeCapacity * 1.5d, Integer.MAX_VALUE 
- 8);
+    assertEquals(buffer.size(), expectedSize);
+  }
+
+  @Test
+  public void testDefaultMemoryAllocatorDataPreservation() {
+    MemoryAllocator defaultAllocator = MemoryBuffer.getGlobalAllocator();
+    MemoryBuffer buffer = defaultAllocator.allocate(100);
+
+    // Write some test data
+    byte[] testData = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
+    buffer.writeBytes(testData);
+    buffer.writeInt32(42);
+    buffer.writeInt64(123456789L);
+
+    int writerIndexBeforeGrowth = buffer.writerIndex();
+
+    // Grow the buffer
+    defaultAllocator.grow(buffer, 500);
+
+    // Verify data is preserved
+    buffer.readerIndex(0);
+    byte[] readData = new byte[testData.length];
+    buffer.readBytes(readData);
+    for (int i = 0; i < testData.length; i++) {
+      assertEquals(readData[i], testData[i]);
+    }
+
+    assertEquals(buffer.readInt32(), 42);
+    assertEquals(buffer.readInt64(), 123456789L);
+    assertEquals(buffer.writerIndex(), writerIndexBeforeGrowth);
+  }
+
+  @Test
+  public void testDefaultMemoryAllocatorGrowthSameInstance() {
+    MemoryAllocator defaultAllocator = MemoryBuffer.getGlobalAllocator();
+    MemoryBuffer buffer = defaultAllocator.allocate(100);
+
+    // Growth should return the same instance
+    MemoryBuffer grownBuffer = defaultAllocator.grow(buffer, 200);
+    assertSame(buffer, grownBuffer);
+  }
+
+  @Test
+  public void testCustomAllocator() {
+    // Create a custom allocator that adds a marker
+    MemoryAllocator customAllocator =
+        new MemoryAllocator() {
+          @Override
+          public MemoryBuffer allocate(int initialCapacity) {
+            // Use larger capacity as a marker
+            return MemoryBuffer.fromByteArray(new byte[initialCapacity + 10]);
+          }
+
+          @Override
+          public MemoryBuffer grow(MemoryBuffer buffer, int newCapacity) {
+            if (newCapacity <= buffer.size()) {
+              return buffer;
+            }
+
+            // Use default grow logic but with custom marker
+            int newSize = newCapacity + 10; // Add 10 as marker
+            byte[] data = new byte[newSize];
+            buffer.copyToUnsafe(0, data, Platform.BYTE_ARRAY_OFFSET, 
buffer.size());
+            buffer.initHeapBuffer(data, 0, data.length);
+            return buffer;
+          }
+        };
+
+    // Set the custom allocator
+    MemoryBuffer.setGlobalAllocator(customAllocator);
+    assertSame(MemoryBuffer.getGlobalAllocator(), customAllocator);
+
+    // Test allocation
+    MemoryBuffer buffer = MemoryBuffer.newHeapBuffer(100);
+    assertEquals(buffer.size(), 110); // 100 + 10 marker
+
+    // Test growth
+    buffer.writerIndex(50);
+    buffer.readerIndex(10);
+    buffer.ensure(200); // This should trigger growth
+    assertEquals(buffer.writerIndex(), 50);
+    assertEquals(buffer.readerIndex(), 10);
+    assertTrue(buffer.size() >= 210); // Should be at least 200 + 10 marker
+  }
+
+  @Test(expectedExceptions = NullPointerException.class)
+  public void testSetNullAllocator() {
+    MemoryBuffer.setGlobalAllocator(null);
+  }
+}


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to