Add configurable coordinator buffer limit for per-query Arrow allocator

[bowenlan-amzn]

Summary

• Replace hardcoded 256MB per-query Arrow allocator limit with dynamic cluster setting analytics.coordinator.buffer_limit
• Add debug/warn logging at allocator close to detect memory leaks (WARN if bytes > 0 at close)

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit bd85809)

Here are some key observations to aid the review process:

🧪 No relevant tests

🔒 No security concerns identified

📝 TODO sections sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [78]: close on shutdown sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [149]: accept request type from FrontEnd including cancelAfterTimeInterval sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java [30]: make configurable via cluster setting sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryExecution.java [121]: Re-evaluate this per query child allocator

🔀 No multiple PR themes

⚡ Recommended focus areas for review Allocator Leak If QueryContext constructor throws after queryAllocator is created but before ownsAllocator=true is set, the allocator is closed on line 169. However, if the constructor succeeds but an exception occurs between lines 167-171 (after the try-catch), the allocator is never closed because context.close() is only called via taskManager.unregister(queryTask) in the batchesListener, which won't run if an exception is thrown before the listener is attached. This leaves the child allocator open indefinitely. final BufferAllocator queryAllocator; final boolean ownsAllocator; if (perQueryBufferLimit <= 0) { queryAllocator = coordinatorAllocator; ownsAllocator = false; } else { queryAllocator = coordinatorAllocator.newChildAllocator("query-" + dag.queryId(), 0, perQueryBufferLimit); ownsAllocator = true; } logger.debug("[query-{}] Arrow allocator created, limit={}B", dag.queryId(), perQueryBufferLimit); final QueryContext context; try { context = new QueryContext(dag, searchExecutor, queryTask, queryAllocator, ownsAllocator); } catch (Exception e) { if (ownsAllocator) queryAllocator.close(); throw e; } Race Condition The close() method checks ownsAllocator and closes the allocator without synchronization, but ownsAllocator is a final field read outside the synchronized block. If close() is called concurrently from multiple threads, both threads could pass the if (closed) return check before either sets closed=true, resulting in allocator.close() being called twice. Arrow allocators typically throw or log errors on double-close. public void close() { synchronized (this) { if (closed) return; closed = true; if (ownsAllocator) { allocator.close(); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to bd85809

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Unregister task on context creation failure The taskManager.unregister(queryTask) call in the batchesListener won't execute if QueryContext construction throws. This leaves the task registered indefinitely. Wrap the entire query setup (including task registration) in a try-catch that ensures taskManager.unregister(queryTask) is called on any failure before the listener chain is established. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [155-171] final BufferAllocator queryAllocator; final boolean ownsAllocator; if (perQueryBufferLimit <= 0) { queryAllocator = coordinatorAllocator; ownsAllocator = false; } else { queryAllocator = coordinatorAllocator.newChildAllocator("query-" + dag.queryId(), 0, perQueryBufferLimit); ownsAllocator = true; } logger.debug("[query-{}] Arrow allocator created, limit={}B", dag.queryId(), perQueryBufferLimit); final QueryContext context; try { context = new QueryContext(dag, searchExecutor, queryTask, queryAllocator, ownsAllocator); } catch (Exception e) { if (ownsAllocator) queryAllocator.close(); + taskManager.unregister(queryTask); throw e; } Suggestion importance[1-10]: 9 __ Why: This is a critical resource leak fix. If QueryContext construction fails, the queryTask remains registered in taskManager indefinitely since the batchesListener cleanup won't execute. Adding taskManager.unregister(queryTask) to the catch block ensures proper cleanup on failure.	High
	Fix race condition in buffer limit The perQueryBufferLimit field is accessed by multiple threads (read in executeInternal, written by settings consumer) without synchronization. This creates a race condition where queries might see inconsistent buffer limits. Use AtomicLong or synchronize access to ensure thread-safe reads and writes. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [105-107] -this.perQueryBufferLimit = AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT.get(clusterService.getSettings()); +this.perQueryBufferLimit = new AtomicLong(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT.get(clusterService.getSettings())); clusterService.getClusterSettings() - .addSettingsUpdateConsumer(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT, v -> perQueryBufferLimit = v); + .addSettingsUpdateConsumer(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT, v -> perQueryBufferLimit.set(v)); Suggestion importance[1-10]: 8 __ Why: The perQueryBufferLimit field is marked volatile but is accessed by multiple threads without proper synchronization. Using AtomicLong ensures thread-safe reads and writes, preventing potential race conditions where queries might see inconsistent buffer limits during dynamic setting updates.	Medium
General	Prevent resource leak on cleanup failure If allocator.close() throws an exception, localTaskExecutor.shutdown() won't execute, causing a resource leak. Wrap each cleanup operation in its own try-catch block to ensure all resources are released even if one cleanup fails. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java [126-137] public void close() { synchronized (this) { if (closed) return; closed = true; if (ownsAllocator) { - allocator.close(); + try { + allocator.close(); + } catch (Exception e) { + // Log but continue cleanup + } } if (localTaskExecutor != null) { - localTaskExecutor.shutdown(); + try { + localTaskExecutor.shutdown(); + } catch (Exception e) { + // Log but continue cleanup + } localTaskExecutor = null; } } } Suggestion importance[1-10]: 7 __ Why: If allocator.close() throws an exception, localTaskExecutor.shutdown() won't execute, causing a resource leak. Wrapping each cleanup operation in try-catch blocks ensures all resources are released even if one cleanup fails, improving robustness.	Medium

---

Previous suggestions

Suggestions up to commit 196a47d

Category	Suggestion	Impact
Possible issue	Unregister task on early failure The taskManager.unregister(queryTask) call in the batchesListener may not execute if an exception occurs before the listener is attached. This leaves the task registered indefinitely. Wrap the entire setup in a try-catch block and ensure taskManager.unregister(queryTask) is called on any failure path before the listener is established. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [155-171] final BufferAllocator queryAllocator; final boolean ownsAllocator; if (perQueryBufferLimit <= 0) { queryAllocator = coordinatorAllocator; ownsAllocator = false; } else { queryAllocator = coordinatorAllocator.newChildAllocator("query-" + dag.queryId(), 0, perQueryBufferLimit); ownsAllocator = true; } logger.debug("[query-{}] Arrow allocator created, limit={}B", dag.queryId(), perQueryBufferLimit); final QueryContext context; try { context = new QueryContext(dag, searchExecutor, queryTask, queryAllocator, ownsAllocator); } catch (Exception e) { if (ownsAllocator) queryAllocator.close(); + taskManager.unregister(queryTask); throw e; } Suggestion importance[1-10]: 9 __ Why: Critical resource leak fix. If QueryContext construction fails, the queryTask remains registered in taskManager indefinitely since the batchesListener cleanup won't execute. Adding taskManager.unregister(queryTask) to the catch block ensures proper cleanup on all failure paths.	High
General	Ensure all resources close If allocator.close() throws an exception, localTaskExecutor.shutdown() will not execute, leaving the executor running. Wrap each cleanup operation in a try-catch block to ensure all resources are released even if one fails. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java [126-138] public void close() { synchronized (this) { if (closed) return; closed = true; if (ownsAllocator) { - allocator.close(); + try { + allocator.close(); + } catch (Exception e) { + // Log but continue cleanup + } } if (localTaskExecutor != null) { - localTaskExecutor.shutdown(); + try { + localTaskExecutor.shutdown(); + } catch (Exception e) { + // Log but continue cleanup + } localTaskExecutor = null; } } } Suggestion importance[1-10]: 7 __ Why: Important robustness improvement. If allocator.close() throws an exception, localTaskExecutor.shutdown() won't execute, leaving the executor running. Wrapping each cleanup operation in try-catch ensures all resources are released even if one fails, preventing resource leaks.	Medium
	Add upper bound validation The minimum value of 0L allows disabling per-query limits, but negative values should be explicitly rejected. Consider using Setting.Property.NonNegative or adjusting the minimum to prevent potential misconfigurations that could lead to unexpected behavior. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java [66-72] public static final Setting<Long> COORDINATOR_BUFFER_LIMIT = Setting.longSetting( "analytics.coordinator.buffer_limit", 256L * 1024 * 1024, 0L, + Long.MAX_VALUE, Setting.Property.NodeScope, Setting.Property.Dynamic ); Suggestion importance[1-10]: 3 __ Why: Minor improvement. The Setting.longSetting already validates the minimum value of 0L, and the suggestion to add Long.MAX_VALUE as an upper bound is reasonable but has minimal practical impact since such large values are unlikely to be configured and would fail naturally at allocation time.	Low

Suggestions up to commit 037df7a

Category	Suggestion	Impact
General	Unregister task on context creation failure If QueryContext construction fails, the queryTask remains registered in the TaskManager and will never be unregistered, causing a resource leak. Ensure taskManager.unregister(queryTask) is called in the catch block before rethrowing the exception. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [165-171] final QueryContext context; try { context = new QueryContext(dag, searchExecutor, queryTask, queryAllocator, ownsAllocator); } catch (Exception e) { if (ownsAllocator) queryAllocator.close(); + taskManager.unregister(queryTask); throw e; } Suggestion importance[1-10]: 9 __ Why: This identifies a critical resource leak where queryTask remains registered in TaskManager if QueryContext construction fails. The task would never be unregistered since the ActionListener.runAfter cleanup won't execute. This is a significant bug that could accumulate leaked task registrations over time.	High
	Always create child allocator per query When perQueryBufferLimit is zero, the code shares the coordinatorAllocator across all queries without isolation. This can lead to memory contention and makes it impossible to track per-query memory usage. Consider always creating a child allocator with appropriate limits to maintain query isolation and observability. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [155-163] final BufferAllocator queryAllocator; final boolean ownsAllocator; if (perQueryBufferLimit <= 0) { - queryAllocator = coordinatorAllocator; - ownsAllocator = false; + queryAllocator = coordinatorAllocator.newChildAllocator("query-" + dag.queryId(), 0, Long.MAX_VALUE); + ownsAllocator = true; } else { queryAllocator = coordinatorAllocator.newChildAllocator("query-" + dag.queryId(), 0, perQueryBufferLimit); ownsAllocator = true; } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies that sharing coordinatorAllocator when perQueryBufferLimit <= 0 prevents per-query memory tracking and isolation. Creating a child allocator with Long.MAX_VALUE limit maintains isolation while allowing unlimited memory, which is a valid improvement for observability and resource management.	Medium
	Use AtomicLong for thread-safe limit The perQueryBufferLimit field is marked volatile but is read in executeInternal without synchronization while being updated by the settings consumer. This creates a race condition where a query might observe an inconsistent limit value during concurrent updates. Use AtomicLong for thread-safe reads and updates. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [105-107] -this.perQueryBufferLimit = AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT.get(clusterService.getSettings()); +private final AtomicLong perQueryBufferLimit = new AtomicLong(); +// In constructor: +this.perQueryBufferLimit.set(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT.get(clusterService.getSettings())); clusterService.getClusterSettings() - .addSettingsUpdateConsumer(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT, v -> perQueryBufferLimit = v); + .addSettingsUpdateConsumer(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT, perQueryBufferLimit::set); +// In executeInternal, use: perQueryBufferLimit.get() Suggestion importance[1-10]: 6 __ Why: While volatile provides visibility guarantees for long reads/writes on most platforms, using AtomicLong is more explicit and portable for thread-safe operations. However, the current implementation is likely safe in practice since volatile long operations are atomic on 64-bit JVMs, making this a minor improvement rather than a critical fix.	Low

Suggestions up to commit abefa15

Category	Suggestion	Impact
Possible issue	Prevent task manager leak on failure The exception handler only closes queryAllocator but doesn't unregister the queryTask from the task manager. If QueryContext construction fails, the task remains registered, causing a resource leak. Ensure taskManager.unregister(queryTask) is called in the catch block. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [162-168] final QueryContext context; try { context = new QueryContext(dag, searchExecutor, queryTask, queryAllocator, ownsAllocator); } catch (Exception e) { if (ownsAllocator) queryAllocator.close(); + taskManager.unregister(queryTask); throw e; } Suggestion importance[1-10]: 9 __ Why: This identifies a critical resource leak where queryTask remains registered in the task manager if QueryContext construction fails. The task was registered at line 147 but won't be unregistered on failure, causing a resource leak. This is a significant bug that should be fixed.	High
General	Adjust minimum buffer limit value The minimum value of 0L for COORDINATOR_BUFFER_LIMIT may cause issues since a zero-byte limit would prevent any allocations. Consider setting the minimum to -1L to explicitly support "unlimited" semantics, or to a small positive value if zero-byte limits are not intended. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java [67-73] public static final Setting<Long> COORDINATOR_BUFFER_LIMIT = Setting.longSetting( "analytics.coordinator.buffer_limit", 256L * 1024 * 1024, - 0L, + -1L, Setting.Property.NodeScope, Setting.Property.Dynamic ); Suggestion importance[1-10]: 6 __ Why: Changing the minimum from 0L to -1L better supports "unlimited" semantics and prevents potential confusion with zero-byte limits. This is a reasonable improvement that aligns with the code logic in DefaultPlanExecutor where <= 0 triggers unlimited mode, though the current implementation is still functional.	Low
	Clarify unlimited buffer condition The condition perQueryBufferLimit <= 0 allows zero as a valid value to disable per-query limits, but zero is semantically ambiguous. Consider using perQueryBufferLimit < 0 to explicitly indicate "unlimited" mode, making the intent clearer and preventing accidental zero-limit configurations. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [154-160] final BufferAllocator queryAllocator; final boolean ownsAllocator; -if (perQueryBufferLimit <= 0) { +if (perQueryBufferLimit < 0) { queryAllocator = coordinatorAllocator; ownsAllocator = false; } else { queryAllocator = coordinatorAllocator.newChildAllocator("query-" + dag.queryId(), 0, perQueryBufferLimit); ownsAllocator = true; } Suggestion importance[1-10]: 5 __ Why: The suggestion to use < 0 instead of <= 0 improves semantic clarity by explicitly distinguishing "unlimited" mode from a zero-byte limit. However, the current implementation with <= 0 is functionally correct and the impact is moderate since it's primarily a code clarity improvement.	Low

Suggestions up to commit 4f0a067

Category	Suggestion	Impact
Possible issue	Set reasonable minimum buffer limit The minimum value of 0L for buffer limit is problematic as it would prevent any Arrow allocations from succeeding. Consider setting a reasonable minimum (e.g., 1MB) to ensure queries can allocate at least minimal buffers for operation. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java [66-72] public static final Setting<Long> COORDINATOR_BUFFER_LIMIT = Setting.longSetting( "analytics.coordinator.buffer_limit", 256L * 1024 * 1024, - 0L, + 1024L * 1024, Setting.Property.NodeScope, Setting.Property.Dynamic ); Suggestion importance[1-10]: 8 __ Why: Setting a minimum of 0L for COORDINATOR_BUFFER_LIMIT would prevent Arrow allocations from succeeding, causing query failures. The suggestion to use a reasonable minimum like 1MB is valid and addresses a potential runtime issue.	Medium
General	Validate dynamic buffer limit updates The perQueryBufferLimit field is marked volatile but the settings update consumer directly assigns to it without validation. If a dynamic update sets an invalid value (e.g., 0), in-flight queries could fail. Consider adding validation in the consumer or using a validator in the Setting definition. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [102-103] this.perQueryBufferLimit = AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT.get(clusterService.getSettings()); -clusterService.getClusterSettings().addSettingsUpdateConsumer(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT, v -> perQueryBufferLimit = v); +clusterService.getClusterSettings().addSettingsUpdateConsumer(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT, v -> { + if (v <= 0) { + logger.warn("Invalid buffer limit {} ignored, keeping current value {}", v, perQueryBufferLimit); + return; + } + perQueryBufferLimit = v; +}); Suggestion importance[1-10]: 7 __ Why: While the suggestion to validate dynamic updates is reasonable, the minimum value validation should ideally be enforced at the Setting definition level (as noted in suggestion 1). Adding runtime validation here provides defense-in-depth but doesn't fully address the root cause.	Medium

Suggestions up to commit f877ed5

Category	Suggestion	Impact
Possible issue	Set reasonable minimum buffer limit The minimum value of 0L for buffer limit is problematic as it would prevent any Arrow allocations. Consider setting a reasonable minimum (e.g., 1MB) to ensure the allocator can function properly and avoid runtime allocation failures. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java [66-72] public static final Setting<Long> COORDINATOR_BUFFER_LIMIT = Setting.longSetting( "analytics.coordinator.buffer_limit", 256L * 1024 * 1024, - 0L, + 1024L * 1024, Setting.Property.NodeScope, Setting.Property.Dynamic ); Suggestion importance[1-10]: 8 __ Why: Setting a minimum of 0L for COORDINATOR_BUFFER_LIMIT could cause Arrow allocation failures. The suggestion to use a reasonable minimum like 1MB is valid and prevents runtime issues.	Medium
General	Validate buffer limit updates The perQueryBufferLimit field is marked volatile but the settings update consumer directly assigns to it without validation. If a dynamic update sets an invalid value (e.g., 0), in-flight queries could fail. Consider adding validation in the consumer or using a validator in the Setting definition. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java [102-103] this.perQueryBufferLimit = AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT.get(clusterService.getSettings()); -clusterService.getClusterSettings().addSettingsUpdateConsumer(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT, v -> perQueryBufferLimit = v); +clusterService.getClusterSettings().addSettingsUpdateConsumer(AnalyticsPlugin.COORDINATOR_BUFFER_LIMIT, v -> { + if (v <= 0) { + logger.warn("Invalid buffer limit {} ignored, keeping current value {}", v, perQueryBufferLimit); + return; + } + perQueryBufferLimit = v; +}); Suggestion importance[1-10]: 7 __ Why: While the suggestion to validate dynamic updates is reasonable, the Setting definition already enforces a minimum of 0L. However, adding validation in the consumer provides defense-in-depth against invalid values affecting in-flight queries.	Medium

[github-actions]

Persistent review updated to latest commit f877ed5

[github-actions]

Persistent review updated to latest commit 4f0a067

[github-actions]

❌ Gradle check result for 4f0a067: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

PR Code Analyzer ❗

AI-powered 'Code-Diff-Analyzer' found issues on commit abefa15.

Path	Line	Severity	Description
sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java	118	low	Coordinator allocator is created with Long.MAX_VALUE as its memory limit, meaning there is no effective cap on coordinator-level Arrow memory. While per-query limits are applied as child allocators, an unbounded parent could allow a coordinated set of queries to exhaust JVM heap. This appears intentional (the coordinator acts as a shared root), but the absence of any ceiling is worth a deliberate review.
sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java	142	low	A static RootAllocator (TEST_ROOT) with Long.MAX_VALUE limit is initialized at class-load time in production source. It is only reachable through test-only factory methods, but it is never explicitly closed, creating a permanent resource that could mask Arrow memory leaks in test suites and, if a class-loader edge case causes production code to invoke a forTest path, would provide an uncapped allocation pool.

The table above displays the top 10 most important findings.

Total: 2 | Critical: 0 | High: 0 | Medium: 0 | Low: 2

---

Pull Requests Author(s): Please update your Pull Request according to the report above.

Repository Maintainer(s): You can bypass diff analyzer by adding label skip-diff-analyzer after reviewing the changes carefully, then re-run failed actions. To re-enable the analyzer, remove the label, then re-run all actions.

---

⚠️ Note: The Code-Diff-Analyzer helps protect against potentially harmful code patterns. Please ensure you have thoroughly reviewed the changes beforehand.

Thanks.

[github-actions]

Persistent review updated to latest commit abefa15

[github-actions]

❌ Gradle check result for abefa15: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit 037df7a

[github-actions]

Persistent review updated to latest commit 196a47d

[github-actions]

Persistent review updated to latest commit bd85809

[github-actions]

✅ Gradle check result for bd85809: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.52%. Comparing base (8f2d058) to head (bd85809).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21726      +/-   ##
============================================
+ Coverage     73.46%   73.52%   +0.05%     
+ Complexity    74825    74818       -7     
============================================
  Files          5997     5997              
  Lines        339688   339688              
  Branches      48961    48961              
============================================
+ Hits         249558   249745     +187     
+ Misses        70272    70023     -249     
- Partials      19858    19920      +62     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.