Add completion count in SearchBackpressure nodeStats API

CHANGELOG.md

@@ -48,6 +48,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Migrate client transports to Apache HttpClient / Core 5.x ([#4459](https://github.com/opensearch-project/OpenSearch/pull/4459))
 - Change http code on create index API with bad input raising NotXContentException from 500 to 400 ([#4773](https://github.com/opensearch-project/OpenSearch/pull/4773))
 - Improve summary error message for invalid setting updates ([#4792](https://github.com/opensearch-project/OpenSearch/pull/4792))
+- Add completionCount in search_backpressure stats API
 - Return 409 Conflict HTTP status instead of 503 on failure to concurrently execute snapshots ([#8986](https://github.com/opensearch-project/OpenSearch/pull/5855))
 
 ### Deprecated

server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java

@@ -399,6 +399,7 @@ public SearchBackpressureStats nodeStats() {
         SearchTaskStats searchTaskStats = new SearchTaskStats(
             searchBackpressureStates.get(SearchTask.class).getCancellationCount(),
             searchBackpressureStates.get(SearchTask.class).getLimitReachedCount(),
+            searchBackpressureStates.get(SearchTask.class).getCompletionCount(),
             taskTrackers.get(SearchTask.class)
                 .stream()
                 .collect(Collectors.toUnmodifiableMap(t -> TaskResourceUsageTrackerType.fromName(t.name()), t -> t.stats(searchTasks)))
@@ -407,6 +408,7 @@ public SearchBackpressureStats nodeStats() {
         SearchShardTaskStats searchShardTaskStats = new SearchShardTaskStats(
             searchBackpressureStates.get(SearchShardTask.class).getCancellationCount(),
             searchBackpressureStates.get(SearchShardTask.class).getLimitReachedCount(),
+            searchBackpressureStates.get(SearchShardTask.class).getCompletionCount(),
             taskTrackers.get(SearchShardTask.class)
                 .stream()
                 .collect(Collectors.toUnmodifiableMap(t -> TaskResourceUsageTrackerType.fromName(t.name()), t -> t.stats(searchShardTasks)))

server/src/main/java/org/opensearch/search/backpressure/stats/SearchShardTaskStats.java

@@ -8,6 +8,7 @@
 
 package org.opensearch.search.backpressure.stats;
 
+import org.opensearch.Version;
 import org.opensearch.common.collect.MapBuilder;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -30,21 +31,27 @@
 public class SearchShardTaskStats implements ToXContentObject, Writeable {
     private final long cancellationCount;
     private final long limitReachedCount;
+    private Long completionCount;
     private final Map<TaskResourceUsageTrackerType, TaskResourceUsageTracker.Stats> resourceUsageTrackerStats;
 
     public SearchShardTaskStats(
         long cancellationCount,
         long limitReachedCount,
+        long completionCount,
         Map<TaskResourceUsageTrackerType, TaskResourceUsageTracker.Stats> resourceUsageTrackerStats
     ) {
         this.cancellationCount = cancellationCount;
         this.limitReachedCount = limitReachedCount;
+        this.completionCount = completionCount;
         this.resourceUsageTrackerStats = resourceUsageTrackerStats;
     }
 
     public SearchShardTaskStats(StreamInput in) throws IOException {
         this.cancellationCount = in.readVLong();
         this.limitReachedCount = in.readVLong();
+        if (in.getVersion().onOrAfter(Version.V_3_0_0)) {
+            this.completionCount = in.readOptionalVLong();
+        }
 
         MapBuilder<TaskResourceUsageTrackerType, TaskResourceUsageTracker.Stats> builder = new MapBuilder<>();
         builder.put(TaskResourceUsageTrackerType.CPU_USAGE_TRACKER, in.readOptionalWriteable(CpuUsageTracker.Stats::new));
@@ -62,6 +69,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(entry.getKey().getName(), entry.getValue());
         }
         builder.endObject();
+        if (completionCount != null) {
+            builder.field("completion_count", completionCount);
+        }
 
         builder.startObject("cancellation_stats")
             .field("cancellation_count", cancellationCount)
@@ -75,6 +85,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     public void writeTo(StreamOutput out) throws IOException {
         out.writeVLong(cancellationCount);
         out.writeVLong(limitReachedCount);
+        if (out.getVersion().onOrAfter(Version.V_3_0_0)) {
+            out.writeOptionalVLong(completionCount);
+        }
 
         out.writeOptionalWriteable(resourceUsageTrackerStats.get(TaskResourceUsageTrackerType.CPU_USAGE_TRACKER));
         out.writeOptionalWriteable(resourceUsageTrackerStats.get(TaskResourceUsageTrackerType.HEAP_USAGE_TRACKER));
@@ -88,11 +101,12 @@ public boolean equals(Object o) {
         SearchShardTaskStats that = (SearchShardTaskStats) o;
         return cancellationCount == that.cancellationCount
             && limitReachedCount == that.limitReachedCount
+            && Objects.equals(completionCount, that.completionCount)
             && resourceUsageTrackerStats.equals(that.resourceUsageTrackerStats);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(cancellationCount, limitReachedCount, resourceUsageTrackerStats);
+        return Objects.hash(cancellationCount, limitReachedCount, completionCount, resourceUsageTrackerStats);
     }
 }

server/src/main/java/org/opensearch/search/backpressure/stats/SearchTaskStats.java

@@ -8,6 +8,7 @@
 
 package org.opensearch.search.backpressure.stats;
 
+import org.opensearch.Version;
 import org.opensearch.common.collect.MapBuilder;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -31,21 +32,29 @@
 public class SearchTaskStats implements ToXContentObject, Writeable {
     private final long cancellationCount;
     private final long limitReachedCount;
+    private final long completionCount;
     private final Map<TaskResourceUsageTrackerType, TaskResourceUsageTracker.Stats> resourceUsageTrackerStats;
 
     public SearchTaskStats(
         long cancellationCount,
         long limitReachedCount,
+        long completionCount,
         Map<TaskResourceUsageTrackerType, TaskResourceUsageTracker.Stats> resourceUsageTrackerStats
     ) {
         this.cancellationCount = cancellationCount;
         this.limitReachedCount = limitReachedCount;
+        this.completionCount = completionCount;
         this.resourceUsageTrackerStats = resourceUsageTrackerStats;
     }
 
     public SearchTaskStats(StreamInput in) throws IOException {
         this.cancellationCount = in.readVLong();
         this.limitReachedCount = in.readVLong();
+        if (in.getVersion().onOrAfter(Version.V_3_0_0)) {
+            this.completionCount = in.readVLong();
+        } else {
+            this.completionCount = -1;
+        }
 
         MapBuilder<TaskResourceUsageTrackerType, TaskResourceUsageTracker.Stats> builder = new MapBuilder<>();
         builder.put(TaskResourceUsageTrackerType.CPU_USAGE_TRACKER, in.readOptionalWriteable(CpuUsageTracker.Stats::new));
@@ -63,6 +72,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(entry.getKey().getName(), entry.getValue());
         }
         builder.endObject();
+        if (completionCount != -1) {
+            builder.field("completion_count", completionCount);
+        }
+
 
         builder.startObject("cancellation_stats")
             .field("cancellation_count", cancellationCount)
@@ -76,6 +89,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     public void writeTo(StreamOutput out) throws IOException {
         out.writeVLong(cancellationCount);
         out.writeVLong(limitReachedCount);
+        if (out.getVersion().onOrAfter(Version.V_3_0_0)) {
+            out.writeVLong(completionCount);
+        }
 
         out.writeOptionalWriteable(resourceUsageTrackerStats.get(TaskResourceUsageTrackerType.CPU_USAGE_TRACKER));
         out.writeOptionalWriteable(resourceUsageTrackerStats.get(TaskResourceUsageTrackerType.HEAP_USAGE_TRACKER));
@@ -87,13 +103,18 @@ public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         SearchTaskStats that = (SearchTaskStats) o;
+        boolean isCompletionCountEqual = false;
+        if (that.completionCount != -1) {
+            isCompletionCountEqual = completionCount == that.completionCount;
+        }
         return cancellationCount == that.cancellationCount
             && limitReachedCount == that.limitReachedCount
+            && isCompletionCountEqual
             && resourceUsageTrackerStats.equals(that.resourceUsageTrackerStats);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(cancellationCount, limitReachedCount, resourceUsageTrackerStats);
+        return Objects.hash(cancellationCount, limitReachedCount, completionCount, resourceUsageTrackerStats);
     }
 }

server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java

@@ -249,10 +249,11 @@ public void testSearchTaskInFlightCancellation() {
         verify(mockTaskManager, times(10)).cancelTaskAndDescendants(any(), anyString(), anyBoolean(), any());
         assertEquals(3, service.getSearchBackpressureState(SearchTask.class).getLimitReachedCount());
 
-        // Verify search backpressure stats.
+        // Verify search backpressure stats. Since we are not marking any task as completed the completionCount will be 0
+        // for SearchTaskStats here.
         SearchBackpressureStats expectedStats = new SearchBackpressureStats(
-            new SearchTaskStats(10, 3, Map.of(TaskResourceUsageTrackerType.CPU_USAGE_TRACKER, new MockStats(10))),
-            new SearchShardTaskStats(0, 0, Collections.emptyMap()),
+            new SearchTaskStats(10, 3, 0, Map.of(TaskResourceUsageTrackerType.CPU_USAGE_TRACKER, new MockStats(10))),
+            new SearchShardTaskStats(0, 0, 0, Collections.emptyMap()),
             SearchBackpressureMode.ENFORCED
         );
         SearchBackpressureStats actualStats = service.nodeStats();
@@ -323,10 +324,11 @@ public void testSearchShardTaskInFlightCancellation() {
         verify(mockTaskManager, times(12)).cancelTaskAndDescendants(any(), anyString(), anyBoolean(), any());
         assertEquals(3, service.getSearchBackpressureState(SearchShardTask.class).getLimitReachedCount());
 
-        // Verify search backpressure stats.
+        // Verify search backpressure stats. We are marking 20 SearchShardTasks as completed this should get
+        // reflected in SearchShardTaskStats.
         SearchBackpressureStats expectedStats = new SearchBackpressureStats(
-            new SearchTaskStats(0, 0, Collections.emptyMap()),
-            new SearchShardTaskStats(12, 3, Map.of(TaskResourceUsageTrackerType.CPU_USAGE_TRACKER, new MockStats(12))),
+            new SearchTaskStats(0, 0, 0, Collections.emptyMap()),
+            new SearchShardTaskStats(12, 3, 20, Map.of(TaskResourceUsageTrackerType.CPU_USAGE_TRACKER, new MockStats(12))),
             SearchBackpressureMode.ENFORCED
         );
         SearchBackpressureStats actualStats = service.nodeStats();

server/src/test/java/org/opensearch/search/backpressure/stats/SearchShardTaskStatsTests.java

@@ -39,6 +39,6 @@ public static SearchShardTaskStats randomInstance() {
             new ElapsedTimeTracker.Stats(randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong())
         );
 
-        return new SearchShardTaskStats(randomNonNegativeLong(), randomNonNegativeLong(), resourceUsageTrackerStats);
+        return new SearchShardTaskStats(randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong(), resourceUsageTrackerStats);
     }
 }

server/src/test/java/org/opensearch/search/backpressure/stats/SearchTaskStatsTests.java

@@ -40,6 +40,6 @@ public static SearchTaskStats randomInstance() {
             new ElapsedTimeTracker.Stats(randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong())
         );
 
-        return new SearchTaskStats(randomNonNegativeLong(), randomNonNegativeLong(), resourceUsageTrackerStats);
+        return new SearchTaskStats(randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong(), resourceUsageTrackerStats);
     }
 }