Add unified native allocator framework with elastic rebalancing (#21703)

Signed-off-by: Bukhtawar Khan <bukhtawa@amazon.com>

Author: Bukhtawar

libs/arrow-spi/build.gradle

@@ -0,0 +1,23 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+apply plugin: 'opensearch.publish'
+
+dependencies {
+    api project(':libs:opensearch-core')
+    api project(':libs:opensearch-common')
+    testImplementation project(':test:framework')
+}
+
+tasks.named('forbiddenApisMain').configure {
+    replaceSignatureFiles 'jdk-signatures'
+}
+
+tasks.named('forbiddenApisTest').configure {
+    replaceSignatureFiles 'jdk-signatures'
+}

libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocator.java

@@ -0,0 +1,93 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.arrow.spi;
+
+import java.io.Closeable;
+
+/**
+ * Arrow-agnostic interface for a hierarchical native memory allocator.
+ *
+ * <p>The implementation (backed by Arrow's {@code RootAllocator}) is provided by
+ * a plugin. The SPI allows other subsystems to interact with the allocator
+ * without depending on Arrow classes.
+ *
+ * <p>Plugins that need Arrow allocators obtain the implementation via
+ * service lookup or plugin extension and call {@link #getOrCreatePool} to
+ * register their pool.
+ *
+ * @opensearch.api
+ */
+public interface NativeAllocator extends Closeable {
+
+    /**
+     * Returns the named pool, creating it on first access with the given limit.
+     * Subsequent calls with the same name return the same pool (first-call limit wins).
+     *
+     * @param poolName logical pool name (e.g., "query", "flight")
+     * @param limit maximum bytes this pool can allocate in aggregate
+     * @return an opaque pool handle
+     */
+    PoolHandle getOrCreatePool(String poolName, long limit);
+
+    /**
+     * Updates the limit of an existing pool.
+     *
+     * @param poolName logical pool name
+     * @param newLimit new maximum bytes for the pool
+     */
+    void setPoolLimit(String poolName, long newLimit);
+
+    /**
+     * Sets the root-level memory limit for the entire allocator.
+     *
+     * @param limit new maximum bytes for the root allocator
+     */
+    void setRootLimit(long limit);
+
+    /**
+     * Collects a point-in-time stats snapshot across all pools.
+     */
+    NativeAllocatorPoolStats stats();
+
+    /**
+     * Opaque handle to a memory pool. Plugins downcast to the concrete type
+     * (e.g., Arrow's {@code BufferAllocator}) in the implementation layer.
+     */
+    interface PoolHandle {
+
+        /**
+         * Creates a named child allocation within this pool.
+         *
+         * @param childName name for debugging
+         * @param childLimit maximum bytes for the child
+         * @return an opaque child handle (downcast to BufferAllocator in Arrow impl)
+         */
+        PoolHandle newChild(String childName, long childLimit);
+
+        /**
+         * Returns the current allocated bytes for this pool/child.
+         */
+        long allocatedBytes();
+
+        /**
+         * Returns the peak memory allocation.
+         */
+        long peakBytes();
+
+        /**
+         * Returns the configured limit.
+         */
+        long limit();
+
+        /**
+         * Releases this allocation handle.
+         */
+        void close();
+    }
+}

libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfig.java

@@ -0,0 +1,47 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.arrow.spi;
+
+/**
+ * Pool name constants and setting keys for native allocator pools.
+ *
+ * <p>Each pool has a min (guaranteed reservation) and max (burst limit).
+ * The rebalancer ensures every pool can always allocate up to its min,
+ * and distributes unused capacity up to each pool's max.
+ *
+ * <p>Limits are provided via {@code opensearch.yml} or the cluster settings API.
+ * The setting keys follow the pattern
+ * {@code native.allocator.pool.<name>.min} and {@code native.allocator.pool.<name>.max}.
+ *
+ * <p>This class is Arrow-agnostic — it defines the logical pool topology
+ * without referencing any Arrow classes.
+ *
+ * @opensearch.api
+ */
+public final class NativeAllocatorPoolConfig {
+
+    /** Pool name for Arrow Flight RPC memory. */
+    public static final String POOL_FLIGHT = "flight";
+    /** Pool name for ingest pipeline memory. */
+    public static final String POOL_INGEST = "ingest";
+
+    /** Setting key for the root allocator limit. */
+    public static final String SETTING_ROOT_LIMIT = "native.allocator.root.limit";
+
+    /** Setting key for the Flight pool minimum. */
+    public static final String SETTING_FLIGHT_MIN = "native.allocator.pool.flight.min";
+    /** Setting key for the Flight pool maximum. */
+    public static final String SETTING_FLIGHT_MAX = "native.allocator.pool.flight.max";
+    /** Setting key for the ingest pool minimum. */
+    public static final String SETTING_INGEST_MIN = "native.allocator.pool.ingest.min";
+    /** Setting key for the ingest pool maximum. */
+    public static final String SETTING_INGEST_MAX = "native.allocator.pool.ingest.max";
+
+    private NativeAllocatorPoolConfig() {}
+}

libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolStats.java

@@ -0,0 +1,208 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.arrow.spi;
+
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.common.unit.ByteSizeValue;
+import org.opensearch.core.xcontent.ToXContentFragment;
+import org.opensearch.core.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Point-in-time snapshot of native memory pool statistics for a node.
+ *
+ * <p>Arrow-agnostic — describes pool utilization without referencing
+ * Arrow classes. The concrete implementation (backed by Arrow's
+ * {@code BufferAllocator}) lives in the plugin that owns the pools.
+ *
+ * @opensearch.api
+ */
+public class NativeAllocatorPoolStats implements Writeable, ToXContentFragment {
+
+    private final long rootAllocatedBytes;
+    private final long rootPeakBytes;
+    private final long rootLimitBytes;
+    private final List<PoolStats> pools;
+
+    /**
+     * Creates a new stats snapshot from the given values.
+     *
+     * @param rootAllocatedBytes current bytes allocated by the root
+     * @param rootPeakBytes peak bytes allocated by the root
+     * @param rootLimitBytes configured root limit
+     * @param pools per-pool stats
+     */
+    public NativeAllocatorPoolStats(long rootAllocatedBytes, long rootPeakBytes, long rootLimitBytes, List<PoolStats> pools) {
+        this.rootAllocatedBytes = rootAllocatedBytes;
+        this.rootPeakBytes = rootPeakBytes;
+        this.rootLimitBytes = rootLimitBytes;
+        this.pools = Collections.unmodifiableList(pools);
+    }
+
+    /**
+     * Deserializes from stream.
+     *
+     * @param in the stream input
+     */
+    public NativeAllocatorPoolStats(StreamInput in) throws IOException {
+        this.rootAllocatedBytes = in.readVLong();
+        this.rootPeakBytes = in.readVLong();
+        this.rootLimitBytes = in.readVLong();
+        int count = in.readVInt();
+        List<PoolStats> list = new ArrayList<>(count);
+        for (int i = 0; i < count; i++) {
+            list.add(new PoolStats(in));
+        }
+        this.pools = Collections.unmodifiableList(list);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeVLong(rootAllocatedBytes);
+        out.writeVLong(rootPeakBytes);
+        out.writeVLong(rootLimitBytes);
+        out.writeVInt(pools.size());
+        for (PoolStats pool : pools) {
+            pool.writeTo(out);
+        }
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject("native_allocator");
+
+        builder.startObject("root");
+        builder.humanReadableField("allocated_bytes", "allocated", new ByteSizeValue(rootAllocatedBytes));
+        builder.humanReadableField("peak_bytes", "peak", new ByteSizeValue(rootPeakBytes));
+        builder.humanReadableField("limit_bytes", "limit", new ByteSizeValue(rootLimitBytes));
+        builder.endObject();
+
+        builder.startObject("pools");
+        for (PoolStats pool : pools) {
+            pool.toXContent(builder, params);
+        }
+        builder.endObject();
+
+        builder.endObject();
+        return builder;
+    }
+
+    /** Returns the root allocator's currently allocated bytes. */
+    public long getRootAllocatedBytes() {
+        return rootAllocatedBytes;
+    }
+
+    /** Returns the root allocator's peak allocated bytes. */
+    public long getRootPeakBytes() {
+        return rootPeakBytes;
+    }
+
+    /** Returns the root allocator's configured limit in bytes. */
+    public long getRootLimitBytes() {
+        return rootLimitBytes;
+    }
+
+    /** Returns the per-pool statistics. */
+    public List<PoolStats> getPools() {
+        return pools;
+    }
+
+    /**
+     * Per-pool statistics snapshot.
+     */
+    public static class PoolStats implements Writeable, ToXContentFragment {
+
+        private final String name;
+        private final long allocatedBytes;
+        private final long peakBytes;
+        private final long limitBytes;
+        private final int childCount;
+
+        /**
+         * Creates a new pool stats snapshot.
+         *
+         * @param name pool name
+         * @param allocatedBytes current allocated bytes
+         * @param peakBytes peak allocated bytes
+         * @param limitBytes configured limit
+         * @param childCount number of child allocators
+         */
+        public PoolStats(String name, long allocatedBytes, long peakBytes, long limitBytes, int childCount) {
+            this.name = name;
+            this.allocatedBytes = allocatedBytes;
+            this.peakBytes = peakBytes;
+            this.limitBytes = limitBytes;
+            this.childCount = childCount;
+        }
+
+        /**
+         * Deserializes from stream.
+         *
+         * @param in the stream input
+         */
+        public PoolStats(StreamInput in) throws IOException {
+            this.name = in.readString();
+            this.allocatedBytes = in.readVLong();
+            this.peakBytes = in.readVLong();
+            this.limitBytes = in.readVLong();
+            this.childCount = in.readVInt();
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeString(name);
+            out.writeVLong(allocatedBytes);
+            out.writeVLong(peakBytes);
+            out.writeVLong(limitBytes);
+            out.writeVInt(childCount);
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            builder.startObject(name);
+            builder.humanReadableField("allocated_bytes", "allocated", new ByteSizeValue(allocatedBytes));
+            builder.humanReadableField("peak_bytes", "peak", new ByteSizeValue(peakBytes));
+            builder.humanReadableField("limit_bytes", "limit", new ByteSizeValue(limitBytes));
+            builder.field("child_count", childCount);
+            builder.endObject();
+            return builder;
+        }
+
+        /** Returns the pool name. */
+        public String getName() {
+            return name;
+        }
+
+        /** Returns the currently allocated bytes. */
+        public long getAllocatedBytes() {
+            return allocatedBytes;
+        }
+
+        /** Returns the peak allocated bytes. */
+        public long getPeakBytes() {
+            return peakBytes;
+        }
+
+        /** Returns the configured limit in bytes. */
+        public long getLimitBytes() {
+            return limitBytes;
+        }
+
+        /** Returns the number of child allocators. */
+        public int getChildCount() {
+            return childCount;
+        }
+    }
+}

libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/package-info.java

@@ -0,0 +1,12 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+/**
+ * Service Provider Interface for Arrow-agnostic native memory allocation.
+ */
+package org.opensearch.arrow.spi;