Adding stream request flag to SearchRequestContext

CHANGELOG.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Changed
 - Move Randomness from server to libs/common ([#20570](https://github.com/opensearch-project/OpenSearch/pull/20570))
 - Use env variable (OPENSEARCH_FIPS_MODE) to enable opensearch to run in FIPS enforced mode instead of checking for existence of bcFIPS jars ([#20625](https://github.com/opensearch-project/OpenSearch/pull/20625))
+- Update streaming flag to use search request context ([#20530](https://github.com/opensearch-project/OpenSearch/pull/20530))
 
 ### Fixed
 - Fix flaky test failures in ShardsLimitAllocationDeciderIT ([#20375](https://github.com/opensearch-project/OpenSearch/pull/20375))

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

@@ -739,6 +739,10 @@ public final SearchRequest getRequest() {
         return request;
     }
 
+    protected SearchRequestContext getSearchRequestContext() {
+        return searchRequestContext;
+    }
+
     protected final SearchResponse buildSearchResponse(
         InternalSearchResponse internalSearchResponse,
         ShardSearchFailure[] failures,

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

@@ -43,6 +43,7 @@ public class SearchRequestContext {
     private final SearchRequest searchRequest;
     private final LinkedBlockingQueue<TaskResourceInfo> phaseResourceUsage;
     private final Supplier<TaskResourceInfo> taskResourceUsageSupplier;
+    private boolean streamingRequest;
 
     SearchRequestContext(
         final SearchRequestOperationsListener searchRequestOperationsListener,
@@ -156,6 +157,14 @@ void setSuccessfulSearchShardIndices(Set<Index> successfulSearchShardIndices) {
     public Set<Index> getSuccessfulSearchShardIndices() {
         return successfulSearchShardIndices;
     }
+
+    void setStreamingRequest(boolean streamingRequest) {
+        this.streamingRequest = streamingRequest;
+    }
+
+    public boolean isStreamingRequest() {
+        return streamingRequest;
+    }
 }
 
 enum ShardStatsFieldNames {

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

@@ -154,6 +154,9 @@ void successfulShardExecution(SearchShardIterator shardsIt) {
             try {
                 shardResultsConsumed.set(true);
                 if (streamResultsReceived.get() == streamResultsConsumeCallback.get()) {
+                    if (streamResultsReceived.get() > 0) {
+                        getSearchRequestContext().setStreamingRequest(true);
+                    }
                     getLogger().debug("Stream results consumption has called back, let shard consumption callback trigger onPhaseDone");
                     onPhaseDone();
                 } else {
@@ -180,6 +183,7 @@ private void successfulStreamExecution() {
         try {
             if (streamResultsReceived.get() == streamResultsConsumeCallback.incrementAndGet()) {
                 if (shardResultsConsumed.get()) {
+                    getSearchRequestContext().setStreamingRequest(true);
                     getLogger().debug("Stream consumption trigger onPhaseDone");
                     onPhaseDone();
                 }