Address comments on PR.

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>

server/src/main/java/org/opensearch/index/SegmentReplicationStatsTracker.java

@@ -33,18 +33,12 @@ public SegmentReplicationStatsTracker(IndicesService indicesService) {
         rejectionCount = ConcurrentCollections.newConcurrentMap();
     }
 
-    public SegmentReplicationRejectionStats getRejectionStats() {
-        long rejectionCount = 0;
-        for (IndexService indexService : indicesService) {
-            if (indexService.getIndexSettings().isSegRepEnabled()) {
-                for (IndexShard indexShard : indexService) {
-                    if (indexShard.indexSettings().isSegRepEnabled() && indexShard.routingEntry().primary()) {
-                        rejectionCount += getStatsForShard(indexShard).getRejectedRequestCount();
-                    }
-                }
-            }
-        }
-        return new SegmentReplicationRejectionStats(rejectionCount);
+    public SegmentReplicationRejectionStats getTotalRejectionStats() {
+        return new SegmentReplicationRejectionStats(this.rejectionCount.values().stream().mapToInt(AtomicInteger::get).sum());
+    }
+
+    protected Map<ShardId, AtomicInteger> getRejectionCount() {
+        return rejectionCount;
     }
 
     public SegmentReplicationStats getStats() {

server/src/main/java/org/opensearch/node/NodeService.java

@@ -262,7 +262,7 @@ public NodeStats stats(
             fileCacheStats && fileCache != null ? fileCache.fileCacheStats() : null,
             taskCancellation ? this.taskCancellationMonitoringService.stats() : null,
             searchPipelineStats ? this.searchPipelineService.stats() : null,
-            segmentReplicationTrackerStats ? this.segmentReplicationStatsTracker.getRejectionStats() : null,
+            segmentReplicationTrackerStats ? this.segmentReplicationStatsTracker.getTotalRejectionStats() : null,
             repositoriesStats ? this.repositoriesService.getRepositoriesStats() : null
         );
     }

server/src/test/java/org/opensearch/index/SegmentReplicationStatsTrackerTests.java

@@ -0,0 +1,28 @@
+/*
+ * 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.index;
+
+import org.opensearch.indices.IndicesService;
+import org.opensearch.test.OpenSearchTestCase;
+
+import static org.mockito.Mockito.mock;
+
+public class SegmentReplicationStatsTrackerTests extends OpenSearchTestCase {
+
+    private IndicesService indicesService = mock(IndicesService.class);
+
+    public void testRejectedCountWhenEmpty() {
+        SegmentReplicationStatsTracker segmentReplicationStatsTracker = new SegmentReplicationStatsTracker(indicesService);
+
+        // Verify that total rejection count is 0 on an empty rejectionCount map in statsTracker.
+        assertTrue(segmentReplicationStatsTracker.getRejectionCount().isEmpty());
+        assertEquals(segmentReplicationStatsTracker.getTotalRejectionStats().getTotalRejectionCount(), 0L);
+    }
+
+}