Skip to content

Commit cf2bcca

Browse files
gaurav-amzbkhishor
authored andcommitted
Native allocator: dynamic settings, query/datafusion pools, plugin nodeStats (opensearch-project#21732)
* Make pool min/max settings dynamic and grouped-validate cross-setting invariants on every cluster-state update (sum of pool mins <= root, per-pool min <= max). FLIGHT_MIN/INGEST_MIN defaults are 0L (the prior Long.MAX_VALUE defaults caused the new validator to reject any non-MAX root). setPoolMin now updates the live BufferAllocator.setLimit so the Dynamic property has observable effect even when the rebalancer is off. * Derive ROOT_LIMIT_SETTING default from the AC node-native-memory budget (limit minus buffer-percent) so the framework respects the same budget AC throttles on, with a Long.MAX_VALUE fallback when AC is unconfigured. * Add QUERY pool alongside FLIGHT and INGEST. Pools init at min when the rebalancer is enabled; otherwise at max so non-rebalanced nodes can still allocate. * Wire arrow-flight-rpc to the FLIGHT pool, parquet-data-format to the INGEST pool, and analytics-engine to the QUERY pool via the framework's allocator service. Hard-fail if the framework plugin is missing — silently skipping the wire-up is the silent-misconfiguration class of bug we want to prevent. * Cleanup ad-hoc allocator fallbacks in parquet-data-format / analytics-engine so all Arrow consumers go through the unified pool hierarchy. * Rebalancer now distributes headroom across all pools (not only those with current allocation > 0). Avoids the dead-pool corner case where a pool with min=0 starts at limit=0, can never make a first allocation, and never receives a bonus. Signed-off-by: Gaurav Singh <snghsvn@amazon.com>
1 parent 050f2bf commit cf2bcca

76 files changed

Lines changed: 2806 additions & 398 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,9 @@ public interface NativeAllocator extends Closeable {
3636
PoolHandle getOrCreatePool(String poolName, long limit);
3737

3838
/**
39-
* Updates the limit of an existing pool.
39+
* Updates the limit of an existing pool. Children of the pool allocator
40+
* inherit the change automatically via Arrow's parent-cap check at
41+
* allocation time — no notification SPI is needed.
4042
*
4143
* @param poolName logical pool name
4244
* @param newLimit new maximum bytes for the pool

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ public final class NativeAllocatorPoolConfig {
3030
public static final String POOL_FLIGHT = "flight";
3131
/** Pool name for ingest pipeline memory. */
3232
public static final String POOL_INGEST = "ingest";
33+
/** Pool name for query-execution memory (analytics-engine fragments and per-query allocators). */
34+
public static final String POOL_QUERY = "query";
3335

3436
/** Setting key for the root allocator limit. */
3537
public static final String SETTING_ROOT_LIMIT = "native.allocator.root.limit";
@@ -42,6 +44,10 @@ public final class NativeAllocatorPoolConfig {
4244
public static final String SETTING_INGEST_MIN = "native.allocator.pool.ingest.min";
4345
/** Setting key for the ingest pool maximum. */
4446
public static final String SETTING_INGEST_MAX = "native.allocator.pool.ingest.max";
47+
/** Setting key for the query pool minimum. */
48+
public static final String SETTING_QUERY_MIN = "native.allocator.pool.query.min";
49+
/** Setting key for the query pool maximum. */
50+
public static final String SETTING_QUERY_MAX = "native.allocator.pool.query.max";
4551

4652
private NativeAllocatorPoolConfig() {}
4753
}

libs/arrow-spi/src/test/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfigTests.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,16 @@ public class NativeAllocatorPoolConfigTests extends OpenSearchTestCase {
1515
public void testPoolConstants() {
1616
assertEquals("flight", NativeAllocatorPoolConfig.POOL_FLIGHT);
1717
assertEquals("ingest", NativeAllocatorPoolConfig.POOL_INGEST);
18+
assertEquals("query", NativeAllocatorPoolConfig.POOL_QUERY);
1819
}
1920

2021
public void testSettingKeys() {
2122
assertEquals("native.allocator.pool.flight.min", NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN);
2223
assertEquals("native.allocator.pool.flight.max", NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX);
2324
assertEquals("native.allocator.pool.ingest.min", NativeAllocatorPoolConfig.SETTING_INGEST_MIN);
2425
assertEquals("native.allocator.pool.ingest.max", NativeAllocatorPoolConfig.SETTING_INGEST_MAX);
26+
assertEquals("native.allocator.pool.query.min", NativeAllocatorPoolConfig.SETTING_QUERY_MIN);
27+
assertEquals("native.allocator.pool.query.max", NativeAllocatorPoolConfig.SETTING_QUERY_MAX);
2528
}
2629

2730
public void testRootSettingKey() {
Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.arrow.allocator;
10+
11+
import org.apache.arrow.memory.OutOfMemoryException;
12+
import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
13+
14+
import java.util.function.Supplier;
15+
16+
/**
17+
* Translates Arrow's {@link OutOfMemoryException} into OpenSearch's standard
18+
* {@link OpenSearchRejectedExecutionException} at allocation boundaries.
19+
*
20+
* <p>Arrow throws {@code OutOfMemoryException} when an allocation would exceed
21+
* a parent's {@code allocationLimit}. This is a per-request, recoverable
22+
* condition — the caller should back off and retry — but Arrow's exception type
23+
* is not recognized by OpenSearch's REST layer, so it surfaces to clients as a
24+
* generic 500. {@link OpenSearchRejectedExecutionException} (a subclass of
25+
* {@link java.util.concurrent.RejectedExecutionException}) maps to HTTP 429
26+
* (Too Many Requests) and is the framework's standard rejection signal.
27+
*
28+
* <p>Apply this wrapper at <em>per-request</em> allocation sites — places where
29+
* a failure represents "this request didn't fit" rather than "the framework
30+
* failed to start". Examples:
31+
* <ul>
32+
* <li>per-query child allocator creation in
33+
* {@code DefaultPlanExecutor.executeInternal}</li>
34+
* <li>per-VSR child allocator creation in
35+
* {@code ArrowBufferPool.createChildAllocator}</li>
36+
* <li>request-handler-side allocations in transport/RPC layers</li>
37+
* </ul>
38+
*
39+
* <p><strong>Do not</strong> wrap startup-time / lifetime-of-component
40+
* allocations (e.g. transport server/client allocators created in
41+
* {@code doStart}, service-level allocators created in plugin constructors).
42+
* Those failing during node init is a configuration error, not a per-request
43+
* back-pressure signal — they should propagate as the original exception.
44+
*/
45+
public final class AllocationRejection {
46+
47+
private AllocationRejection() {}
48+
49+
/**
50+
* Runs {@code body} and translates any Arrow {@link OutOfMemoryException}
51+
* into {@link OpenSearchRejectedExecutionException}.
52+
*
53+
* @param context short label included in the rejection message (typically the
54+
* pool or operation name; e.g. "query-pool", "ingest-vsr")
55+
* @param body the allocation site to invoke
56+
* @param <T> return type of {@code body}
57+
* @return whatever {@code body} returns on success
58+
* @throws OpenSearchRejectedExecutionException if {@code body} throws
59+
* {@link OutOfMemoryException}; the original Arrow exception is
60+
* attached as the cause
61+
*/
62+
public static <T> T wrap(String context, Supplier<T> body) {
63+
try {
64+
return body.get();
65+
} catch (OutOfMemoryException e) {
66+
throw rejection(context, e);
67+
}
68+
}
69+
70+
/**
71+
* Runs {@code body} and translates any Arrow {@link OutOfMemoryException}
72+
* into {@link OpenSearchRejectedExecutionException}. Use this overload for
73+
* void allocation sites.
74+
*
75+
* @param context short label included in the rejection message (typically the
76+
* pool or operation name; e.g. "query-pool", "ingest-vsr")
77+
* @param body the allocation site to invoke
78+
* @throws OpenSearchRejectedExecutionException if {@code body} throws
79+
* {@link OutOfMemoryException}; the original Arrow exception is
80+
* attached as the cause
81+
*/
82+
public static void wrap(String context, Runnable body) {
83+
try {
84+
body.run();
85+
} catch (OutOfMemoryException e) {
86+
throw rejection(context, e);
87+
}
88+
}
89+
90+
private static OpenSearchRejectedExecutionException rejection(String context, OutOfMemoryException cause) {
91+
OpenSearchRejectedExecutionException rejection = new OpenSearchRejectedExecutionException(
92+
"native memory allocation rejected at [" + context + "]: " + cause.getMessage()
93+
);
94+
rejection.initCause(cause);
95+
return rejection;
96+
}
97+
}

0 commit comments

Comments
 (0)