diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsCompositeListenerFactory.java b/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsCompositeListenerFactory.java index db487bf945889..842a0e2dc61fe 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsCompositeListenerFactory.java +++ b/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsCompositeListenerFactory.java @@ -69,8 +69,8 @@ public SearchRequestOperationsListener.CompositeListener buildCompositeListener( final SearchRequestOperationsListener... perRequestListeners ) { final List searchListenersList = Stream.concat( - searchRequestListenersList.stream(), - Arrays.stream(perRequestListeners) + Arrays.stream(perRequestListeners), + searchRequestListenersList.stream() ) .filter((searchRequestOperationsListener -> searchRequestOperationsListener.isEnabled(searchRequest))) .collect(Collectors.toList()); diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java index 6e380775355a2..423ce249fb99a 100644 --- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java +++ b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java @@ -86,6 +86,7 @@ import org.opensearch.search.profile.ProfileShardResult; import org.opensearch.search.profile.SearchProfileShardResults; import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.SearchTaskResourceOperationsListener; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskResourceTrackingService; import org.opensearch.telemetry.metrics.MetricsRegistry; @@ -453,7 +454,8 @@ private void executeRequest( requestOperationsListeners = searchRequestOperationsCompositeListenerFactory.buildCompositeListener( originalSearchRequest, logger, - TraceableSearchRequestOperationsListener.create(tracer, requestSpan) + TraceableSearchRequestOperationsListener.create(tracer, requestSpan), + new SearchTaskResourceOperationsListener(taskResourceTrackingService) ); SearchRequestContext searchRequestContext = new SearchRequestContext( requestOperationsListeners, diff --git a/server/src/main/java/org/opensearch/tasks/SearchTaskResourceOperationsListener.java b/server/src/main/java/org/opensearch/tasks/SearchTaskResourceOperationsListener.java new file mode 100644 index 0000000000000..56686388f7aa1 --- /dev/null +++ b/server/src/main/java/org/opensearch/tasks/SearchTaskResourceOperationsListener.java @@ -0,0 +1,43 @@ +/* + * 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.tasks; + +import org.opensearch.action.search.SearchPhaseContext; +import org.opensearch.action.search.SearchRequestContext; +import org.opensearch.action.search.SearchRequestOperationsListener; + +/** + * SearchTaskResourceOperationsListener subscriber for operations on search tasks resource usages + * + * @opensearch.internal + */ +public final class SearchTaskResourceOperationsListener extends SearchRequestOperationsListener { + private final TaskResourceTrackingService taskResourceTrackingService; + + public SearchTaskResourceOperationsListener(TaskResourceTrackingService taskResourceTrackingService) { + this.taskResourceTrackingService = taskResourceTrackingService; + } + + @Override + protected void onPhaseStart(SearchPhaseContext context) {} + + @Override + protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {} + + @Override + protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {} + + @Override + public void onRequestStart(SearchRequestContext searchRequestContext) {} + + @Override + public void onRequestEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) { + taskResourceTrackingService.refreshResourceStats(context.getTask()); + } +}