Improve/native memory admission control's precision addresses multiple bug fixes

plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java

@@ -179,6 +179,14 @@ public void close() {
             }
         });
         pools.clear();
+        // Close any remaining child allocators (e.g., ad-hoc children created via ArrowAllocatorService)
+        for (BufferAllocator child : new ArrayList<>(root.getChildAllocators())) {
+            try {
+                child.close();
+            } catch (Exception e) {
+                // best-effort — log but don't block shutdown
+            }
+        }
         root.close();
         INSTANCE = null;
     }

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java

@@ -9,7 +9,10 @@
 package org.opensearch.analytics.exec;
 
 import org.apache.arrow.memory.BufferAllocator;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.opensearch.analytics.backend.AnalyticsOperationListener;
+import org.opensearch.analytics.backend.EngineResultBatch;
 import org.opensearch.analytics.backend.EngineResultStream;
 import org.opensearch.analytics.backend.SearchExecEngine;
 import org.opensearch.analytics.backend.ShardScanExecutionContext;
@@ -34,8 +37,10 @@
 import org.opensearch.tasks.TaskResourceTrackingService;
 
 import java.io.IOException;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Executor;
 
 /**
  * Data-node service that executes plan fragments against local shards.
@@ -56,6 +61,8 @@
  */
 public class AnalyticsSearchService implements AutoCloseable {
 
+    private static final Logger LOGGER = LogManager.getLogger(AnalyticsSearchService.class);
+
     private final Map<String, AnalyticsSearchBackendPlugin> backends;
     private final AnalyticsOperationListener listener;
     private final NamedWriteableRegistry namedWriteableRegistry;
@@ -108,6 +115,45 @@ public FragmentResources executeFragmentStreaming(FragmentExecutionRequest reque
         }
     }
 
+    /**
+     * Async variant that forks fragment execution onto the given executor and streams
+     * batches back through the channel. The transport thread returns immediately.
+     */
+    public void executeFragmentStreamingAsync(
+        FragmentExecutionRequest request,
+        IndexShard shard,
+        AnalyticsShardTask task,
+        StreamingFragmentResponseHandler responseHandler,
+        Executor executor
+    ) {
+        try {
+            executor.execute(() -> {
+                try (FragmentResources ctx = executeFragmentStreaming(request, shard, task)) {
+                    Iterator<EngineResultBatch> it = ctx.stream().iterator();
+                    while (it.hasNext()) {
+                        responseHandler.onBatch(it.next());
+                    }
+                    responseHandler.onComplete();
+                } catch (Exception e) {
+                    responseHandler.onFailure(e);
+                }
+            });
+        } catch (Exception e) {
+            responseHandler.onFailure(e);
+        }
+    }
+
+    /**
+     * Callback interface for async fragment streaming results.
+     */
+    public interface StreamingFragmentResponseHandler {
+        void onBatch(EngineResultBatch batch) throws Exception;
+
+        void onComplete();
+
+        void onFailure(Exception e);
+    }
+
     private FragmentResources startFragment(FragmentExecutionRequest request, ResolvedFragment resolved, IndexShard shard, Task task)
         throws IOException {
         GatedCloseable<Reader> gatedReader = resolved.readerProvider.acquireReader();
@@ -163,6 +209,15 @@ public void trackEnd(long threadId) {
             stream = engine.execute(ctx);
             return new FragmentResources(gatedReader, engine, stream, trackerCleanup);
         } catch (Exception e) {
+            LOGGER.error(
+                () -> new org.apache.logging.log4j.message.ParameterizedMessage(
+                    "startFragment failed [queryId={}, stageId={}, shardId={}]",
+                    resolved.queryId,
+                    resolved.stageId,
+                    resolved.shardIdStr
+                ),
+                e
+            );
             try {
                 new FragmentResources(gatedReader, engine, stream, trackerCleanup).close();
             } catch (Exception suppressed) {

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java

@@ -35,7 +35,6 @@
 import org.opensearch.transport.stream.StreamTransportResponse;
 
 import java.io.IOException;
-import java.util.Iterator;
 
 /**
  * Stateless transport dispatch component for fragment requests. Owns the
@@ -88,21 +87,33 @@ private static void registerStreamingFragmentHandler(
             FragmentExecutionRequest::new,
             (request, channel, task) -> {
                 IndexShard shard = indicesService.indexServiceSafe(request.getShardId().getIndex()).getShard(request.getShardId().id());
-                try (FragmentResources ctx = searchService.executeFragmentStreaming(request, shard, (AnalyticsShardTask) task)) {
-                    Iterator<EngineResultBatch> it = ctx.stream().iterator();
-                    while (it.hasNext()) {
-                        EngineResultBatch batch = it.next();
-                        channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot()));
-                    }
-                    channel.completeStream();
-                } catch (StreamException e) {
-                    if (e.getErrorCode() != StreamErrorCode.CANCELLED) {
-                        channel.sendResponse(e);
-                    }
-                    // CANCELLED: channel already torn down — exit silently
-                } catch (Exception e) {
-                    channel.sendResponse(e);
-                }
+                searchService.executeFragmentStreamingAsync(
+                    request,
+                    shard,
+                    (AnalyticsShardTask) task,
+                    new AnalyticsSearchService.StreamingFragmentResponseHandler() {
+                        @Override
+                        public void onBatch(EngineResultBatch batch) throws Exception {
+                            channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot()));
+                        }
+
+                        @Override
+                        public void onComplete() {
+                            channel.completeStream();
+                        }
+
+                        @Override
+                        public void onFailure(Exception e) {
+                            if (e instanceof StreamException se && se.getErrorCode() == StreamErrorCode.CANCELLED) {
+                                return;
+                            }
+                            try {
+                                channel.sendResponse(e);
+                            } catch (Exception ignored) {}
+                        }
+                    },
+                    transportService.getThreadPool().executor(ThreadPool.Names.SEARCH)
+                );
             }
         );
     }

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/AnalyticsSearchTransportServiceTests.java

@@ -0,0 +1,51 @@
+/*
+ * 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.analytics.exec;
+
+import org.opensearch.analytics.exec.action.FragmentExecutionAction;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.indices.IndicesService;
+import org.opensearch.tasks.TaskResourceTrackingService;
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.StreamTransportService;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+/**
+ * Verifies that the fragment execution handler is registered with SAME executor
+ * (admission control runs on transport thread) but forks execution internally
+ * via {@link AnalyticsSearchService#executeFragmentStreamingAsync}.
+ */
+public class AnalyticsSearchTransportServiceTests extends OpenSearchTestCase {
+
+    public void testFragmentHandlerRegisteredWithSameExecutor() {
+        StreamTransportService transportService = mock(StreamTransportService.class);
+        AnalyticsSearchService searchService = mock(AnalyticsSearchService.class);
+        IndicesService indicesService = mock(IndicesService.class);
+        ClusterService clusterService = mock(ClusterService.class);
+        TaskResourceTrackingService taskResourceTrackingService = mock(TaskResourceTrackingService.class);
+
+        new AnalyticsSearchTransportService(transportService, clusterService, searchService, indicesService, taskResourceTrackingService);
+
+        verify(transportService).registerRequestHandler(
+            eq(FragmentExecutionAction.NAME),
+            eq(ThreadPool.Names.SAME),
+            anyBoolean(),
+            anyBoolean(),
+            any(),
+            any(),
+            any()
+        );
+    }
+}

server/src/internalClusterTest/java/org/opensearch/cluster/ClusterInfoServiceIT.java

@@ -261,7 +261,7 @@ public void testClusterInfoServiceCollectsNodeResourceStatsInformation() throws
             .put(ResourceTrackerSettings.GLOBAL_JVM_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueSeconds(1))
             .put(ResourceTrackerSettings.GLOBAL_IO_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueMillis(5000))
             .put(ResourceTrackerSettings.GLOBAL_CPU_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueSeconds(1))
-            .put(ResourceTrackerSettings.GLOBAL_NATIVE_MEMORY_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueMillis(1000));
+            .put(ResourceTrackerSettings.GLOBAL_NATIVE_MEMORY_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueSeconds(5));
 
         internalCluster().startClusterManagerOnlyNode(settingsBuilder.build());
         internalCluster().startWarmOnlyNodes(1, settingsBuilder.build());
@@ -293,7 +293,7 @@ public void testClusterInfoServiceInformationClearOnError() throws InterruptedEx
                 .put(ResourceTrackerSettings.GLOBAL_JVM_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueMillis(500))
                 .put(ResourceTrackerSettings.GLOBAL_IO_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueMillis(5000))
                 .put(ResourceTrackerSettings.GLOBAL_CPU_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueMillis(500))
-                .put(ResourceTrackerSettings.GLOBAL_NATIVE_MEMORY_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueMillis(500))
+                .put(ResourceTrackerSettings.GLOBAL_NATIVE_MEMORY_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueSeconds(5))
                 .build()
         );
         prepareCreate("test").setSettings(Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)).get();