From bcceb86d8d4611531a4a5a4ffc902a413b1e3d88 Mon Sep 17 00:00:00 2001
From: Pranshu Shukla <pranshushukla06@gmail.com>
Date: Wed, 28 Aug 2024 10:26:29 +0530
Subject: [PATCH] Addressing comments

Signed-off-by: Pranshu Shukla <pranshushukla06@gmail.com>
---
 CHANGELOG.md                                  |  1 +
 .../opensearch/indices/NodeIndicesStats.java  | 25 +++++++++++--------
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3dff44ed96dfd..39706af949416 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Adding access to noSubMatches and noOverlappingMatches in Hyphenation ([#13895](https://github.com/opensearch-project/OpenSearch/pull/13895))
 - Add support for index level max slice count setting for concurrent segment search ([#15336](https://github.com/opensearch-project/OpenSearch/pull/15336))
 - Add concurrent search support for Derived Fields ([#15326](https://github.com/opensearch-project/OpenSearch/pull/15326))
+- Optimize NodeIndicesStats output behind flag ([#14454](https://github.com/opensearch-project/OpenSearch/pull/14454))
 
 ### Dependencies
 - Bump `netty` from 4.1.111.Final to 4.1.112.Final ([#15081](https://github.com/opensearch-project/OpenSearch/pull/15081))
diff --git a/server/src/main/java/org/opensearch/indices/NodeIndicesStats.java b/server/src/main/java/org/opensearch/indices/NodeIndicesStats.java
index c629eaf867fc8..83a759cdb71c5 100644
--- a/server/src/main/java/org/opensearch/indices/NodeIndicesStats.java
+++ b/server/src/main/java/org/opensearch/indices/NodeIndicesStats.java
@@ -86,13 +86,11 @@ public NodeIndicesStats(StreamInput in) throws IOException {
         if (in.getVersion().onOrAfter(Version.V_3_0_0)) {
             // contains statsByIndex
             if (in.readBoolean()) {
-                statsByIndex = new HashMap<>();
-                readStatsByIndex(in);
+                statsByIndex = readStatsByIndex(in);
             }
         }
         if (in.readBoolean()) {
-            statsByShard = new HashMap<>();
-            readStatsByShards(in);
+            statsByShard = readStatsByShard(in);
         }
     }
 
@@ -158,8 +156,12 @@ public NodeIndicesStats(
 
     /**
      * By default, the levels passed from the transport action will be a list of strings, since NodeIndicesStats can
-     * only aggregate on one level, we pick the first accepted level else we ignore if no known level is passed.
+     * only aggregate on one level, we pick the first accepted level else we ignore if no known level is passed. Level is
+     * selected based on enum defined in {@link StatsLevel}
+     *
+     * Note - we are picking the first level as multiple levels are not supported in the previous versions.
      * @param levels - levels sent in the request.
+     *
      * @return Corresponding identified enum {@link StatsLevel}
      */
     public static StatsLevel getAcceptedLevel(String[] levels) {
@@ -172,16 +174,19 @@ public static StatsLevel getAcceptedLevel(String[] levels) {
         return null;
     }
 
-    private void readStatsByIndex(StreamInput in) throws IOException {
+    private Map<Index, CommonStats> readStatsByIndex(StreamInput in) throws IOException {
+        Map<Index, CommonStats> statsByIndex = new HashMap<>();
         int indexEntries = in.readVInt();
         for (int i = 0; i < indexEntries; i++) {
             Index index = new Index(in);
             CommonStats commonStats = new CommonStats(in);
             statsByIndex.put(index, commonStats);
         }
+        return statsByIndex;
     }
 
-    private void readStatsByShards(StreamInput in) throws IOException {
+    private Map<Index, List<IndexShardStats>> readStatsByShard(StreamInput in) throws IOException {
+        Map<Index, List<IndexShardStats>> statsByShard = new HashMap<>();
         int entries = in.readVInt();
         for (int i = 0; i < entries; i++) {
             Index index = new Index(in);
@@ -192,6 +197,7 @@ private void readStatsByShards(StreamInput in) throws IOException {
             }
             statsByShard.put(index, indexShardStats);
         }
+        return statsByShard;
     }
 
     @Nullable
@@ -324,8 +330,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             throw new IllegalArgumentException(
                 "level parameter must be one of ["
                     + StatsLevel.INDICES.getRestName()
-                    + "] or "
-                    + "["
+                    + "] or ["
                     + StatsLevel.NODE.getRestName()
                     + "] or ["
                     + StatsLevel.SHARDS.getRestName()
@@ -405,7 +410,7 @@ public List<IndexShardStats> getShardStats(Index index) {
      *
      * @opensearch.internal
      */
-    @PublicApi(since = "2.0.0")
+    @PublicApi(since = "3.0.0")
     public enum StatsLevel {
         INDICES("indices"),
         SHARDS("shards"),