diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java index 8c06116a1ee26..dd9866723651b 100644 --- a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java @@ -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; @@ -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() {} @@ -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 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 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 ────────────────────────────────────────────────────────── @@ -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), @@ -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 -> { diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java index baa791282ef68..44f9d0e95584a 100644 --- a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java @@ -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; @@ -53,6 +54,8 @@ public class ArrowNativeAllocator implements NativeAllocator { private final ConcurrentMap pools = new ConcurrentHashMap<>(); private final ConcurrentMap virtualPools = new ConcurrentHashMap<>(); private final ConcurrentMap poolConfigs = new ConcurrentHashMap<>(); + /** Pools excluded from budget validation, the rebalancer, and pool-group limit sums. */ + private final Set unmanagedPools = ConcurrentHashMap.newKeySet(); private final ConcurrentMap>> poolGroupLimitListeners = new ConcurrentHashMap<>(); private final List statsRefreshers = new CopyOnWriteArrayList<>(); private volatile Supplier nativeMemoryStatsSupplier; @@ -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 getManagedPoolNames() { + Set managed = new HashSet<>(getAllPoolNames()); + managed.removeAll(unmanagedPools); + return Collections.unmodifiableSet(managed); + } + @Override public void setPoolLimit(String poolName, long newLimit) { PoolConfig config = poolConfigs.get(poolName); @@ -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); @@ -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()); } } @@ -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); } } } @@ -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; } } diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java index 5a5294b2cb4b0..57a3014b13705 100644 --- a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java @@ -104,7 +104,10 @@ public void run() { } void rebalance() { - Set 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 allPools = allocator.getManagedPoolNames(); if (allPools.isEmpty()) return; long budget = budgetSupplier.get(); diff --git a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java index 189cee5b33f39..7dea7f23290a5 100644 --- a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java +++ b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java @@ -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 }); } // ----------------------------------------------------------------- @@ -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(); @@ -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(); diff --git a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java index c9e370bc22413..7ef09cd7663a3 100644 --- a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java +++ b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java @@ -14,6 +14,8 @@ import org.opensearch.plugin.stats.NativeAllocatorPoolStats; import org.opensearch.test.OpenSearchTestCase; +import java.util.concurrent.atomic.AtomicLong; + public class ArrowNativeAllocatorTests extends OpenSearchTestCase { private ArrowNativeAllocator allocator; @@ -142,7 +144,7 @@ public void testPoolGroupLimitListenerFiresWithCorrectSum() { allocator.getOrCreatePool("ingest", 0L, 80 * 1024 * 1024, PoolGroup.INDEXING); allocator.getOrCreatePool("write", 0L, 50 * 1024 * 1024, PoolGroup.INDEXING); - java.util.concurrent.atomic.AtomicLong received = new java.util.concurrent.atomic.AtomicLong(-1); + AtomicLong received = new AtomicLong(-1); allocator.addPoolGroupLimitListener(PoolGroup.INDEXING, received::set); // Change limits @@ -160,7 +162,7 @@ public void testPoolGroupLimitListenerNotFiredForOtherGroups() { allocator.getOrCreatePool("flight", 0L, 50 * 1024 * 1024, PoolGroup.TRANSPORT); allocator.getOrCreatePool("ingest", 0L, 80 * 1024 * 1024, PoolGroup.INDEXING); - java.util.concurrent.atomic.AtomicLong transportReceived = new java.util.concurrent.atomic.AtomicLong(-1); + AtomicLong transportReceived = new AtomicLong(-1); allocator.addPoolGroupLimitListener(PoolGroup.TRANSPORT, transportReceived::set); // Fire INDEXING group — should NOT trigger TRANSPORT listener @@ -171,4 +173,71 @@ public void testPoolGroupLimitListenerNotFiredForOtherGroups() { allocator.firePoolGroupListeners(PoolGroup.TRANSPORT); assertEquals(50L * 1024 * 1024, transportReceived.get()); } + + // ─── Unmanaged (special, unbounded) pool ─────────────────────────────────────── + + public void testRegisterUnmanagedPoolIsUnboundedAndFlagged() { + NativeAllocator.PoolHandle handle = allocator.registerUnmanagedPool("query", PoolGroup.SEARCH); + assertNotNull(handle); + assertEquals("unmanaged pool is fixed at Long.MAX_VALUE", Long.MAX_VALUE, handle.limit()); + assertTrue(allocator.isUnmanagedPool("query")); + assertFalse(allocator.isUnmanagedPool("some-other-pool")); + } + + public void testRegisterUnmanagedPoolRejectsExistingManagedPool() { + // A pool created as managed (bounded child allocator) must not be flippable to unmanaged. + allocator.getOrCreatePool("query", 0L, 100 * 1024 * 1024, PoolGroup.SEARCH); + IllegalStateException e = expectThrows( + IllegalStateException.class, + () -> allocator.registerUnmanagedPool("query", PoolGroup.SEARCH) + ); + assertTrue(e.getMessage().contains("already exists as a managed pool")); + assertFalse("must not be flagged unmanaged after a rejected re-register", allocator.isUnmanagedPool("query")); + } + + public void testUnmanagedPoolExcludedFromManagedPoolNames() { + allocator.getOrCreatePool("flight", 0L, 50 * 1024 * 1024, PoolGroup.TRANSPORT); + allocator.registerUnmanagedPool("query", PoolGroup.SEARCH); + + // Both appear in the full set, but only the managed one is in getManagedPoolNames. + assertTrue(allocator.getAllPoolNames().contains("query")); + assertTrue(allocator.getAllPoolNames().contains("flight")); + assertTrue(allocator.getManagedPoolNames().contains("flight")); + assertFalse( + "unmanaged query pool must be excluded from the rebalancer's managed set", + allocator.getManagedPoolNames().contains("query") + ); + } + + public void testUnmanagedPoolExcludedFromBudgetValidation() { + // Budget only fits the flight pool's max; the unbounded query pool must NOT be summed + // against the budget (else its Long.MAX_VALUE would trip validation). + allocator.setBudget(100 * 1024 * 1024); + allocator.getOrCreatePool("flight", 0L, 50 * 1024 * 1024, PoolGroup.TRANSPORT); + // Registering the unbounded pool must not throw despite Long.MAX_VALUE > budget. + allocator.registerUnmanagedPool("query", PoolGroup.SEARCH); + + // A managed pool whose max would exceed the remaining budget still trips — proving the + // unmanaged pool simply isn't counted, not that validation was disabled entirely. + expectThrows(IllegalArgumentException.class, () -> allocator.getOrCreatePool("ingest", 0L, 80 * 1024 * 1024, PoolGroup.INDEXING)); + } + + public void testUnmanagedPoolExcludedFromPoolGroupSum() { + // A managed SEARCH pool + an unmanaged SEARCH pool: the group sum must reflect only the + // managed one, not the unmanaged pool's Long.MAX_VALUE limit. + allocator.getOrCreatePool("datafusion", 0L, 70 * 1024 * 1024, PoolGroup.SEARCH); + allocator.registerUnmanagedPool("query", PoolGroup.SEARCH); + + AtomicLong received = new AtomicLong(-1); + allocator.addPoolGroupLimitListener(PoolGroup.SEARCH, received::set); + allocator.firePoolGroupListeners(PoolGroup.SEARCH); + + assertEquals("group sum must exclude the unmanaged pool", 70L * 1024 * 1024, received.get()); + } + + public void testUnmanagedPoolLimitUntouchedByResetAllPoolsToMax() { + allocator.registerUnmanagedPool("query", PoolGroup.SEARCH); + allocator.resetAllPoolsToMax(); + assertEquals(Long.MAX_VALUE, allocator.getPoolAllocator("query").getLimit()); + } } diff --git a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java index b92a477a67ea3..812a4d834c46a 100644 --- a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java +++ b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java @@ -10,6 +10,7 @@ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; +import org.opensearch.arrow.spi.PoolGroup; import org.opensearch.test.OpenSearchTestCase; import java.util.ArrayList; @@ -149,4 +150,26 @@ public void testSumLimitsNeverExceedsBudget() { bufs.forEach(ArrowBuf::close); } } + + public void testUnmanagedPoolUntouchedByRebalance() { + // A pressured managed pool that would normally receive freed capacity... + allocator.getOrCreatePool("pressured", 5 * MB, 20 * MB, PoolGroup.INDEXING); + // ...and the special unbounded query pool, which is idle (0 allocated). + allocator.registerUnmanagedPool("query", PoolGroup.SEARCH); + + BufferAllocator pressuredPool = allocator.getPoolAllocator("pressured"); + ArrowBuf buf = pressuredPool.buffer((long) (20 * MB * 0.9)); // >75% → pressured + try { + long queryLimitBefore = allocator.getPoolAllocator("query").getLimit(); + rebalancer.rebalance(); + long queryLimitAfter = allocator.getPoolAllocator("query").getLimit(); + + // The unmanaged pool's limit is unchanged (still Long.MAX_VALUE) — it was neither + // shrunk as an "idle" pool nor otherwise touched. + assertEquals("unmanaged pool must not be resized by the rebalancer", queryLimitBefore, queryLimitAfter); + assertEquals(Long.MAX_VALUE, queryLimitAfter); + } finally { + buf.close(); + } + } } diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java index e28e373fe69a9..5360cfbf09cf1 100644 --- a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java @@ -97,4 +97,43 @@ public void testIdlePoolShrinksWhenOtherPressured() throws Exception { buf.close(); } } + + /** + * The query pool is registered as a special, unmanaged, unbounded pool on a real node: its limit + * is Long.MAX_VALUE regardless of native.allocator.pool.query.max, it reports as unmanaged, and + * it is excluded from the rebalancer's managed set — so the running rebalancer never resizes it. + */ + public void testQueryPoolIsUnmanagedAndUnbounded() throws Exception { + ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); + + assertTrue( + "query pool must be registered as unmanaged", + allocator.isUnmanagedPool(NativeAllocatorPoolConfig.POOL_QUERY) + ); + assertEquals( + "query pool must be unbounded (Long.MAX_VALUE), ignoring the query.max setting", + Long.MAX_VALUE, + allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit() + ); + assertFalse( + "query pool must be excluded from the rebalancer's managed set", + allocator.getManagedPoolNames().contains(NativeAllocatorPoolConfig.POOL_QUERY) + ); + + // Pressure a managed pool so the (1s-interval) rebalancer actively runs, then confirm the + // query pool's limit is still untouched after several ticks. + BufferAllocator ingestPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST); + ArrowBuf buf = ingestPool.buffer((long) (ingestPool.getLimit() * 0.8)); + try { + assertBusy(() -> { + assertEquals( + "rebalancer must never resize the unmanaged query pool", + Long.MAX_VALUE, + allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit() + ); + }); + } finally { + buf.close(); + } + } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java index c07574b76abed..3c04ec48bf037 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java @@ -27,7 +27,6 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.action.ActionListener; -import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; @@ -76,7 +75,6 @@ public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); - iteratorInstance.reclaimDrainedStaging(); } } finally { try { @@ -99,9 +97,6 @@ static class BatchIterator implements Iterator { private Boolean nextAvailable; private boolean batchEmitted; private boolean nativeStreamExhausted; - // Per-batch staging allocators used by {@link #importBatch}. Each is reclaimed once its batch's - // buffers have been released by the consumer (see {@link #reclaimDrainedStaging}). - private final List stagingAllocators = new ArrayList<>(); BatchIterator(StreamHandle streamHandle, BufferAllocator allocator, CDataDictionaryProvider dictionaryProvider) { this.streamHandle = streamHandle; @@ -139,81 +134,15 @@ private boolean loadNextBatch() { } return false; } + VectorSchemaRoot freshRoot = VectorSchemaRoot.create(schema, allocator); try (ArrowArray arrowArray = ArrowArray.wrap(arrayAddr)) { - nextBatch = importBatch(arrowArray); + Data.importIntoVectorSchemaRoot(allocator, arrowArray, freshRoot, dictionaryProvider); } + nextBatch = freshRoot; batchEmitted = true; return true; } - /** - * Imports one native batch across the Arrow C Data Interface into a per-batch staging allocator - * (an unbounded child of the root) rather than directly into {@code allocator}. - * - *

{@link Data#importIntoVectorSchemaRoot} charges each buffer against the target allocator as it - * walks the array. Against a bounded target that fills part-way through a wide batch the import - * throws, and arrow-java's {@code ReferenceCountedArrowArray#unsafeAssociateAllocation} retains the - * imported array before the throwing {@code wrapForeignAllocation} without rolling back, so - * the C Data release callback never fires and the whole native batch leaks in the producer's native - * allocator — invisible to the JVM heap and the Java Arrow allocator (arrow-java ≤ 18.1.0). An - * unbounded staging child can't OOM mid-array, so the release callback always fires. - * - *

The batch is returned as-is (zero-copy); its buffers are released by the existing consumer close - * paths, which drives the C Data reference count to zero. Each staging allocator is reclaimed once - * drained (see {@link #reclaimDrainedStaging}); on import failure it is closed immediately. - */ - private VectorSchemaRoot importBatch(ArrowArray arrowArray) { - reclaimDrainedStaging(); - BufferAllocator staging = allocator.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE); - try { - VectorSchemaRoot root = importOntoStaging(staging, schema, arrowArray, dictionaryProvider); - stagingAllocators.add(staging); - return root; - } catch (RuntimeException e) { - staging.close(); - throw e; - } - } - - /** - * Closes staging allocators whose batches have been fully released (drained to zero). A batch still - * in flight keeps its staging allocator open so the eventual release callback frees the small C Data - * bookkeeping allocation against a live allocator; that allocator is a leaf child of the root and - * holds no batch data once drained. - */ - private void reclaimDrainedStaging() { - stagingAllocators.removeIf(a -> { - if (a.getAllocatedMemory() == 0) { - a.close(); - return true; - } - return false; - }); - } - - /** - * Imports {@code arrowArray} into a fresh {@link VectorSchemaRoot} on {@code staging}, which MUST be - * an unbounded child of the root so the import cannot OOM part-way through the array. On failure the - * returned root is closed (firing the native release for the whole batch) and the exception rethrown; - * the caller owns {@code staging}. Package-private so the leak regression test can drive the exact - * production import path. - */ - static VectorSchemaRoot importOntoStaging( - BufferAllocator staging, - Schema schema, - ArrowArray arrowArray, - CDataDictionaryProvider dictionaryProvider - ) { - VectorSchemaRoot root = VectorSchemaRoot.create(schema, staging); - try { - Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); - } catch (RuntimeException e) { - root.close(); - throw e; - } - return root; - } - @Override public boolean hasNext() { if (nextAvailable == null) { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionImportLeakTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionImportLeakTests.java deleted file mode 100644 index 5a6efc62d4d4b..0000000000000 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionImportLeakTests.java +++ /dev/null @@ -1,162 +0,0 @@ -/* - * 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.be.datafusion; - -import org.apache.arrow.c.ArrowArray; -import org.apache.arrow.c.ArrowSchema; -import org.apache.arrow.c.CDataDictionaryProvider; -import org.apache.arrow.c.Data; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.OutOfMemoryException; -import org.apache.arrow.memory.RootAllocator; -import org.apache.arrow.vector.FieldVector; -import org.apache.arrow.vector.VarCharVector; -import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.types.pojo.Schema; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; - -/** - * Regression tests for the native batch leak on Arrow C Data Interface import under allocator pressure. - * - *

Self-contained (no native runtime): a batch is built on a "producer" allocator and exported across - * the C Data Interface. The producer allocator stands in for the native (Rust) allocator that owns the - * exported buffers — when import releases the C Data array the producer drains to zero; when it leaks the - * producer stays non-zero. - * - *

{@link #testDirectImportIntoBoundedAllocatorLeaksExportedBatch} pins the arrow-java bug: importing - * directly into a bounded allocator throws {@link OutOfMemoryException} mid-array and strands the whole - * exported batch (producer not drained). {@link #testProductionStagingImportReleasesBatch} drives the - * production path ({@link DatafusionResultStream.BatchIterator#importOntoStaging}) and asserts the batch is - * released even when the caller's allocator is far too small — this test fails on the unfixed direct-import - * code and passes with the staging import. - */ -public class DatafusionImportLeakTests extends OpenSearchTestCase { - - private static final int ROWS = 4096; - private static final int VALUE_BYTES = 512; - private static final int COLUMNS = 4; - private static final long TINY_LIMIT = 64 * 1024; // far smaller than the ~8 MB batch - - private RootAllocator root; - - @Override - public void setUp() throws Exception { - super.setUp(); - root = new RootAllocator(Long.MAX_VALUE); - } - - @Override - public void tearDown() throws Exception { - root.close(); - super.tearDown(); - } - - /** - * Production path: {@link DatafusionResultStream.BatchIterator#importOntoStaging} imports onto an - * unbounded staging child of the root, so a caller allocator far too small to hold the batch does not - * cause a mid-import OOM; the batch imports, and closing it releases the exported buffers — the producer - * allocator drains to zero. - * - *

On the unfixed code (import directly into the bounded caller allocator) this leaks and the assertion - * fails, so this is a genuine regression guard. - */ - public void testProductionStagingImportReleasesBatch() throws Exception { - BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE); - try (ArrowArray array = ArrowArray.allocateNew(producer); ArrowSchema cSchema = ArrowSchema.allocateNew(producer)) { - exportBatch(producer, array, cSchema); - assertTrue("producer holds the exported batch before import", producer.getAllocatedMemory() > 0); - - // Stand-in for the caller allocator: intentionally far smaller than the batch. - BufferAllocator tinyTarget = root.newChildAllocator("tiny-target", 0, TINY_LIMIT); - BufferAllocator staging = tinyTarget.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE); - try (CDataDictionaryProvider dp = new CDataDictionaryProvider()) { - Schema schema = Data.importSchema(staging, cSchema, dp); - VectorSchemaRoot imported = DatafusionResultStream.BatchIterator.importOntoStaging(staging, schema, array, dp); - imported.close(); // consumer releases the batch → C Data release fires - } - staging.close(); - tinyTarget.close(); - } - assertEquals("staging import must release the exported batch (producer drained)", 0L, producer.getAllocatedMemory()); - producer.close(); - } - - /** - * Pins the arrow-java bug this fix works around: importing the exported batch directly into a - * bounded allocator (the pre-fix behaviour) throws {@link OutOfMemoryException} mid-array and - * leaks — the producer allocator is not drained because the C Data release callback never fires. - */ - public void testDirectImportIntoBoundedAllocatorLeaksExportedBatch() throws Exception { - // Dedicated local root, deliberately abandoned: this test demonstrates the arrow-java leak, so the - // producer/target/C-Data structs cannot be drained or closed. Keeping it off the shared #root (and - // never closing anything here) avoids tripping the allocator leak detector on this expected leak. - RootAllocator localRoot = new RootAllocator(Long.MAX_VALUE); - BufferAllocator producer = localRoot.newChildAllocator("producer-direct", 0, Long.MAX_VALUE); - ArrowArray array = ArrowArray.allocateNew(producer); - ArrowSchema cSchema = ArrowSchema.allocateNew(producer); - exportBatch(producer, array, cSchema); - - BufferAllocator tinyTarget = localRoot.newChildAllocator("tiny-target-direct", 0, TINY_LIMIT); - CDataDictionaryProvider dp = new CDataDictionaryProvider(); - Schema schema = Data.importSchema(tinyTarget, cSchema, dp); - VectorSchemaRoot target = VectorSchemaRoot.create(schema, tinyTarget); - // importIntoVectorSchemaRoot wraps the allocator OOM as IllegalArgumentException - // ("Could not load buffers for field ...") with the OutOfMemoryException as its cause. - Exception thrown = expectThrows(Exception.class, () -> Data.importIntoVectorSchemaRoot(tinyTarget, array, target, dp)); - assertTrue("mid-import failure must be an allocator OOM, was: " + thrown, hasOomCause(thrown)); - assertTrue("arrow-java strands the exported batch on a mid-import OOM (producer not drained)", producer.getAllocatedMemory() > 0); - // Nothing is closed: the mid-import OOM strands buffers on producer and target and never fires the - // C-Data release, so no close would succeed. localRoot is abandoned; the JVM reclaims the wrappers. - } - - /** True if {@code t} is, or is caused by, an Arrow {@link OutOfMemoryException}. */ - private static boolean hasOomCause(Throwable t) { - for (Throwable c = t; c != null; c = c.getCause()) { - if (c instanceof OutOfMemoryException) { - return true; - } - } - return false; - } - - /** Builds a multi-column VarChar batch on {@code alloc} and exports it into {@code array}/{@code cSchema}. */ - private void exportBatch(BufferAllocator alloc, ArrowArray array, ArrowSchema cSchema) { - VectorSchemaRoot source = buildBatch(alloc); - try { - long total = 0; - for (FieldVector v : source.getFieldVectors()) { - total += v.getBufferSize(); - } - assertTrue("test setup: batch must exceed the tiny target limit", total > TINY_LIMIT); - Data.exportVectorSchemaRoot(alloc, source, null, array, cSchema); - } finally { - source.close(); - } - } - - private VectorSchemaRoot buildBatch(BufferAllocator alloc) { - byte[] value = new byte[VALUE_BYTES]; - Arrays.fill(value, (byte) 'x'); - List fieldVectors = new ArrayList<>(COLUMNS); - for (int c = 0; c < COLUMNS; c++) { - VarCharVector v = new VarCharVector("col" + c, alloc); - v.allocateNew((long) ROWS * VALUE_BYTES, ROWS); - for (int r = 0; r < ROWS; r++) { - v.setSafe(r, value); - } - v.setValueCount(ROWS); - fieldVectors.add(v); - } - return new VectorSchemaRoot(fieldVectors); - } -} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java index 1136afa93b98b..28db10793bffc 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java @@ -297,37 +297,7 @@ public void testDoubleCloseIsHarmless() throws Exception { stream.close(); } - /** - * End-to-end smoke test: a real native stream consumed through {@link DatafusionResultStream} under a - * small caller allocator imports and releases cleanly, leaving the shared root allocator drained. - * - *

Note: the deterministic regression for the mid-import-OOM native leak lives in - * {@link DatafusionImportLeakTests}; this test's fixture ({@code test.parquet}) is tiny, so it does not - * itself force a mid-import OOM — it verifies the staging import path is wired correctly end-to-end and - * that a normal small batch is fully released. - */ - public void testStreamConsumeAndCloseDrainsAllocator() throws Exception { - try (DatafusionResultStream stream = createStreamWithLimit("SELECT message, message2 FROM test_table", 1024)) { - Iterator it = stream.iterator(); - while (it.hasNext()) { - EngineResultBatch batch = it.next(); - batch.getArrowRoot().close(); - } - } - assertEquals("root allocator not fully drained after import + consume + close", 0L, testRootAllocator.getAllocatedMemory()); - } - - private DatafusionResultStream createStreamWithLimit(String sql, long limitBytes) { - long streamPtr = executeQueryForStream(sql); - BufferAllocator childAllocator = testRootAllocator.newChildAllocator("test-stream-bounded", 0, limitBytes); - allocatorsToClose.add(childAllocator); - return new DatafusionResultStream( - new org.opensearch.be.datafusion.nativelib.StreamHandle(streamPtr, runtimeHandle), - childAllocator - ); - } - - private long executeQueryForStream(String sql) { + private DatafusionResultStream createStream(String sql) { byte[] substrait = NativeBridge.sqlToSubstrait(readerHandle.getPointer(), "test_table", sql, runtimeHandle.get()); CompletableFuture future = new CompletableFuture<>(); NativeBridge.executeQueryAsync( @@ -349,11 +319,7 @@ public void onFailure(Exception e) { } } ); - return future.join(); - } - - private DatafusionResultStream createStream(String sql) { - long streamPtr = executeQueryForStream(sql); + long streamPtr = future.join(); BufferAllocator childAllocator = testRootAllocator.newChildAllocator("test-stream", 0, Long.MAX_VALUE); allocatorsToClose.add(childAllocator); return new DatafusionResultStream( diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java index 357b40fcaea76..a0434251ebf0f 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java @@ -23,7 +23,6 @@ import org.opensearch.analytics.exec.ArrowValues; import org.opensearch.common.annotation.ExperimentalApi; -import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; @@ -80,7 +79,6 @@ public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); - iteratorInstance.reclaimDrainedStaging(); } } finally { try { @@ -112,9 +110,6 @@ static class BatchIterator implements Iterator { private Boolean nextAvailable; private boolean batchEmitted; private boolean exhausted; - // Per-batch staging allocators used by {@link #importBatch}. Each is reclaimed once its batch's - // buffers have been released by the consumer (see {@link #reclaimDrainedStaging}). - private final List stagingAllocators = new ArrayList<>(); BatchIterator( ArrowArray arrowArray, @@ -140,59 +135,14 @@ private void ensureSchema() { private boolean loadNextBatch() { ensureSchema(); if (exhausted) return false; - nextBatch = importBatch(); + VectorSchemaRoot freshRoot = VectorSchemaRoot.create(schema, allocator); + Data.importIntoVectorSchemaRoot(allocator, arrowArray, freshRoot, dictionaryProvider); + nextBatch = freshRoot; batchEmitted = true; exhausted = true; return true; } - /** - * Imports the batch across the Arrow C Data Interface into a per-batch staging allocator (an - * unbounded child of the root) rather than directly into {@code allocator}. - * - *

{@link Data#importIntoVectorSchemaRoot} charges each buffer against the target allocator as it - * walks the array. Against a bounded target that fills part-way through a wide batch the import - * throws, and arrow-java's {@code ReferenceCountedArrowArray#unsafeAssociateAllocation} retains the - * imported array before the throwing {@code wrapForeignAllocation} without rolling back, so - * the C Data release callback never fires and the whole native batch leaks — invisible to the JVM - * heap and the Java Arrow allocator (arrow-java ≤ 18.1.0). An unbounded staging child can't OOM - * mid-array, so the release callback always fires. - * - *

The batch is returned as-is (zero-copy); its buffers are released by the existing consumer close - * paths, which drives the C Data reference count to zero. Each staging allocator is reclaimed once - * drained (see {@link #reclaimDrainedStaging}); on import failure it is closed immediately. - */ - private VectorSchemaRoot importBatch() { - reclaimDrainedStaging(); - BufferAllocator staging = allocator.getRoot().newChildAllocator("lucene-import-staging", 0, Long.MAX_VALUE); - VectorSchemaRoot root = VectorSchemaRoot.create(schema, staging); - try { - Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); - } catch (RuntimeException e) { - root.close(); - staging.close(); - throw e; - } - stagingAllocators.add(staging); - return root; - } - - /** - * Closes staging allocators whose batches have been fully released (drained to zero). A batch still - * in flight keeps its staging allocator open so the eventual release callback frees the small C Data - * bookkeeping allocation against a live allocator; that allocator is a leaf child of the root and - * holds no batch data once drained. - */ - private void reclaimDrainedStaging() { - stagingAllocators.removeIf(a -> { - if (a.getAllocatedMemory() == 0) { - a.close(); - return true; - } - return false; - }); - } - @Override public boolean hasNext() { if (nextAvailable == null) {