Add took time to request nodes stats

Signed-off-by: David Zane <davizane@amazon.com>

Author: dzane17

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

@@ -27,6 +27,7 @@
 @PublicApi(since = "2.11.0")
 public final class SearchRequestStats extends SearchRequestOperationsListener {
     Map<SearchPhaseName, StatsHolder> phaseStatsMap = new EnumMap<>(SearchPhaseName.class);
+    StatsHolder tookStatsHolder;
 
     public static final String SEARCH_REQUEST_STATS_ENABLED_KEY = "search.request_stats_enabled";
     public static final Setting<Boolean> SEARCH_REQUEST_STATS_ENABLED = Setting.boolSetting(
@@ -40,6 +41,7 @@ public final class SearchRequestStats extends SearchRequestOperationsListener {
     public SearchRequestStats(ClusterSettings clusterSettings) {
         this.setEnabled(clusterSettings.get(SEARCH_REQUEST_STATS_ENABLED));
         clusterSettings.addSettingsUpdateConsumer(SEARCH_REQUEST_STATS_ENABLED, this::setEnabled);
+        tookStatsHolder = new StatsHolder();
         for (SearchPhaseName searchPhaseName : SearchPhaseName.values()) {
             phaseStatsMap.put(searchPhaseName, new StatsHolder());
         }
@@ -57,6 +59,18 @@ public long getPhaseMetric(SearchPhaseName searchPhaseName) {
         return phaseStatsMap.get(searchPhaseName).timing.sum();
     }
 
+    public long getTookCurrent() {
+        return tookStatsHolder.current.count();
+    }
+
+    public long getTookTotal() {
+        return tookStatsHolder.total.count();
+    }
+
+    public long getTookMetric() {
+        return tookStatsHolder.timing.sum();
+    }
+
     @Override
     protected void onPhaseStart(SearchPhaseContext context) {
         phaseStatsMap.get(context.getCurrentPhase().getSearchPhaseName()).current.inc();
@@ -75,6 +89,23 @@ protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {
         phaseStatsMap.get(context.getCurrentPhase().getSearchPhaseName()).current.dec();
     }
 
+    @Override
+    protected void onRequestStart(SearchRequestContext searchRequestContext) {
+        tookStatsHolder.current.inc();
+    }
+
+    @Override
+    protected void onRequestEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        tookStatsHolder.current.dec();
+        tookStatsHolder.total.inc();
+        tookStatsHolder.timing.inc(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - searchRequestContext.getAbsoluteStartNanos()));
+    }
+
+    @Override
+    protected void onRequestFailure(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        tookStatsHolder.current.dec();
+    }
+
     /**
      * Holder of statistics values
      *

server/src/main/java/org/opensearch/index/search/stats/SearchStats.java

@@ -110,7 +110,7 @@ public void writeTo(StreamOutput out) throws IOException {
     }
 
     /**
-     * Holds requests stats for different phases.
+     * Holds all requests stats.
      *
      * @opensearch.api
      */
@@ -124,6 +124,7 @@ public Map<String, PhaseStatsLongHolder> getRequestStatsHolder() {
         }
 
         RequestStatsLongHolder() {
+            requestStatsHolder.put(Fields.TOOK, new PhaseStatsLongHolder());
             for (SearchPhaseName searchPhaseName : SearchPhaseName.values()) {
                 requestStatsHolder.put(searchPhaseName.getName(), new PhaseStatsLongHolder());
             }
@@ -512,6 +513,15 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             if (requestStatsLongHolder != null) {
                 builder.startObject(Fields.REQUEST);
 
+                PhaseStatsLongHolder tookStatsLongHolder = requestStatsLongHolder.requestStatsHolder.get(Fields.TOOK);
+                if (tookStatsLongHolder != null) {
+                    builder.startObject(Fields.TOOK);
+                    builder.humanReadableField(Fields.TIME_IN_MILLIS, Fields.TIME, new TimeValue(tookStatsLongHolder.timeInMillis));
+                    builder.field(Fields.CURRENT, tookStatsLongHolder.current);
+                    builder.field(Fields.TOTAL, tookStatsLongHolder.total);
+                    builder.endObject();
+                }
+
                 for (SearchPhaseName searchPhaseName : SearchPhaseName.values()) {
                     PhaseStatsLongHolder statsLongHolder = requestStatsLongHolder.requestStatsHolder.get(searchPhaseName.getName());
                     if (statsLongHolder == null) {
@@ -545,6 +555,17 @@ public void setSearchRequestStats(SearchRequestStats searchRequestStats) {
             totalStats.requestStatsLongHolder = new RequestStatsLongHolder();
         }
 
+        // Set took stats
+        totalStats.requestStatsLongHolder.requestStatsHolder.put(
+            Fields.TOOK,
+            new PhaseStatsLongHolder(
+                searchRequestStats.getTookCurrent(),
+                searchRequestStats.getTookTotal(),
+                searchRequestStats.getTookMetric()
+            )
+        );
+
+        // Set phase stats
         for (SearchPhaseName searchPhaseName : SearchPhaseName.values()) {
             totalStats.requestStatsLongHolder.requestStatsHolder.put(
                 searchPhaseName.getName(),
@@ -678,6 +699,7 @@ static final class Fields {
         static final String CURRENT = "current";
         static final String TOTAL = "total";
         static final String SEARCH_IDLE_REACTIVATE_COUNT_TOTAL = "search_idle_reactivate_count_total";
+        static final String TOOK = "took";
 
     }
 

server/src/test/java/org/opensearch/action/search/SearchRequestStatsTests.java

@@ -25,6 +25,41 @@
 import static org.mockito.Mockito.when;
 
 public class SearchRequestStatsTests extends OpenSearchTestCase {
+    public void testSearchRequestStats_OnRequestFailure() {
+        ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
+        SearchRequestStats testRequestStats = new SearchRequestStats(clusterSettings);
+        SearchPhaseContext mockSearchPhaseContext = mock(SearchPhaseContext.class);
+        SearchRequestContext mockSearchRequestContext = mock(SearchRequestContext.class);
+
+        testRequestStats.onRequestStart(mockSearchRequestContext);
+        assertEquals(1, testRequestStats.getTookCurrent());
+        testRequestStats.onRequestFailure(mockSearchPhaseContext, mockSearchRequestContext);
+        assertEquals(0, testRequestStats.getTookCurrent());
+        assertEquals(0, testRequestStats.getTookTotal());
+    }
+
+    public void testSearchRequestStats_OnRequestEnd() {
+        ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
+        SearchRequestStats testRequestStats = new SearchRequestStats(clusterSettings);
+        SearchPhaseContext mockSearchPhaseContext = mock(SearchPhaseContext.class);
+        SearchRequestContext mockSearchRequestContext = mock(SearchRequestContext.class);
+
+        // Start request
+        testRequestStats.onRequestStart(mockSearchRequestContext);
+        assertEquals(1, testRequestStats.getTookCurrent());
+
+        // Mock start time
+        long tookTimeInMillis = randomIntBetween(1, 10);
+        long startTimeInNanos = System.nanoTime() - TimeUnit.MILLISECONDS.toNanos(tookTimeInMillis);
+        when(mockSearchRequestContext.getAbsoluteStartNanos()).thenReturn(startTimeInNanos);
+
+        // End request
+        testRequestStats.onRequestEnd(mockSearchPhaseContext, mockSearchRequestContext);
+        assertEquals(0, testRequestStats.getTookCurrent());
+        assertEquals(1, testRequestStats.getTookTotal());
+        assertThat(testRequestStats.getTookMetric(), greaterThanOrEqualTo(tookTimeInMillis));
+    }
+
     public void testSearchRequestPhaseFailure() {
         ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
         SearchRequestStats testRequestStats = new SearchRequestStats(clusterSettings);