From 9e183e929f663ccb00cec5b19583a961efce0713 Mon Sep 17 00:00:00 2001
From: Rohit Ashiwal <rashiwal@amazon.com>
Date: Tue, 4 Jul 2023 14:51:07 +0530
Subject: [PATCH] introduce rest accounting for indexing actions

---
 .idea/runConfigurations/Debug_OpenSearch.xml  |  6 +-
 .../action/bulk/TransportBulkAction.java      | 10 +++
 .../opensearch/index/shard/IndexingStats.java | 81 ++++++++++++++++++-
 .../opensearch/indices/IndicesService.java    |  9 +++
 .../action/bulk/TransportBulkActionTests.java |  1 +
 .../opensearch/test/InternalTestCluster.java  |  4 +-
 6 files changed, 107 insertions(+), 4 deletions(-)

diff --git a/.idea/runConfigurations/Debug_OpenSearch.xml b/.idea/runConfigurations/Debug_OpenSearch.xml
index 0d8bf59823acf..c18046f873477 100644
--- a/.idea/runConfigurations/Debug_OpenSearch.xml
+++ b/.idea/runConfigurations/Debug_OpenSearch.xml
@@ -6,6 +6,10 @@
     <option name="HOST" value="localhost" />
     <option name="PORT" value="5005" />
     <option name="AUTO_RESTART" value="true" />
+    <RunnerSettings RunnerId="Debug">
+      <option name="DEBUG_PORT" value="5005" />
+      <option name="LOCAL" value="false" />
+    </RunnerSettings>
     <method v="2" />
   </configuration>
-</component>
+</component>
\ No newline at end of file
diff --git a/server/src/main/java/org/opensearch/action/bulk/TransportBulkAction.java b/server/src/main/java/org/opensearch/action/bulk/TransportBulkAction.java
index 19f610908677f..c87cdb2e4296f 100644
--- a/server/src/main/java/org/opensearch/action/bulk/TransportBulkAction.java
+++ b/server/src/main/java/org/opensearch/action/bulk/TransportBulkAction.java
@@ -36,6 +36,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.util.SparseFixedBitSet;
 import org.opensearch.LegacyESVersion;
+import org.opensearch.action.admin.indices.stats.CommonStatsFlags;
 import org.opensearch.core.Assertions;
 import org.opensearch.OpenSearchParseException;
 import org.opensearch.ExceptionsHelper;
@@ -78,8 +79,10 @@
 import org.opensearch.index.IndexingPressureService;
 import org.opensearch.index.VersionType;
 import org.opensearch.index.seqno.SequenceNumbers;
+import org.opensearch.index.shard.IndexingStats;
 import org.opensearch.index.shard.ShardId;
 import org.opensearch.indices.IndexClosedException;
+import org.opensearch.indices.IndicesService;
 import org.opensearch.indices.SystemIndices;
 import org.opensearch.ingest.IngestService;
 import org.opensearch.node.NodeClosedException;
@@ -130,6 +133,7 @@ public class TransportBulkAction extends HandledTransportAction<BulkRequest, Bul
     private final IndexNameExpressionResolver indexNameExpressionResolver;
     private static final String DROPPED_ITEM_WITH_AUTO_GENERATED_ID = "auto-generated";
     private final IndexingPressureService indexingPressureService;
+    private final IndicesService indicesService;
     private final SystemIndices systemIndices;
 
     @Inject
@@ -144,6 +148,7 @@ public TransportBulkAction(
         IndexNameExpressionResolver indexNameExpressionResolver,
         AutoCreateIndex autoCreateIndex,
         IndexingPressureService indexingPressureService,
+        IndicesService indicesService,
         SystemIndices systemIndices
     ) {
         this(
@@ -157,6 +162,7 @@ public TransportBulkAction(
             indexNameExpressionResolver,
             autoCreateIndex,
             indexingPressureService,
+            indicesService,
             systemIndices,
             System::nanoTime
         );
@@ -173,6 +179,7 @@ public TransportBulkAction(
         IndexNameExpressionResolver indexNameExpressionResolver,
         AutoCreateIndex autoCreateIndex,
         IndexingPressureService indexingPressureService,
+        IndicesService indicesService,
         SystemIndices systemIndices,
         LongSupplier relativeTimeProvider
     ) {
@@ -188,6 +195,7 @@ public TransportBulkAction(
         this.client = client;
         this.indexNameExpressionResolver = indexNameExpressionResolver;
         this.indexingPressureService = indexingPressureService;
+        this.indicesService = indicesService;
         this.systemIndices = systemIndices;
         clusterService.addStateApplier(this.ingestForwarder);
     }
@@ -637,6 +645,8 @@ public void onResponse(BulkShardResponse bulkShardResponse) {
                             if (bulkItemResponse.getResponse() != null) {
                                 bulkItemResponse.getResponse().setShardInfo(bulkShardResponse.getShardInfo());
                             }
+
+                            indicesService.updateBulkItemRestStatus(bulkItemResponse.status());
                             responses.set(bulkItemResponse.getItemId(), bulkItemResponse);
                         }
                         if (counter.decrementAndGet() == 0) {
diff --git a/server/src/main/java/org/opensearch/index/shard/IndexingStats.java b/server/src/main/java/org/opensearch/index/shard/IndexingStats.java
index d40878674ac13..aac69d29f252e 100644
--- a/server/src/main/java/org/opensearch/index/shard/IndexingStats.java
+++ b/server/src/main/java/org/opensearch/index/shard/IndexingStats.java
@@ -44,6 +44,8 @@
 
 import java.io.IOException;
 import java.util.Map;
+import java.util.TreeMap;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Tracks indexing statistics
@@ -59,6 +61,69 @@ public class IndexingStats implements Writeable, ToXContentFragment {
      */
     public static class Stats implements Writeable, ToXContentFragment {
 
+        /**
+         * Tracks item level rest status codes during indexing
+         *
+         * @opensearch.internal
+         */
+        public static class BulkStats implements Writeable, ToXContentFragment {
+
+            private final Map<Integer, AtomicLong> stats;
+
+            public BulkStats() {
+                this.stats = new TreeMap<>();
+            }
+
+            public BulkStats(StreamInput in) throws IOException {
+                int size = in.readInt();
+                stats = new TreeMap<>();
+
+                for (int i = 0; i < size; ++i) {
+                    stats.put(in.readInt(), new AtomicLong(in.readLong()));
+                }
+            }
+
+            public void put(int status) {
+                synchronized (this) {
+                    stats.computeIfAbsent(status, s -> new AtomicLong(0));
+                    stats.get(status).incrementAndGet();
+                }
+            }
+
+            @Override
+            public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+                builder.startObject("bulk");
+
+                for (Map.Entry<Integer, AtomicLong> entry : stats.entrySet()) {
+                    builder.field(String.valueOf(entry.getKey()), entry.getValue().longValue());
+                }
+
+                return builder.endObject();
+            }
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                out.writeInt(stats.size());
+
+                for (Map.Entry<Integer, AtomicLong> entry : stats.entrySet()) {
+                    out.writeInt(entry.getKey());
+                    out.writeLong(entry.getValue().longValue());
+                }
+            }
+
+            public void add(BulkStats other) {
+                for (Map.Entry<Integer, AtomicLong> entry : other.stats.entrySet()) {
+                    synchronized (this) {
+                        int k = entry.getKey();
+                        AtomicLong v = entry.getValue();
+
+                        stats.putIfAbsent(k, new AtomicLong(0));
+                        stats.get(k).addAndGet(v.longValue());
+                    }
+                }
+            }
+        }
+
         private long indexCount;
         private long indexTimeInMillis;
         private long indexCurrent;
@@ -70,7 +135,11 @@ public static class Stats implements Writeable, ToXContentFragment {
         private long throttleTimeInMillis;
         private boolean isThrottled;
 
-        Stats() {}
+        private final BulkStats bulkStats;
+
+        Stats() {
+            bulkStats = new BulkStats();
+        }
 
         public Stats(StreamInput in) throws IOException {
             indexCount = in.readVLong();
@@ -83,6 +152,7 @@ public Stats(StreamInput in) throws IOException {
             noopUpdateCount = in.readVLong();
             isThrottled = in.readBoolean();
             throttleTimeInMillis = in.readLong();
+            bulkStats = new BulkStats(in);
         }
 
         public Stats(
@@ -107,6 +177,7 @@ public Stats(
             this.noopUpdateCount = noopUpdateCount;
             this.isThrottled = isThrottled;
             this.throttleTimeInMillis = throttleTimeInMillis;
+            this.bulkStats = new BulkStats();
         }
 
         public void add(Stats stats) {
@@ -124,6 +195,7 @@ public void add(Stats stats) {
             if (isThrottled != stats.isThrottled) {
                 isThrottled = true; // When combining if one is throttled set result to throttled.
             }
+            bulkStats.add(stats.getBulkStats());
         }
 
         /**
@@ -193,6 +265,10 @@ public long getNoopUpdateCount() {
             return noopUpdateCount;
         }
 
+        public BulkStats getBulkStats() {
+            return bulkStats;
+        }
+
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeVLong(indexCount);
@@ -205,7 +281,7 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeVLong(noopUpdateCount);
             out.writeBoolean(isThrottled);
             out.writeLong(throttleTimeInMillis);
-
+            bulkStats.writeTo(out);
         }
 
         @Override
@@ -223,6 +299,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
 
             builder.field(Fields.IS_THROTTLED, isThrottled);
             builder.humanReadableField(Fields.THROTTLED_TIME_IN_MILLIS, Fields.THROTTLED_TIME, getThrottleTime());
+            bulkStats.toXContent(builder, params);
             return builder;
         }
     }
diff --git a/server/src/main/java/org/opensearch/indices/IndicesService.java b/server/src/main/java/org/opensearch/indices/IndicesService.java
index 81b1be0f58804..6d77fca9d52e5 100644
--- a/server/src/main/java/org/opensearch/indices/IndicesService.java
+++ b/server/src/main/java/org/opensearch/indices/IndicesService.java
@@ -152,6 +152,7 @@
 import org.opensearch.plugins.IndexStorePlugin;
 import org.opensearch.plugins.PluginsService;
 import org.opensearch.repositories.RepositoriesService;
+import org.opensearch.rest.RestStatus;
 import org.opensearch.script.ScriptService;
 import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
 import org.opensearch.search.internal.AliasFilter;
@@ -1015,6 +1016,14 @@ public IndicesQueryCache getIndicesQueryCache() {
         return indicesQueryCache;
     }
 
+    public void updateBulkItemRestStatus(final RestStatus status) {
+        oldShardsStats
+            .indexingStats
+            .getTotal()
+            .getBulkStats()
+            .put(status.getStatus());
+    }
+
     /**
      * Statistics for old shards
      *
diff --git a/server/src/test/java/org/opensearch/action/bulk/TransportBulkActionTests.java b/server/src/test/java/org/opensearch/action/bulk/TransportBulkActionTests.java
index d53b860e6524a..580bc46765809 100644
--- a/server/src/test/java/org/opensearch/action/bulk/TransportBulkActionTests.java
+++ b/server/src/test/java/org/opensearch/action/bulk/TransportBulkActionTests.java
@@ -114,6 +114,7 @@ class TestTransportBulkAction extends TransportBulkAction {
                 new Resolver(),
                 new AutoCreateIndex(Settings.EMPTY, clusterService.getClusterSettings(), new Resolver(), new SystemIndices(emptyMap())),
                 new IndexingPressureService(Settings.EMPTY, clusterService),
+                null,
                 new SystemIndices(emptyMap())
             );
         }
diff --git a/test/framework/src/main/java/org/opensearch/test/InternalTestCluster.java b/test/framework/src/main/java/org/opensearch/test/InternalTestCluster.java
index 714b0ab0b9bf4..7dd32a2136ed9 100644
--- a/test/framework/src/main/java/org/opensearch/test/InternalTestCluster.java
+++ b/test/framework/src/main/java/org/opensearch/test/InternalTestCluster.java
@@ -2685,7 +2685,9 @@ public void ensureEstimatedStats() {
                     false,
                     false,
                     false,
-                    false
+                    false,
+                    false,
+                    null
                 );
                 assertThat(
                     "Fielddata size must be 0 on node: " + stats.getNode(),