Correctly updating the followerCheckpoint in stats api (#438)

Summary : We need to update followerCheckpoint after writing to the follower index.
Currently, we are not waiting for the writes and updating it with soon-to-be stale values

Signed-off-by: Gaurav Bafna <gbbafna@amazon.com>

Author: gbbafna

src/main/kotlin/com/amazon/elasticsearch/replication/task/shard/ShardReplicationTask.kt

@@ -222,6 +222,7 @@ class ShardReplicationTask(id: Long, type: String, action: String, description:
                                           TaskId(clusterService.nodeName, id), client, indexShard.localCheckpoint, followerClusterStats)
 
         val changeTracker = ShardReplicationChangesTracker(indexShard, replicationSettings)
+        followerClusterStats.stats[followerShardId]!!.followerCheckpoint = indexShard.localCheckpoint
         coroutineScope {
             while (isActive) {
                 rateLimiter.acquire()
@@ -278,7 +279,6 @@ class ShardReplicationTask(id: Long, type: String, action: String, description:
                 //hence renew retention lease with lastSyncedGlobalCheckpoint + 1 so that any shard that picks up shard replication task has data until then.
                 try {
                     retentionLeaseHelper.renewRetentionLease(leaderShardId, indexShard.lastSyncedGlobalCheckpoint + 1, followerShardId)
-                    followerClusterStats.stats[followerShardId]!!.followerCheckpoint = indexShard.lastSyncedGlobalCheckpoint
                     lastLeaseRenewalMillis = System.currentTimeMillis()
                 } catch (ex: Exception) {
                     when (ex) {

src/main/kotlin/com/amazon/elasticsearch/replication/task/shard/TranslogSequencer.kt

@@ -32,6 +32,7 @@ import org.elasticsearch.client.Client
 import org.elasticsearch.common.logging.Loggers
 import org.elasticsearch.index.shard.ShardId
 import org.elasticsearch.index.translog.Translog
+import com.amazon.elasticsearch.replication.util.indicesService
 import org.elasticsearch.tasks.TaskId
 import java.util.ArrayList
 import java.util.concurrent.ConcurrentHashMap
@@ -59,6 +60,9 @@ class TranslogSequencer(scope: CoroutineScope, private val replicationMetadata:
     private val log = Loggers.getLogger(javaClass, followerShardId)!!
     private val completed = CompletableDeferred<Unit>()
 
+    val followerIndexService = indicesService.indexServiceSafe(followerShardId.index)
+    val indexShard = followerIndexService.getShard(followerShardId.id)
+
     private val sequencer = scope.actor<Unit>(capacity = Channel.UNLIMITED) {
         // Exceptions thrown here will mark the channel as failed and the next attempt to send to the channel will
         // raise the same exception.  See [SendChannel.close] method for details.
@@ -92,6 +96,7 @@ class TranslogSequencer(scope: CoroutineScope, private val replicationMetadata:
                     val tookInNanos = System.nanoTime() - relativeStartNanos
                     followerClusterStats.stats[followerShardId]!!.totalWriteTime.addAndGet(TimeUnit.NANOSECONDS.toMillis(tookInNanos))
                     followerClusterStats.stats[followerShardId]!!.opsWritten.addAndGet(replayRequest.changes.size.toLong())
+                    followerClusterStats.stats[followerShardId]!!.followerCheckpoint = indexShard.localCheckpoint
                 }
                 highWatermark = next.changes.lastOrNull()?.seqNo() ?: highWatermark
             }

src/test/kotlin/com/amazon/elasticsearch/replication/integ/rest/StartReplicationIT.kt

@@ -1055,8 +1055,6 @@ class StartReplicationIT: MultiClusterRestTestCase() {
 
             followerClient.pauseReplication(followerIndex2)
             followerClient.stopReplication(followerIndex3)
-
-
             val stats = followerClient.followerStats()
             assertThat(stats.getValue("num_syncing_indices").toString()).isEqualTo("1")
             assertThat(stats.getValue("num_paused_indices").toString()).isEqualTo("1")
@@ -1066,6 +1064,7 @@ class StartReplicationIT: MultiClusterRestTestCase() {
             assertThat(stats.getValue("operations_read").toString()).isEqualTo("50")
             assertThat(stats.getValue("failed_read_requests").toString()).isEqualTo("0")
             assertThat(stats.getValue("failed_write_requests").toString()).isEqualTo("0")
+            assertThat(stats.getValue("follower_checkpoint").toString()).isEqualTo((docCount-1).toString())
             assertThat(stats.containsKey("index_stats"))
             assertThat(stats.size).isEqualTo(16)
         } finally {

src/test/kotlin/com/amazon/elasticsearch/replication/task/shard/TranslogSequencerTests.kt

@@ -26,21 +26,26 @@ import com.amazon.elasticsearch.replication.metadata.store.ReplicationStoreMetad
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.ObsoleteCoroutinesApi
 import kotlinx.coroutines.test.runBlockingTest
+import org.mockito.Mockito
 import org.assertj.core.api.Assertions.assertThat
 import org.elasticsearch.action.ActionListener
 import org.elasticsearch.action.ActionRequest
 import org.elasticsearch.action.ActionResponse
 import org.elasticsearch.action.ActionType
 import org.elasticsearch.action.support.replication.ReplicationResponse.ShardInfo
 import org.elasticsearch.common.settings.Settings
+import org.elasticsearch.index.IndexService
+import org.elasticsearch.index.shard.IndexShard
 import org.elasticsearch.index.shard.ShardId
 import org.elasticsearch.index.translog.Translog
+import org.elasticsearch.indices.IndicesService
 import org.elasticsearch.tasks.TaskId.EMPTY_TASK_ID
 import org.elasticsearch.test.ESTestCase
 import org.elasticsearch.test.ESTestCase.randomList
 import org.elasticsearch.test.client.NoOpClient
 import java.util.Locale
 
+
 @ObsoleteCoroutinesApi
 class TranslogSequencerTests : ESTestCase() {
 
@@ -87,6 +92,11 @@ class TranslogSequencerTests : ESTestCase() {
         val stats = FollowerClusterStats()
         stats.stats[followerShardId]  = FollowerShardMetric()
         val startSeqNo = randomNonNegativeLong()
+        indicesService = Mockito.mock(IndicesService::class.java)
+        val followerIndexService = Mockito.mock(IndexService::class.java)
+        val indexShard = Mockito.mock(IndexShard::class.java)
+        Mockito.`when`(indicesService.indexServiceSafe(followerShardId.index)).thenReturn(followerIndexService)
+        Mockito.`when`(followerIndexService.getShard(followerShardId.id)).thenReturn(indexShard)
         val sequencer = TranslogSequencer(this, replicationMetadata, followerShardId, leaderAlias, leaderIndex, EMPTY_TASK_ID,
                                           client, startSeqNo, stats)