Skip to content

[analytics-engine] Fix query-pool import leak: register POOL_QUERY as unmanaged/unbounded #22538

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sandeshkr419 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[analytics-engine] Fix query-pool import leak: register POOL_QUERY as unmanaged/unbounded #22538

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 34 additions & 17 deletions plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@

package org.opensearch.arrow.allocator;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.arrow.spi.NativeAllocatorPoolConfig;
import org.opensearch.arrow.spi.PoolGroup;
import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
Expand Down Expand Up @@ -56,6 +58,8 @@
*/
public class ArrowBasePlugin extends Plugin implements ExtensiblePlugin, ActionPlugin {

private static final Logger logger = LogManager.getLogger(ArrowBasePlugin.class);

/** Creates the plugin. */
public ArrowBasePlugin() {}

Expand Down Expand Up @@ -189,22 +193,26 @@ public ArrowBasePlugin() {}
Setting.Property.Dynamic
);

/** Minimum guaranteed bytes for the query pool. Default is 2% of budget. */
/** @deprecated The query pool is unbounded (registered as a special unmanaged pool); this has no effect. */
@Deprecated
public static final Setting<Long> QUERY_MIN_SETTING = new Setting<>(
NativeAllocatorPoolConfig.SETTING_QUERY_MIN,
s -> derivePoolMinDefault(s, 2),
s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_QUERY_MIN),
Setting.Property.NodeScope,
Setting.Property.Dynamic
Setting.Property.Dynamic,
Setting.Property.Deprecated
);

/** Maximum bytes the query pool can allocate. Default is 5% of budget. */
/** @deprecated The query pool is unbounded (registered as a special unmanaged pool); this has no effect. */
@Deprecated
public static final Setting<Long> QUERY_MAX_SETTING = new Setting<>(
NativeAllocatorPoolConfig.SETTING_QUERY_MAX,
s -> derivePoolMaxDefault(s, 5),
s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_QUERY_MAX),
Setting.Property.NodeScope,
Setting.Property.Dynamic
Setting.Property.Dynamic,
Setting.Property.Deprecated
);

// ─── Instance state ──────────────────────────────────────────────────────────
Expand Down Expand Up @@ -303,12 +311,10 @@ ArrowNativeAllocator buildAllocator(Settings settings, ClusterSettings cs, Suppl
allocator.setBudget(nativeBudget);
}

// Validate min < max for each pool
// Validate min < max for enforced pools
validateMinMax(NativeAllocatorPoolConfig.POOL_FLIGHT, FLIGHT_MIN_SETTING.get(settings), FLIGHT_MAX_SETTING.get(settings));
validateMinMax(NativeAllocatorPoolConfig.POOL_INGEST, INGEST_MIN_SETTING.get(settings), INGEST_MAX_SETTING.get(settings));
validateMinMax(NativeAllocatorPoolConfig.POOL_QUERY, QUERY_MIN_SETTING.get(settings), QUERY_MAX_SETTING.get(settings));

// Create pools (always start at max)
// Create pools (always start at max).
allocator.getOrCreatePool(
NativeAllocatorPoolConfig.POOL_FLIGHT,
FLIGHT_MIN_SETTING.get(settings),
Expand All @@ -321,20 +327,31 @@ ArrowNativeAllocator buildAllocator(Settings settings, ClusterSettings cs, Suppl
INGEST_MAX_SETTING.get(settings),
PoolGroup.INDEXING
);
allocator.getOrCreatePool(
NativeAllocatorPoolConfig.POOL_QUERY,
QUERY_MIN_SETTING.get(settings),
QUERY_MAX_SETTING.get(settings),
PoolGroup.SEARCH
);
// POOL_QUERY is unmanaged/unbounded: the C-Data importer retains a ref BEFORE calling
// allocateBytes, so any OOM permanently leaks the native batch. Real enforcement lives
// Rust-side (DataFusion MemoryPool). Do NOT add a limit or AllocationListener here.
allocator.registerUnmanagedPool(NativeAllocatorPoolConfig.POOL_QUERY, PoolGroup.SEARCH);

// Register dynamic setting consumers for min/max changes
// Register dynamic setting consumers for min/max changes (enforced pools only)
cs.addSettingsUpdateConsumer(FLIGHT_MIN_SETTING, newMin -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_FLIGHT, newMin));
cs.addSettingsUpdateConsumer(FLIGHT_MAX_SETTING, newMax -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_FLIGHT, newMax));
cs.addSettingsUpdateConsumer(INGEST_MIN_SETTING, newMin -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_INGEST, newMin));
cs.addSettingsUpdateConsumer(INGEST_MAX_SETTING, newMax -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_INGEST, newMax));
cs.addSettingsUpdateConsumer(QUERY_MIN_SETTING, newMin -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_QUERY, newMin));
cs.addSettingsUpdateConsumer(QUERY_MAX_SETTING, newMax -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_QUERY, newMax));
// QUERY min/max are deprecated no-ops — the pool is unbounded. Warn if explicitly configured.
if (QUERY_MIN_SETTING.exists(settings) || QUERY_MAX_SETTING.exists(settings)) {
logger.warn(
"native.allocator.pool.query.min/max are configured but have NO EFFECT: the query pool is "
+ "unbounded (native enforcement lives Rust-side). Remove these settings."
);
}
cs.addSettingsUpdateConsumer(
QUERY_MIN_SETTING,
v -> logger.warn("native.allocator.pool.query.min has no effect: query pool is unbounded")
);
cs.addSettingsUpdateConsumer(
QUERY_MAX_SETTING,
v -> logger.warn("native.allocator.pool.query.max has no effect: query pool is unbounded")
);

// Register dynamic consumer for rebalancer enable/disable
cs.addSettingsUpdateConsumer(REBALANCER_ENABLED_SETTING, enabled -> {
Expand Down
63 changes: 54 additions & 9 deletions plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import org.apache.arrow.memory.RootAllocator;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.logging.log4j.message.ParameterizedMessage;
import org.opensearch.arrow.spi.NativeAllocator;
import org.opensearch.arrow.spi.PoolGroup;
import org.opensearch.common.SetOnce;
Expand Down Expand Up @@ -53,6 +54,8 @@ public class ArrowNativeAllocator implements NativeAllocator {
private final ConcurrentMap<String, ArrowPoolHandle> pools = new ConcurrentHashMap<>();
private final ConcurrentMap<String, VirtualPoolHandleImpl> virtualPools = new ConcurrentHashMap<>();
private final ConcurrentMap<String, PoolConfig> poolConfigs = new ConcurrentHashMap<>();
/** Pools excluded from budget validation, the rebalancer, and pool-group limit sums. */
private final Set<String> unmanagedPools = ConcurrentHashMap.newKeySet();
private final ConcurrentMap<PoolGroup, List<Consumer<Long>>> poolGroupLimitListeners = new ConcurrentHashMap<>();
private final List<Runnable> statsRefreshers = new CopyOnWriteArrayList<>();
private volatile Supplier<long[]> nativeMemoryStatsSupplier;
Expand Down Expand Up @@ -267,6 +270,44 @@ public PoolHandle getOrCreatePool(String poolName, long min, long max, PoolGroup
});
}

/**
* Registers an unbounded pool excluded from budget validation, the rebalancer, and pool-group
* limit sums. Intended for POOL_QUERY: its bytes are zero-copy foreign wraps of pre-existing
* native memory, so limiting it would leak imported batches. Real enforcement lives Rust-side.
*
* @param poolName name of the pool
* @param group pool group (for stats/reporting only; excluded from group limit sums)
* @throws IllegalStateException if a managed pool with this name already exists
*/
public PoolHandle registerUnmanagedPool(String poolName, PoolGroup group) {
if (pools.containsKey(poolName)) {
throw new IllegalStateException(
"Pool ["
+ poolName
+ "] already exists as a managed pool and cannot be re-registered as unmanaged; "
+ "its child allocator has a bounded limit. Register it unmanaged from the start."
);
}
unmanagedPools.add(poolName);
poolConfigs.putIfAbsent(poolName, new PoolConfig(0, Long.MAX_VALUE, group));
return pools.computeIfAbsent(poolName, name -> {
BufferAllocator child = root.newChildAllocator(name, 0, Long.MAX_VALUE);
return new ArrowPoolHandle(child);
});
}

/** Whether a pool is a special/unmanaged unbounded pool (excluded from sizing math). */
public boolean isUnmanagedPool(String poolName) {
return unmanagedPools.contains(poolName);
}

/** Pool names the rebalancer should manage — all pools minus the unmanaged/special ones. */
public Set<String> getManagedPoolNames() {
Set<String> managed = new HashSet<>(getAllPoolNames());
managed.removeAll(unmanagedPools);
return Collections.unmodifiableSet(managed);
}

@Override
public void setPoolLimit(String poolName, long newLimit) {
PoolConfig config = poolConfigs.get(poolName);
Expand Down Expand Up @@ -375,6 +416,9 @@ public void setPoolEffectiveLimit(String poolName, long newLimit) {
*/
public void resetAllPoolsToMax() {
for (String name : getAllPoolNames()) {
if (unmanagedPools.contains(name)) {
continue; // special unbounded pool — nothing to reset
}
PoolConfig config = poolConfigs.get(name);
long max = config != null ? config.max : Long.MAX_VALUE;
long current = getEffectiveLimit(name);
Expand Down Expand Up @@ -646,7 +690,9 @@ public void firePoolGroupListeners(PoolGroup group) {

long groupSum = 0;
for (var entry : poolConfigs.entrySet()) {
if (entry.getValue().group == group) {
// Skip unmanaged/special pools: their Long.MAX_VALUE effective limit would swamp the
// grouped total that group listeners (e.g. the DataFusion pool sizer) consume.
if (entry.getValue().group == group && unmanagedPools.contains(entry.getKey()) == false) {
groupSum += getEffectiveLimit(entry.getKey());
}
}
Expand All @@ -655,13 +701,7 @@ public void firePoolGroupListeners(PoolGroup group) {
try {
listener.accept(finalSum);
} catch (Exception e) {
logger.warn(
() -> new org.apache.logging.log4j.message.ParameterizedMessage(
"Pool group limit listener failed for group [{}]",
group
),
e
);
logger.warn(() -> new ParameterizedMessage("Pool group limit listener failed for group [{}]", group), e);
}
}
}
Expand All @@ -672,9 +712,14 @@ private void validateSumMaxesWithinBudget(String newPoolName, long newPoolMax) {
if (budget == Long.MAX_VALUE || budget <= 0) {
return;
}
// Unmanaged/special pools (e.g. POOL_QUERY) are unbounded by design; their Long.MAX_VALUE
// "max" is not real budget consumption and must not be summed against the node budget.
if (unmanagedPools.contains(newPoolName)) {
return;
}
long sumMaxes = newPoolMax;
for (var entry : poolConfigs.entrySet()) {
if (entry.getKey().equals(newPoolName) == false) {
if (entry.getKey().equals(newPoolName) == false && unmanagedPools.contains(entry.getKey()) == false) {
sumMaxes += entry.getValue().max;
}
}
Expand Down
5 changes: 4 additions & 1 deletion plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,10 @@ public void run() {
}

void rebalance() {
Set<String> allPools = allocator.getAllPoolNames();
// Managed pools only: unmanaged/special pools (e.g. the unbounded POOL_QUERY) are excluded so
// their Long.MAX_VALUE limit is never treated as idle capacity to redistribute, and they are
// never shrunk/grown.
Set<String> allPools = allocator.getManagedPoolNames();
if (allPools.isEmpty()) return;

long budget = budgetSupplier.get();
Expand Down
17 changes: 12 additions & 5 deletions plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,8 @@ public void testPoolMaxRejectsNegative() {
IllegalArgumentException eQuery = expectThrows(IllegalArgumentException.class, () -> ArrowBasePlugin.QUERY_MAX_SETTING.get(query));
assertTrue(eQuery.getMessage().contains("must be >= 0"));
assertTrue("query error must name its setting", eQuery.getMessage().contains(NativeAllocatorPoolConfig.SETTING_QUERY_MAX));
// QUERY_MAX is deprecated (the query pool is unbounded) — reading an explicitly-set value warns.
assertSettingDeprecationsAndWarnings(new Setting<?>[] { ArrowBasePlugin.QUERY_MAX_SETTING });
}

// -----------------------------------------------------------------
Expand Down Expand Up @@ -119,11 +121,16 @@ public void testBuildAllocatorWiresAllPools() throws Exception {
assertTrue(poolNames.contains(NativeAllocatorPoolConfig.POOL_INGEST));
assertTrue(poolNames.contains(NativeAllocatorPoolConfig.POOL_QUERY));

// Pool maxes match the operator-set values (rebalancer disabled,
// so initial limit == max).
// Flight/ingest maxes match the operator-set values (rebalancer disabled, so limit == max).
assertEquals(1L * 1024 * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT).getLimit());
assertEquals(2L * 1024 * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST).getLimit());
assertEquals(1L * 1024 * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit());
// POOL_QUERY is unmanaged/unbounded — pool.query.max has no effect.
assertTrue(allocator.isUnmanagedPool(NativeAllocatorPoolConfig.POOL_QUERY));
assertEquals(Long.MAX_VALUE, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit());
assertFalse(
"unmanaged query pool must be excluded from the rebalancer's managed set",
allocator.getManagedPoolNames().contains(NativeAllocatorPoolConfig.POOL_QUERY)
);
} finally {
allocator.close();
plugin.close();
Expand All @@ -148,10 +155,10 @@ public void testBuildAllocatorWithRebalancerPoolsStartAtMax() throws Exception {
long budget2 = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(nodeSettings).getBytes();
ArrowNativeAllocator allocator = plugin.buildAllocator(nodeSettings, cs, () -> budget2);
try {
// Pools always start at max regardless of rebalancer state
// Managed pools always start at max regardless of rebalancer state.
assertEquals(200L * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT).getLimit());
assertEquals(200L * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST).getLimit());
assertEquals(200L * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit());
assertEquals(Long.MAX_VALUE, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit());
} finally {
allocator.close();
plugin.close();
Expand Down
Loading
Loading