And more stuff

Signed-off-by: Atri Sharma <atrisharma@Atris-Mac-Studio.local>

Author: Atri Sharma

server/src/main/java/org/opensearch/action/search/SearchPhaseController.java

@@ -246,14 +246,7 @@ static TopDocs mergeTopDocs(Collection<TopDocs> results, int topN, int from) {
     }
 
     static void setShardIndex(TopDocs topDocs, int shardIndex) {
-        // Idempotent assignment: in streaming flows partial reductions may touch the same TopDocs more than once.
-        if (topDocs.scoreDocs.length == 0) {
-            return;
-        }
-        if (topDocs.scoreDocs[0].shardIndex != -1) {
-            // Already set by a previous pass; avoid reassigning to prevent assertion failures
-            return;
-        }
+        assert topDocs.scoreDocs.length == 0 || topDocs.scoreDocs[0].shardIndex == -1 : "shardIndex is already set";
         for (ScoreDoc doc : topDocs.scoreDocs) {
             doc.shardIndex = shardIndex;
         }

server/src/main/java/org/opensearch/action/search/StreamQueryPhaseResultConsumer.java

@@ -48,7 +48,6 @@ public StreamQueryPhaseResultConsumer(
 
     @Override
     int getBatchReduceSize(int requestBatchedReduceSize, int minBatchReduceSize) {
-        // Reduce immediately for fastest TTFB
         return Math.min(requestBatchedReduceSize, 1);
     }
 }

server/src/main/java/org/opensearch/action/search/StreamSearchQueryThenFetchAsyncAction.java

@@ -28,8 +28,6 @@
  */
 public class StreamSearchQueryThenFetchAsyncAction extends SearchQueryThenFetchAsyncAction {
 
-    private final Logger logger;
-
     StreamSearchQueryThenFetchAsyncAction(
         Logger logger,
         SearchTransportService searchTransportService,
@@ -70,7 +68,6 @@ public class StreamSearchQueryThenFetchAsyncAction extends SearchQueryThenFetchA
             searchRequestContext,
             tracer
         );
-        this.logger = logger;
     }
 
     @Override
@@ -87,10 +84,7 @@ SearchActionListener<SearchPhaseResult> createShardActionListener(
             @Override
             protected void innerOnStreamResponse(SearchPhaseResult result) {
                 try {
-                    if (getLogger().isTraceEnabled()) {
-                        getLogger().trace("coordinator received partial from shard {}", shard);
-                    }
-                    onStreamResult(result, shardIt, () -> successfulStreamExecution());
+                    onStreamResult(result, shardIt, () -> {});
                 } finally {
                     executeNext(pendingExecutions, thread);
                 }
@@ -99,9 +93,6 @@ protected void innerOnStreamResponse(SearchPhaseResult result) {
             @Override
             protected void innerOnCompleteResponse(SearchPhaseResult result) {
                 try {
-                    if (getLogger().isTraceEnabled()) {
-                        getLogger().trace("coordinator received final for shard {}", shard);
-                    }
                     onShardResult(result, shardIt);
                 } finally {
                     executeNext(pendingExecutions, thread);
@@ -111,8 +102,6 @@ protected void innerOnCompleteResponse(SearchPhaseResult result) {
             @Override
             public void onFailure(Exception t) {
                 try {
-                    // It only happens when onPhaseDone() is called and executePhaseOnShard() fails
-                    // hard with an exception.
                     if (totalOps.get() == expectedTotalOps) {
                         onPhaseFailure(phase, "The phase has failed", t);
                     } else {
@@ -131,31 +120,10 @@ public void onFailure(Exception t) {
     protected void onStreamResult(SearchPhaseResult result, SearchShardIterator shardIt, Runnable next) {
         assert result.getShardIndex() != -1 : "shard index is not set";
         assert result.getSearchShardTarget() != null : "search shard target must not be null";
-        if (getLogger().isTraceEnabled()) {
-            getLogger().trace("got streaming result from {}", result != null ? result.getSearchShardTarget() : null);
-        }
         this.setPhaseResourceUsages();
-        if (result.queryResult() != null) {
-            result.queryResult().setPartial(true);
-        }
         results.consumeResult(result, next);
     }
 
-    @Override
-    protected void onShardResult(SearchPhaseResult result, SearchShardIterator shardIt) {
-        // Trace final shard responses to diagnose coordinator sequencing.
-        if (logger.isTraceEnabled()) {
-            logger.trace(
-                "COORDINATOR: received final shard result from shard={}, target={}, totalOps={}, expectedOps={}",
-                result.getShardIndex(),
-                result.getSearchShardTarget(),
-                totalOps.get(),
-                expectedTotalOps
-            );
-        }
-        super.onShardResult(result, shardIt);
-    }
-
     @Override
     void successfulShardExecution(SearchShardIterator shardsIt) {
         final int remainingOpsOnIterator;
@@ -167,7 +135,6 @@ void successfulShardExecution(SearchShardIterator shardsIt) {
         final int xTotalOps = totalOps.addAndGet(remainingOpsOnIterator);
         if (xTotalOps == expectedTotalOps) {
             try {
-                // All final shard results have been processed; partials are not reduced.
                 onPhaseDone();
             } catch (final Exception ex) {
                 onPhaseFailure(this, "The phase has failed", ex);
@@ -180,8 +147,4 @@ void successfulShardExecution(SearchShardIterator shardsIt) {
         }
     }
 
-    private void successfulStreamExecution() {
-        // No-op.
-    }
-
 }

server/src/main/java/org/opensearch/action/search/TransportSearchAction.java

@@ -60,7 +60,6 @@
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Setting.Property;
 import org.opensearch.common.unit.TimeValue;
-import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.util.concurrent.AtomicArray;
 import org.opensearch.common.util.concurrent.CountDown;
 import org.opensearch.core.action.ActionListener;
@@ -87,7 +86,6 @@
 import org.opensearch.search.pipeline.SearchPipelineService;
 import org.opensearch.search.profile.ProfileShardResult;
 import org.opensearch.search.profile.SearchProfileShardResults;
-import org.opensearch.search.query.StreamingSearchMode;
 import org.opensearch.search.slice.SliceBuilder;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.Task;
@@ -103,6 +101,7 @@
 import org.opensearch.transport.RemoteClusterAware;
 import org.opensearch.transport.RemoteClusterService;
 import org.opensearch.transport.RemoteTransportException;
+import org.opensearch.transport.StreamTransportService;
 import org.opensearch.transport.Transport;
 import org.opensearch.transport.TransportService;
 import org.opensearch.transport.client.Client;
@@ -175,8 +174,6 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
     private final ThreadPool threadPool;
     final ClusterService clusterService;
     final SearchTransportService searchTransportService;
-    @Nullable
-    final StreamSearchTransportService streamSearchTransportService;
     private final RemoteClusterService remoteClusterService;
     final SearchPhaseController searchPhaseController;
     private final SearchService searchService;
@@ -198,10 +195,8 @@ public TransportSearchAction(
         ThreadPool threadPool,
         CircuitBreakerService circuitBreakerService,
         TransportService transportService,
-        @Nullable org.opensearch.transport.StreamTransportService streamTransportService,
         SearchService searchService,
         SearchTransportService searchTransportService,
-        @Nullable StreamSearchTransportService streamSearchTransportService,
         SearchPhaseController searchPhaseController,
         ClusterService clusterService,
         ActionFilters actionFilters,
@@ -220,13 +215,12 @@ public TransportSearchAction(
         this.circuitBreaker = circuitBreakerService.getBreaker(CircuitBreaker.REQUEST);
         this.searchPhaseController = searchPhaseController;
         this.searchTransportService = searchTransportService;
-        this.streamSearchTransportService = streamSearchTransportService;
         this.remoteClusterService = searchTransportService.getRemoteClusterService();
-        // Register request handlers - always classic; add streaming if available
-        if (streamTransportService != null && streamSearchTransportService != null) {
-            StreamSearchTransportService.registerStreamRequestHandler(streamTransportService, searchService);
+        if (transportService instanceof StreamTransportService) {
+            StreamSearchTransportService.registerStreamRequestHandler((StreamTransportService) transportService, searchService);
+        } else {
+            SearchTransportService.registerRequestHandler(transportService, searchService);
         }
-        SearchTransportService.registerRequestHandler(transportService, searchService);
         this.clusterService = clusterService;
         this.searchService = searchService;
         this.indexNameExpressionResolver = indexNameExpressionResolver;
@@ -1147,21 +1141,11 @@ private void executeSearch(
             }
         }
         final DiscoveryNodes nodes = clusterState.nodes();
-        final boolean isStreamingCandidate = (searchRequest.getStreamingSearchMode() != null)
-            && searchRequest.scroll() == null;
-        final boolean streamingEnabledSetting = clusterService.getClusterSettings().get(StreamSearchTransportService.STREAM_SEARCH_ENABLED);
-        final boolean streamingEnabledEffective = streamingEnabledSetting || FeatureFlags.isEnabled(FeatureFlags.STREAM_TRANSPORT);
-        final boolean useStreamingTransportForConnection = isStreamingCandidate
-            && streamSearchTransportService != null
-            && streamingEnabledEffective;
-        final SearchTransportService connectionTransport = useStreamingTransportForConnection
-            ? streamSearchTransportService
-            : searchTransportService;
         BiFunction<String, String, Transport.Connection> connectionLookup = buildConnectionLookup(
             searchRequest.getLocalClusterAlias(),
             nodes::get,
             remoteConnections,
-            connectionTransport::getConnection
+            searchTransportService::getConnection
         );
         final Executor asyncSearchExecutor = asyncSearchExecutor(concreteLocalIndices, clusterState);
         final boolean preFilterSearchShards = shouldPreFilterSearchShards(
@@ -1290,31 +1274,10 @@ AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction(
         SearchResponse.Clusters clusters,
         SearchRequestContext searchRequestContext
     ) {
-        // Determine if this request should use streaming transport
-        final boolean isStreamingCandidate = (searchRequest.getStreamingSearchMode() != null)
-            && searchRequest.scroll() == null;
-
-        // Check if streaming transport is actually available and enabled (cluster setting OR feature flag)
-        final boolean streamingEnabledSetting = clusterService.getClusterSettings().get(StreamSearchTransportService.STREAM_SEARCH_ENABLED);
-        final boolean streamingEnabledEffective = streamingEnabledSetting || FeatureFlags.isEnabled(FeatureFlags.STREAM_TRANSPORT);
-        final boolean canUseStreamingTransport = (streamSearchTransportService != null) && streamingEnabledEffective;
-
-        // Use streaming transport for streaming search requests
-        final boolean useStreamingTransport = isStreamingCandidate && canUseStreamingTransport;
-
         if (preFilter) {
-            if (logger.isTraceEnabled()) {
-                logger.trace(
-                    "prefilter using transport [{}] (streaming={}, enabled={}, canUse={})",
-                    ((isStreamingCandidate && canUseStreamingTransport) ? "stream" : "classic"),
-                    isStreamingCandidate,
-                    streamingEnabledSetting,
-                    canUseStreamingTransport
-                );
-            }
             return new CanMatchPreFilterSearchPhase(
                 logger,
-                (isStreamingCandidate && useStreamingTransport) ? streamSearchTransportService : searchTransportService,
+                searchTransportService,
                 connectionLookup,
                 aliasFilter,
                 concreteIndexBoosts,
@@ -1350,34 +1313,21 @@ AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction(
                 tracer
             );
         } else {
-            final boolean isStreamingRequest = (searchRequest.getStreamingSearchMode() != null);
-
-            final SearchProgressListener progressListener = task.getProgressListener();
-
             final QueryPhaseResultConsumer queryResultConsumer = searchPhaseController.newSearchPhaseResults(
                 executor,
                 circuitBreaker,
-                progressListener,
+                task.getProgressListener(),
                 searchRequest,
                 shardIterators.size(),
                 exc -> cancelTask(task, exc),
                 task::isCancelled
             );
-            if (logger.isTraceEnabled()) {
-                logger.trace(
-                    "query phase using transport [{}] (streamingRequest={}, enabled={}, canUse={})",
-                    ((isStreamingRequest && useStreamingTransport) ? "stream" : "classic"),
-                    isStreamingRequest,
-                    streamingEnabledSetting,
-                    canUseStreamingTransport
-                );
-            }
             AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction;
             switch (searchRequest.searchType()) {
                 case DFS_QUERY_THEN_FETCH:
                     searchAsyncAction = new SearchDfsQueryThenFetchAsyncAction(
                         logger,
-                        (isStreamingRequest && canUseStreamingTransport) ? streamSearchTransportService : searchTransportService,
+                        searchTransportService,
                         connectionLookup,
                         aliasFilter,
                         concreteIndexBoosts,
@@ -1397,49 +1347,26 @@ AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction(
                     );
                     break;
                 case QUERY_THEN_FETCH:
-                    if (isStreamingRequest && canUseStreamingTransport) {
-                        searchAsyncAction = new StreamSearchQueryThenFetchAsyncAction(
-                            logger,
-                            streamSearchTransportService,
-                            connectionLookup,
-                            aliasFilter,
-                            concreteIndexBoosts,
-                            indexRoutings,
-                            searchPhaseController,
-                            executor,
-                            queryResultConsumer,
-                            searchRequest,
-                            listener,
-                            shardIterators,
-                            timeProvider,
-                            clusterState,
-                            task,
-                            clusters,
-                            searchRequestContext,
-                            tracer
-                        );
-                    } else {
-                        searchAsyncAction = new SearchQueryThenFetchAsyncAction(
-                            logger,
-                            searchTransportService,
-                            connectionLookup,
-                            aliasFilter,
-                            concreteIndexBoosts,
-                            indexRoutings,
-                            searchPhaseController,
-                            executor,
-                            queryResultConsumer,
-                            searchRequest,
-                            listener,
-                            shardIterators,
-                            timeProvider,
-                            clusterState,
-                            task,
-                            clusters,
-                            searchRequestContext,
-                            tracer
-                        );
-                    }
+                    searchAsyncAction = new SearchQueryThenFetchAsyncAction(
+                        logger,
+                        searchTransportService,
+                        connectionLookup,
+                        aliasFilter,
+                        concreteIndexBoosts,
+                        indexRoutings,
+                        searchPhaseController,
+                        executor,
+                        queryResultConsumer,
+                        searchRequest,
+                        listener,
+                        shardIterators,
+                        timeProvider,
+                        clusterState,
+                        task,
+                        clusters,
+                        searchRequestContext,
+                        tracer
+                    );
                     break;
                 default:
                     throw new IllegalStateException("Unknown search type: [" + searchRequest.searchType() + "]");

server/src/main/java/org/opensearch/rest/action/search/RestSearchAction.java

@@ -151,17 +151,13 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
             parser -> parseSearchRequest(searchRequest, request, parser, client.getNamedWriteableRegistry(), setSize)
         );
 
-        // Honor streaming when cluster setting is enabled OR feature flag is enabled
-        final boolean streamingEnabledEffective = (clusterSettings != null && clusterSettings.get(STREAM_SEARCH_ENABLED))
-            || FeatureFlags.isEnabled(FeatureFlags.STREAM_TRANSPORT);
-        if (streamingEnabledEffective) {
+        if (clusterSettings != null && clusterSettings.get(STREAM_SEARCH_ENABLED)) {
             if (FeatureFlags.isEnabled(FeatureFlags.STREAM_TRANSPORT)) {
                 if (canUseStreamSearch(searchRequest)) {
                     String scoringMode = request.param("stream_scoring_mode");
                     if (scoringMode != null) {
                         searchRequest.setStreamingSearchMode(scoringMode);
                     }
-                    // StreamSearchAction should always execute with an explicit mode.
                     if (searchRequest.getStreamingSearchMode() == null) {
                         searchRequest.setStreamingSearchMode("no_scoring");
                     }
@@ -172,7 +168,7 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
                 } else {
                     logger.debug("Stream search requested but search contains unsupported aggregations. Falling back to normal search.");
                 }
-            } else {
+            } else if (searchRequest.getStreamingSearchMode() != null || request.hasParam("stream_scoring_mode")) {
                 throw new IllegalArgumentException("You need to enable stream transport first to use stream search.");
             }
         }
@@ -182,14 +178,6 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
         };
     }
 
-    /**
-     * Factory method for creating RestCancellableNodeClient instances.
-     * This method is protected to allow tests to override it and inject spyable clients.
-     *
-     * @param client the NodeClient to wrap
-     * @param httpChannel the HTTP channel for cancellation tracking
-     * @return a RestCancellableNodeClient instance
-     */
     protected RestCancellableNodeClient createRestCancellableNodeClient(NodeClient client, HttpChannel httpChannel) {
         return new RestCancellableNodeClient(client, httpChannel);
     }