Wire Arrow consumers to unified native allocator - #225
Closed
gaurav-amz wants to merge 3 commits into
Closed
Conversation
Introduces a new plugin (native-allocator-arrow) that owns a single Arrow RootAllocator for the node, with named pool-level children for each subsystem (flight, ingest). Each pool has a min (guaranteed floor) and max (burst ceiling). A background rebalancer (disabled by default) redistributes unused capacity across pools so active pools can grow beyond their min up to their max. The SPI (libs/arrow-spi) defines the Arrow-agnostic interface, pool config constants, and Writeable stats shape. The plugin implements it with Arrow's BufferAllocator and registers configurable limits as dynamic cluster settings. Also wires the Parquet ingest path to use the unified ingest pool instead of a standalone RootAllocator. Signed-off-by: Bukhtawar Khan <bukhtawa@amazon.com>
Signed-off-by: Bukhtawar Khan <bukhtawa@amazon.com>
Builds on opensearch-project#21703 with the framework-side changes that were not part of its initial scope. Consumer wiring for arrow-flight-rpc and analytics-engine is handled by opensearch-project#21465 (arrow-base, already merged) and not duplicated here. Pools added: POOL_QUERY analytics-engine query execution POOL_DATAFUSION DataFusion native MemoryPool mirror Both registered with min/max settings dynamically updatable. The flight and ingest pools registered by opensearch-project#21703 are unchanged. Cross-cutting changes: * NativeAllocatorListener SPI in arrow-spi for pool-resize callbacks (allocator-agnostic, no Arrow types in the signature). DataFusion uses this to mirror datafusion-pool resize to the Rust MemoryPool via df_set_memory_pool_limit (R1). * Derive native.allocator.root.limit from node.native_memory.limit * 0.8 when unset, leaving 20% headroom for non-Arrow native usage that admission control still needs to throttle on (R5). Reconciles with opensearch-project#21191's admission control. * Implement CircuitBreakerPlugin: register a native_arrow breaker with Durability.PERMANENT so root usage is reflected in _nodes/stats?breaker and the parent breaker rolls up off-heap pressure. Allocator periodically syncs root.getAllocatedMemory into the breaker counter via the rebalance hook. * Probe df_set_spill_limit at NativeBridge static init. When the symbol is present, datafusion.spill_memory_limit_bytes becomes Dynamic and a listener calls the FFM symbol on resize. When absent, the setting stays NodeScope-only and OpenSearch rejects runtime PUTs cleanly. Lights up automatically once the upstream datafusion crate carrying df_set_spill_limit is picked up here. * Delete dead DataFusionService.rootAllocator field and newChildAllocator() method (no production callers). * Call rebalance() once after pool creation in createComponents so pools reach their max capacity even when the rebalancer is disabled by default. Without this, pools sit at min=0 and silently fail allocations. * Add ensureForTesting helper for unit tests that bring up consumers without the plugin lifecycle. Tests cover listener fan-out, multi-listener invocation, exception isolation, breaker sync via rebalance, AC-derived default resolution. Signed-off-by: Gaurav Singh <gauravsg@amazon.com>
gaurav-amz
force-pushed
the
wire-unified-allocator
branch
from
May 18, 2026 08:43
4cdd366 to
71b9e9f
Compare
Bukhtawar
force-pushed
the
unified-arrow-allocator
branch
6 times, most recently
from
May 18, 2026 15:38
9762df2 to
a40258a
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Stacked on top of opensearch-project#21703. Wires every existing Arrow consumer in OpenSearch to obtain memory through the unified
ArrowNativeAllocatorintroduced there. Today three independentRootAllocatorinstances exist that don't see each other; this PR collapses them into one node-wide hierarchy with elastic budgeting.What changes
Per consumer:
arrow-flight-rpc— Replace staticRootAllocator(Long.MAX_VALUE)with a thin shim over theflightpool. All public-API callers unchanged.analytics-engine—QueryContext(per-query, 256 MB cap) andAnalyticsSearchService(service lifetime) carve children from thequerypool.analytics-backend-datafusion— Delete deadDataFusionService.rootAllocatorfield andnewChildAllocator()method. Register aNativeAllocatorListenerthat mirrorsdatafusion-pool resize to the RustMemoryPoolviadf_set_memory_pool_limit(R1).Cross-cutting:
NativeAllocatorListenerSPI inarrow-spifor pool-resize callbacks (allocator-agnostic).POOL_QUERYandPOOL_DATAFUSIONconstants + min/max settings.native.allocator.root.limitdefaults tonode.native_memory.limit × 0.8when unset, leaving 20% headroom for non-Arrow native (jemalloc, Netty direct, JNI). Reconciles with PR native memory based admission control opensearch-project/OpenSearch#21191's admission control.CircuitBreakerPluginintegration —native_arrowbreaker,Durability.PERMANENT. Allocator periodically syncs root usage into the breaker via the rebalance hook, surfacing native usage in_nodes/stats?breakerand rolling up into the parent breaker.datafusion.spill_memory_limit_bytesbecomesDynamiconly whendf_set_spill_limitis exported by the loaded datafusion library (probed at static init). Otherwise staysNodeScope-only and OpenSearch rejects runtime PUTs cleanly. Lights up automatically when the upstream datafusion crate carrying the symbol is picked up here.extendedPlugins = ['native-allocator-arrow']declared on all migrated plugins.rebalance()after pool creation increateComponentsso pools reachmaxcapacity even when the rebalancer is disabled by default. Without this, pools sit atmin=0and silently fail allocations — applies to parquet'singestpool too. Worth folding into Unified native allocator framework to track and monitor Arrow allocations opensearch-project/OpenSearch#21703 itself.Test plan
:libs:opensearch-arrow-spi:test— passes:plugins:native-allocator-arrow:test— passes (15 tests including 5 new for listener fan-out, breaker sync, AC-derived default):plugins:arrow-flight-rpc:test— passes (12 tests in FlightClientChannelTests, all green):sandbox:plugins:analytics-engine:test— passes:sandbox:plugins:analytics-backend-datafusion:test— passes (sandbox tests requiring native lib excluded; native lib not built in this checkout)Notes for reviewers
This PR targets
Bukhtawar:unified-arrow-allocatordirectly. After opensearch-project#21703 merges to main, this rebases onto main and the diff stays the same.Overlap with opensearch-project#21465 (bowen's
arrow-base): there is meaningful overlap ifarrow-baselands first —ArrowAllocatorProviderwould be deleted there, my shim becomes redundant, andArrowAllocatorServicebecomes the canonical injection target. Coordinating with bowen separately on whetherArrowAllocatorServiceshould expose pool semantics so this work slots in cleanly.Out of scope:
parquet.max_native_allocationanddatafusion.memory_pool_limit_bytes(settings still work; aliases are a follow-up)._nodes/stats?native_allocator) — today thenative_arrowbreaker gives aggregate visibility.df_set_spill_limitandRwLock<Arc<RuntimeEnv>>work — separate PR on the datafusion side.