Implemented computation of segment replication stats at shard level

CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Introduce a setting to disable download of full cluster state from remote on term mismatch([#16798](https://github.com/opensearch-project/OpenSearch/pull/16798/))
 - Added ability to retrieve value from DocValues in a flat_object filed([#16802](https://github.com/opensearch-project/OpenSearch/pull/16802))
 - Improve performace of NumericTermAggregation by avoiding unnecessary sorting([#17252](https://github.com/opensearch-project/OpenSearch/pull/17252))
+- Implemented computation of segment replication stats at shard level ([#17055](https://github.com/opensearch-project/OpenSearch/pull/17055))
 - [Rule Based Auto-tagging] Add in-memory attribute value store ([#17342](https://github.com/opensearch-project/OpenSearch/pull/17342))
 
 ### Dependencies

server/src/internalClusterTest/java/org/opensearch/index/shard/IndexShardIT.java

@@ -114,6 +114,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
@@ -136,6 +137,7 @@
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static com.carrotsearch.randomizedtesting.RandomizedTest.randomAsciiLettersOfLength;
+import static org.mockito.Mockito.mock;
 
 public class IndexShardIT extends OpenSearchSingleNodeTestCase {
 
@@ -716,7 +718,8 @@ public static final IndexShard newIndexShard(
             null,
             DefaultRemoteStoreSettings.INSTANCE,
             false,
-            IndexShardTestUtils.getFakeDiscoveryNodes(initializingShardRouting)
+            IndexShardTestUtils.getFakeDiscoveryNodes(initializingShardRouting),
+            mock(Function.class)
         );
     }
 

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationStatsIT.java

@@ -404,19 +404,17 @@ public void testSegmentReplicationNodeAndIndexStats() throws Exception {
 
             for (NodeStats nodeStats : nodesStatsResponse.getNodes()) {
                 ReplicationStats replicationStats = nodeStats.getIndices().getSegments().getReplicationStats();
-                // primary node - should hold replication statistics
+                // primary node - do not have any replication statistics
                 if (nodeStats.getNode().getName().equals(primaryNode)) {
+                    assertTrue(replicationStats.getMaxBytesBehind() == 0);
+                    assertTrue(replicationStats.getTotalBytesBehind() == 0);
+                    assertTrue(replicationStats.getMaxReplicationLag() == 0);
+                }
+                // replica nodes - should hold replication statistics
+                if (nodeStats.getNode().getName().equals(replicaNode1) || nodeStats.getNode().getName().equals(replicaNode2)) {
                     assertTrue(replicationStats.getMaxBytesBehind() > 0);
                     assertTrue(replicationStats.getTotalBytesBehind() > 0);
                     assertTrue(replicationStats.getMaxReplicationLag() > 0);
-                    // 2 replicas so total bytes should be double of max
-                    assertEquals(replicationStats.getMaxBytesBehind() * 2, replicationStats.getTotalBytesBehind());
-                }
-                // replica nodes - should hold empty replication statistics
-                if (nodeStats.getNode().getName().equals(replicaNode1) || nodeStats.getNode().getName().equals(replicaNode2)) {
-                    assertEquals(0, replicationStats.getMaxBytesBehind());
-                    assertEquals(0, replicationStats.getTotalBytesBehind());
-                    assertEquals(0, replicationStats.getMaxReplicationLag());
                 }
             }
             // get replication statistics at index level

server/src/main/java/org/opensearch/common/io/IndexIOStreamHandlerFactory.java

@@ -0,0 +1,25 @@
+/*
+ * 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.common.io;
+
+/**
+ * Interface for factory to provide handler implementation for type {@link T}
+ * @param <T> The type of content to be read/written to stream
+ *
+ * @opensearch.internal
+ */
+public interface IndexIOStreamHandlerFactory<T> {
+
+    /**
+     * Implements logic to provide handler based on the stream versions
+     * @param version stream version
+     * @return Handler for reading/writing content streams to/from - {@link T}
+     */
+    IndexIOStreamHandler<T> getHandler(int version);
+}

server/src/main/java/org/opensearch/common/io/VersionedCodecStreamWrapper.java

@@ -28,18 +28,25 @@
 public class VersionedCodecStreamWrapper<T> {
     private static final Logger logger = LogManager.getLogger(VersionedCodecStreamWrapper.class);
 
-    // TODO This can be updated to hold a streamReadWriteHandlerFactory and get relevant handler based on the stream versions
-    private final IndexIOStreamHandler<T> indexIOStreamHandler;
+    private final IndexIOStreamHandlerFactory<T> indexIOStreamHandlerFactory;
+    private final int minVersion;
     private final int currentVersion;
     private final String codec;
 
     /**
-     * @param indexIOStreamHandler handler to read/write stream from T
+     * @param indexIOStreamHandlerFactory factory for providing handler to read/write stream from T
+     * @param minVersion earliest supported version of the stream
      * @param currentVersion latest supported version of the stream
      * @param codec: stream codec
      */
-    public VersionedCodecStreamWrapper(IndexIOStreamHandler<T> indexIOStreamHandler, int currentVersion, String codec) {
-        this.indexIOStreamHandler = indexIOStreamHandler;
+    public VersionedCodecStreamWrapper(
+        IndexIOStreamHandlerFactory<T> indexIOStreamHandlerFactory,
+        int minVersion,
+        int currentVersion,
+        String codec
+    ) {
+        this.indexIOStreamHandlerFactory = indexIOStreamHandlerFactory;
+        this.minVersion = minVersion;
         this.currentVersion = currentVersion;
         this.codec = codec;
     }
@@ -87,7 +94,7 @@ public void writeStream(IndexOutput indexOutput, T content) throws IOException {
      */
     private int checkHeader(IndexInput indexInput) throws IOException {
         // TODO Once versioning strategy is decided we'll add support for min/max supported versions
-        return CodecUtil.checkHeader(indexInput, this.codec, this.currentVersion, this.currentVersion);
+        return CodecUtil.checkHeader(indexInput, this.codec, minVersion, this.currentVersion);
     }
 
     /**
@@ -120,8 +127,6 @@ private void writeFooter(IndexOutput indexOutput) throws IOException {
      * @param version stream content version
      */
     private IndexIOStreamHandler<T> getHandlerForVersion(int version) {
-        // TODO implement factory and pick relevant handler based on version.
-        // It should also take into account min and max supported versions
-        return this.indexIOStreamHandler;
+        return this.indexIOStreamHandlerFactory.getHandler(version);
     }
 }

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

@@ -57,6 +57,7 @@
 import org.opensearch.common.util.io.IOUtils;
 import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.core.index.Index;
+import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.indices.breaker.CircuitBreakerService;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.env.NodeEnvironment;
@@ -652,7 +653,8 @@ public IndexService newIndexService(
             clusterDefaultRefreshIntervalSupplier,
             recoverySettings,
             remoteStoreSettings,
-            (s) -> {}
+            (s) -> {},
+            shardId -> ReplicationStats.empty()
         );
     }
 
@@ -678,7 +680,8 @@ public IndexService newIndexService(
         Supplier<TimeValue> clusterDefaultRefreshIntervalSupplier,
         RecoverySettings recoverySettings,
         RemoteStoreSettings remoteStoreSettings,
-        Consumer<IndexShard> replicator
+        Consumer<IndexShard> replicator,
+        Function<ShardId, ReplicationStats> segmentReplicationStatsProvider
     ) throws IOException {
         final IndexEventListener eventListener = freeze();
         Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrapperFactory = indexReaderWrapper
@@ -740,7 +743,8 @@ public IndexService newIndexService(
                 remoteStoreSettings,
                 fileCache,
                 compositeIndexSettings,
-                replicator
+                replicator,
+                segmentReplicationStatsProvider
             );
             success = true;
             return indexService;

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

@@ -197,6 +197,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust
     private final FileCache fileCache;
     private final CompositeIndexSettings compositeIndexSettings;
     private final Consumer<IndexShard> replicator;
+    private final Function<ShardId, ReplicationStats> segmentReplicationStatsProvider;
 
     public IndexService(
         IndexSettings indexSettings,
@@ -235,7 +236,8 @@ public IndexService(
         RemoteStoreSettings remoteStoreSettings,
         FileCache fileCache,
         CompositeIndexSettings compositeIndexSettings,
-        Consumer<IndexShard> replicator
+        Consumer<IndexShard> replicator,
+        Function<ShardId, ReplicationStats> segmentReplicationStatsProvider
     ) {
         super(indexSettings);
         this.allowExpensiveQueries = allowExpensiveQueries;
@@ -322,6 +324,7 @@ public IndexService(
         this.compositeIndexSettings = compositeIndexSettings;
         this.fileCache = fileCache;
         this.replicator = replicator;
+        this.segmentReplicationStatsProvider = segmentReplicationStatsProvider;
         updateFsyncTaskIfNecessary();
     }
 
@@ -398,7 +401,8 @@ public IndexService(
             remoteStoreSettings,
             null,
             null,
-            s -> {}
+            s -> {},
+            (shardId) -> ReplicationStats.empty()
         );
     }
 
@@ -694,7 +698,8 @@ protected void closeInternal() {
                 recoverySettings,
                 remoteStoreSettings,
                 seedRemote,
-                discoveryNodes
+                discoveryNodes,
+                segmentReplicationStatsProvider
             );
             eventListener.indexShardStateChanged(indexShard, null, indexShard.state(), "shard created");
             eventListener.afterIndexShardCreated(indexShard);

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

@@ -42,6 +42,10 @@ public ReplicationStats(StreamInput in) throws IOException {
         this.maxReplicationLag = in.readVLong();
     }
 
+    public static ReplicationStats empty() {
+        return new ReplicationStats();
+    }
+
     public ReplicationStats() {
 
     }

server/src/main/java/org/opensearch/index/shard/IndexShard.java

@@ -361,6 +361,7 @@ Runnable getGlobalCheckpointSyncer() {
      */
     private final ShardMigrationState shardMigrationState;
     private DiscoveryNodes discoveryNodes;
+    private final Function<ShardId, ReplicationStats> segmentReplicationStatsProvider;
 
     public IndexShard(
         final ShardRouting shardRouting,
@@ -391,7 +392,8 @@ public IndexShard(
         final RecoverySettings recoverySettings,
         final RemoteStoreSettings remoteStoreSettings,
         boolean seedRemote,
-        final DiscoveryNodes discoveryNodes
+        final DiscoveryNodes discoveryNodes,
+        final Function<ShardId, ReplicationStats> segmentReplicationStatsProvider
     ) throws IOException {
         super(shardRouting.shardId(), indexSettings);
         assert shardRouting.initializing();
@@ -493,6 +495,7 @@ public boolean shouldCache(Query query) {
         this.fileDownloader = new RemoteStoreFileDownloader(shardRouting.shardId(), threadPool, recoverySettings);
         this.shardMigrationState = getShardMigrationState(indexSettings, seedRemote);
         this.discoveryNodes = discoveryNodes;
+        this.segmentReplicationStatsProvider = segmentReplicationStatsProvider;
     }
 
     public ThreadPool getThreadPool() {
@@ -3233,17 +3236,10 @@ public Set<SegmentReplicationShardStats> getReplicationStatsForTrackedReplicas()
     }
 
     public ReplicationStats getReplicationStats() {
-        if (indexSettings.isSegRepEnabledOrRemoteNode() && routingEntry().primary()) {
-            final Set<SegmentReplicationShardStats> stats = getReplicationStatsForTrackedReplicas();
-            long maxBytesBehind = stats.stream().mapToLong(SegmentReplicationShardStats::getBytesBehindCount).max().orElse(0L);
-            long totalBytesBehind = stats.stream().mapToLong(SegmentReplicationShardStats::getBytesBehindCount).sum();
-            long maxReplicationLag = stats.stream()
-                .mapToLong(SegmentReplicationShardStats::getCurrentReplicationLagMillis)
-                .max()
-                .orElse(0L);
-            return new ReplicationStats(maxBytesBehind, totalBytesBehind, maxReplicationLag);
-        }
-        return new ReplicationStats();
+        if (indexSettings.isSegRepEnabledOrRemoteNode() && !routingEntry().primary()) {
+            return segmentReplicationStatsProvider.apply(shardId);
+        }
+        return ReplicationStats.empty();
     }
 
     /**

server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java

@@ -38,7 +38,7 @@
 import org.opensearch.index.store.lockmanager.RemoteStoreLockManager;
 import org.opensearch.index.store.lockmanager.RemoteStoreMetadataLockManager;
 import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadata;
-import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadataHandler;
+import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadataHandlerFactory;
 import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 import org.opensearch.node.remotestore.RemoteStorePinnedTimestampService;
 import org.opensearch.threadpool.ThreadPool;
@@ -104,7 +104,8 @@ public final class RemoteSegmentStoreDirectory extends FilterDirectory implement
     private Map<String, UploadedSegmentMetadata> segmentsUploadedToRemoteStore;
 
     private static final VersionedCodecStreamWrapper<RemoteSegmentMetadata> metadataStreamWrapper = new VersionedCodecStreamWrapper<>(
-        new RemoteSegmentMetadataHandler(),
+        new RemoteSegmentMetadataHandlerFactory(),
+        RemoteSegmentMetadata.VERSION_ONE,
         RemoteSegmentMetadata.CURRENT_VERSION,
         RemoteSegmentMetadata.METADATA_CODEC
     );

server/src/main/java/org/opensearch/index/store/remote/metadata/RemoteSegmentMetadata.java

@@ -30,10 +30,15 @@
  */
 @PublicApi(since = "2.6.0")
 public class RemoteSegmentMetadata {
+
+    public static final int VERSION_ONE = 1;
+
+    public static final int VERSION_TWO = 2;
+
     /**
      * Latest supported version of metadata
      */
-    public static final int CURRENT_VERSION = 1;
+    public static final int CURRENT_VERSION = VERSION_TWO;
     /**
      * Metadata codec
      */
@@ -106,18 +111,30 @@ public static Map<String, RemoteSegmentStoreDirectory.UploadedSegmentMetadata> f
             );
     }
 
+    /**
+     * Write always writes with the latest version of the RemoteSegmentMetadata
+     * @param out file output stream which will store stream content
+     * @throws IOException in case there is a problem writing the file
+     */
     public void write(IndexOutput out) throws IOException {
         out.writeMapOfStrings(toMapOfStrings());
         writeCheckpointToIndexOutput(replicationCheckpoint, out);
         out.writeLong(segmentInfosBytes.length);
         out.writeBytes(segmentInfosBytes, segmentInfosBytes.length);
     }
 
-    public static RemoteSegmentMetadata read(IndexInput indexInput) throws IOException {
+    /**
+     * Read can happen in the upgraded version of replica which needs to support all versions of RemoteSegmentMetadata
+     * @param indexInput file input stream
+     * @param version version of the RemoteSegmentMetadata
+     * @return {@code RemoteSegmentMetadata}
+     * @throws IOException in case there is a problem reading from the file input stream
+     */
+    public static RemoteSegmentMetadata read(IndexInput indexInput, int version) throws IOException {
         Map<String, String> metadata = indexInput.readMapOfStrings();
         final Map<String, RemoteSegmentStoreDirectory.UploadedSegmentMetadata> uploadedSegmentMetadataMap = RemoteSegmentMetadata
             .fromMapOfStrings(metadata);
-        ReplicationCheckpoint replicationCheckpoint = readCheckpointFromIndexInput(indexInput, uploadedSegmentMetadataMap);
+        ReplicationCheckpoint replicationCheckpoint = readCheckpointFromIndexInput(indexInput, uploadedSegmentMetadataMap, version);
         int byteArraySize = (int) indexInput.readLong();
         byte[] segmentInfosBytes = new byte[byteArraySize];
         indexInput.readBytes(segmentInfosBytes, 0, byteArraySize);
@@ -136,11 +153,13 @@ public static void writeCheckpointToIndexOutput(ReplicationCheckpoint replicatio
         out.writeLong(replicationCheckpoint.getSegmentInfosVersion());
         out.writeLong(replicationCheckpoint.getLength());
         out.writeString(replicationCheckpoint.getCodec());
+        out.writeLong(replicationCheckpoint.getCreatedTimeStamp());
     }
 
     private static ReplicationCheckpoint readCheckpointFromIndexInput(
         IndexInput in,
-        Map<String, RemoteSegmentStoreDirectory.UploadedSegmentMetadata> uploadedSegmentMetadataMap
+        Map<String, RemoteSegmentStoreDirectory.UploadedSegmentMetadata> uploadedSegmentMetadataMap,
+        int version
     ) throws IOException {
         return new ReplicationCheckpoint(
             new ShardId(new Index(in.readString(), in.readString()), in.readVInt()),
@@ -149,7 +168,8 @@ private static ReplicationCheckpoint readCheckpointFromIndexInput(
             in.readLong(),
             in.readLong(),
             in.readString(),
-            toStoreFileMetadata(uploadedSegmentMetadataMap)
+            toStoreFileMetadata(uploadedSegmentMetadataMap),
+            version >= VERSION_TWO ? in.readLong() : 0
         );
     }
 

server/src/main/java/org/opensearch/index/store/remote/metadata/RemoteSegmentMetadataHandler.java

@@ -20,14 +20,21 @@
  * @opensearch.internal
  */
 public class RemoteSegmentMetadataHandler implements IndexIOStreamHandler<RemoteSegmentMetadata> {
+
+    private final int version;
+
+    public RemoteSegmentMetadataHandler(int version) {
+        this.version = version;
+    }
+
     /**
      * Reads metadata content from metadata file input stream and parsed into {@link RemoteSegmentMetadata}
      * @param indexInput metadata file input stream with {@link IndexInput#getFilePointer()} pointing to metadata content
      * @return {@link RemoteSegmentMetadata}
      */
     @Override
     public RemoteSegmentMetadata readContent(IndexInput indexInput) throws IOException {
-        return RemoteSegmentMetadata.read(indexInput);
+        return RemoteSegmentMetadata.read(indexInput, version);
     }
 
     /**