Implemented computation of segment replication stats at shard level

CHANGELOG.md

@@ -34,6 +34,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Added a new `time` field to replace the deprecated `getTime` field in `GetStats`. ([#17009](https://github.com/opensearch-project/OpenSearch/pull/17009))
 - Improve flat_object field parsing performance by reducing two passes to a single pass ([#16297](https://github.com/opensearch-project/OpenSearch/pull/16297))
 - Improve performance of the bitmap filtering([#16936](https://github.com/opensearch-project/OpenSearch/pull/16936/))
+- Implemented computation of segment replication stats at shard level ([#17055](https://github.com/opensearch-project/OpenSearch/pull/17055))
 
 ### Dependencies
 - Bump `com.google.cloud:google-cloud-core-http` from 2.23.0 to 2.47.0 ([#16504](https://github.com/opensearch-project/OpenSearch/pull/16504))

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.BiFunction;
 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(BiFunction.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);
-                    // 2 replicas so total bytes should be double of max
-                    assertEquals(replicationStats.getMaxBytesBehind() * 2, replicationStats.getTotalBytesBehind());
+                    assertTrue(replicationStats.getMaxBytesBehind() == 0);
+                    assertTrue(replicationStats.getTotalBytesBehind() == 0);
+                    assertTrue(replicationStats.getMaxReplicationLag() == 0);
                 }
                 // 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());
+                    assertTrue(replicationStats.getMaxBytesBehind() > 0);
+                    assertTrue(replicationStats.getTotalBytesBehind() > 0);
+                    assertTrue(replicationStats.getMaxReplicationLag() > 0);
                 }
             }
             // get replication statistics at index level

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

@@ -59,6 +59,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;
@@ -87,6 +88,7 @@
 import org.opensearch.indices.mapper.MapperRegistry;
 import org.opensearch.indices.recovery.RecoverySettings;
 import org.opensearch.indices.recovery.RecoveryState;
+import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 import org.opensearch.plugins.IndexStorePlugin;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.script.ScriptService;
@@ -653,7 +655,8 @@ public IndexService newIndexService(
             clusterDefaultRefreshIntervalSupplier,
             recoverySettings,
             remoteStoreSettings,
-            (s) -> {}
+            (s) -> {},
+            null
         );
     }
 
@@ -679,7 +682,8 @@ public IndexService newIndexService(
         Supplier<TimeValue> clusterDefaultRefreshIntervalSupplier,
         RecoverySettings recoverySettings,
         RemoteStoreSettings remoteStoreSettings,
-        Consumer<IndexShard> replicator
+        Consumer<IndexShard> replicator,
+        BiFunction<ShardId, ReplicationCheckpoint, ReplicationStats> segmentReplicationStatsProvider
     ) throws IOException {
         final IndexEventListener eventListener = freeze();
         Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrapperFactory = indexReaderWrapper
@@ -741,7 +745,8 @@ public IndexService newIndexService(
                 remoteStoreSettings,
                 fileCache,
                 compositeIndexSettings,
-                replicator
+                replicator,
+                segmentReplicationStatsProvider
             );
             success = true;
             return indexService;

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

@@ -105,6 +105,7 @@
 import org.opensearch.indices.mapper.MapperRegistry;
 import org.opensearch.indices.recovery.RecoverySettings;
 import org.opensearch.indices.recovery.RecoveryState;
+import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
 import org.opensearch.node.remotestore.RemoteStoreNodeAttribute;
 import org.opensearch.plugins.IndexStorePlugin;
@@ -197,6 +198,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust
     private final FileCache fileCache;
     private final CompositeIndexSettings compositeIndexSettings;
     private final Consumer<IndexShard> replicator;
+    private final BiFunction<ShardId, ReplicationCheckpoint, ReplicationStats> segmentReplicationStatsProvider;
 
     public IndexService(
         IndexSettings indexSettings,
@@ -235,7 +237,8 @@ public IndexService(
         RemoteStoreSettings remoteStoreSettings,
         FileCache fileCache,
         CompositeIndexSettings compositeIndexSettings,
-        Consumer<IndexShard> replicator
+        Consumer<IndexShard> replicator,
+        BiFunction<ShardId, ReplicationCheckpoint, ReplicationStats> segmentReplicationStatsProvider
     ) {
         super(indexSettings);
         this.allowExpensiveQueries = allowExpensiveQueries;
@@ -319,6 +322,7 @@ public IndexService(
         this.compositeIndexSettings = compositeIndexSettings;
         this.fileCache = fileCache;
         this.replicator = replicator;
+        this.segmentReplicationStatsProvider = segmentReplicationStatsProvider;
         updateFsyncTaskIfNecessary();
     }
 
@@ -395,7 +399,8 @@ public IndexService(
             remoteStoreSettings,
             null,
             null,
-            s -> {}
+            s -> {},
+            null
         );
     }
 
@@ -691,7 +696,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/shard/IndexShard.java

@@ -361,6 +361,7 @@ Runnable getGlobalCheckpointSyncer() {
      */
     private final ShardMigrationState shardMigrationState;
     private DiscoveryNodes discoveryNodes;
+    private final BiFunction<ShardId, ReplicationCheckpoint, 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 BiFunction<ShardId, ReplicationCheckpoint, 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() {
@@ -1784,7 +1787,8 @@ ReplicationCheckpoint computeReplicationCheckpoint(SegmentInfos segmentInfos) th
             segmentInfos.getVersion(),
             metadataMap.values().stream().mapToLong(StoreFileMetadata::length).sum(),
             getEngine().config().getCodec().getName(),
-            metadataMap
+            metadataMap,
+            System.nanoTime()
         );
         logger.trace("Recomputed ReplicationCheckpoint for shard {}", checkpoint);
         return checkpoint;
@@ -3228,17 +3232,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, this.getLatestReplicationCheckpoint());
+        }
+        return new ReplicationStats(0, 0, 0);
     }
 
     /**

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

@@ -136,6 +136,7 @@ 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(
@@ -149,7 +150,8 @@ private static ReplicationCheckpoint readCheckpointFromIndexInput(
             in.readLong(),
             in.readLong(),
             in.readString(),
-            toStoreFileMetadata(uploadedSegmentMetadataMap)
+            toStoreFileMetadata(uploadedSegmentMetadataMap),
+            in.readLong()
         );
     }
 

server/src/main/java/org/opensearch/indices/IndicesService.java

@@ -105,6 +105,7 @@
 import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.index.IndexService;
 import org.opensearch.index.IndexSettings;
+import org.opensearch.index.ReplicationStats;
 import org.opensearch.index.analysis.AnalysisRegistry;
 import org.opensearch.index.cache.request.ShardRequestCache;
 import org.opensearch.index.compositeindex.CompositeIndexSettings;
@@ -150,6 +151,7 @@
 import org.opensearch.indices.recovery.RecoveryListener;
 import org.opensearch.indices.recovery.RecoverySettings;
 import org.opensearch.indices.recovery.RecoveryState;
+import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
 import org.opensearch.indices.replication.common.ReplicationType;
 import org.opensearch.node.Node;
@@ -361,6 +363,7 @@ public class IndicesService extends AbstractLifecycleComponent
     private final FileCache fileCache;
     private final CompositeIndexSettings compositeIndexSettings;
     private final Consumer<IndexShard> replicator;
+    private final BiFunction<ShardId, ReplicationCheckpoint, ReplicationStats> segmentReplicationStatsProvider;
     private volatile int maxSizeInRequestCache;
 
     @Override
@@ -399,7 +402,8 @@ public IndicesService(
         RemoteStoreSettings remoteStoreSettings,
         FileCache fileCache,
         CompositeIndexSettings compositeIndexSettings,
-        Consumer<IndexShard> replicator
+        Consumer<IndexShard> replicator,
+        BiFunction<ShardId, ReplicationCheckpoint, ReplicationStats> segmentReplicationStatsProvider
     ) {
         this.settings = settings;
         this.threadPool = threadPool;
@@ -509,6 +513,7 @@ protected void closeInternal() {
         this.compositeIndexSettings = compositeIndexSettings;
         this.fileCache = fileCache;
         this.replicator = replicator;
+        this.segmentReplicationStatsProvider = segmentReplicationStatsProvider;
         this.maxSizeInRequestCache = INDICES_REQUEST_CACHE_MAX_SIZE_ALLOWED_IN_CACHE_SETTING.get(clusterService.getSettings());
         clusterService.getClusterSettings()
             .addSettingsUpdateConsumer(INDICES_REQUEST_CACHE_MAX_SIZE_ALLOWED_IN_CACHE_SETTING, this::setMaxSizeInRequestCache);
@@ -573,6 +578,7 @@ public IndicesService(
             remoteStoreSettings,
             null,
             null,
+            null,
             null
         );
     }
@@ -990,7 +996,8 @@ private synchronized IndexService createIndexService(
             this::getClusterDefaultRefreshInterval,
             this.recoverySettings,
             this.remoteStoreSettings,
-            replicator
+            replicator,
+            segmentReplicationStatsProvider
         );
     }
 

server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTarget.java

@@ -24,6 +24,7 @@
 import org.opensearch.common.util.CancellableThreads;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.index.store.Store;
 import org.opensearch.index.store.StoreFileMetadata;
@@ -39,6 +40,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Set;
+import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
 /**
@@ -161,7 +163,7 @@ public void writeFileChunk(
      *
      * @param listener {@link ActionListener} listener.
      */
-    public void startReplication(ActionListener<Void> listener) {
+    public void startReplication(ActionListener<Void> listener, BiConsumer<ReplicationCheckpoint, ShardId> checkpointUpdater) {
         cancellableThreads.setOnCancel((reason, beforeCancelEx) -> {
             throw new CancellableThreads.ExecutionCancelledException("replication was canceled reason [" + reason + "]");
         });
@@ -177,6 +179,8 @@ public void startReplication(ActionListener<Void> listener) {
         source.getCheckpointMetadata(getId(), checkpoint, checkpointInfoListener);
 
         checkpointInfoListener.whenComplete(checkpointInfo -> {
+            checkpointUpdater.accept(checkpointInfo.getCheckpoint(), this.indexShard.shardId());
+
             final List<StoreFileMetadata> filesToFetch = getFiles(checkpointInfo);
             state.setStage(SegmentReplicationState.Stage.GET_FILES);
             cancellableThreads.checkForCancel();

server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java

@@ -294,6 +294,7 @@ public synchronized void onNewCheckpoint(final ReplicationCheckpoint receivedChe
             return;
         }
         updateLatestReceivedCheckpoint(receivedCheckpoint, replicaShard);
+        replicator.updatePrimaryLastRefreshedCheckpoint(receivedCheckpoint, replicaShard.shardId());
         // Checks if replica shard is in the correct STARTED state to process checkpoints (avoids parallel replication events taking place)
         // This check ensures we do not try to process a received checkpoint while the shard is still recovering, yet we stored the latest
         // checkpoint to be replayed once the shard is Active.