@Test
  public void replicaThreadLeaderBasedReplicationFetchMissingKeysInMultipleCyclesTest() throws Exception {
    Map<DataNodeId, MockHost> hosts = new HashMap<>();
    hosts.put(remoteNodeInLocalDC, remoteHostInLocalDC);
    hosts.put(remoteNodeInRemoteDC, remoteHostInRemoteDC);
    int batchSize = 5;
    StoreKeyFactory storeKeyFactory = new BlobIdFactory(clusterMap);
    FindTokenHelper findTokenHelper = new MockFindTokenHelper(storeKeyFactory, replicationConfig);
    MockNetworkClientFactory mockNetworkClientFactory =
        new MockNetworkClientFactory(hosts, clusterMap, batchSize, findTokenHelper);

    Pair<StorageManager, ReplicationManager> managers =
        createStorageManagerAndReplicationManager(clusterMap, clusterMapConfig, mockHelixParticipant,
            mockNetworkClientFactory);
    StorageManager storageManager = managers.getFirst();
    MockReplicationManager replicationManager = (MockReplicationManager) managers.getSecond();

    /*
    Scenario:
      we have 3 nodes that have replicas belonging to same partitions:
        a) local node
        b) remote node in local DC
        c) remote node in remote DC

      Each node have few of its partitions as leaders and others are standby. They are randomly assigned during creation
      of replicas for mock partitions.

      we have 6 records (5 put messages + 1 ttl update message) to be fetched in each partition from remoteNodeInLocalDC and
      remoteNodeInRemoteDC.

      We will first fetch all 6 records in metadata exchange from remoteNodeInRemoteDC and store them for non-leader
      replicas and check if they are tracked correctly as they are fetched from remoteNodeInLocalDC in multiple cycles.

      Records to be fetched:
      1. put b0
      2. put b1
      3. put b2
      4. put b3
      5. put b4
      6. TTLUpdate b0

      Sequence of steps:
        1. Send metadata request to remoteNodeInRemoteDC with batch size of 6.
        2. We should get all messages in the metadata response b0, b1, b2, b3, b4 with b0 having ttlupdate flag set as well.
        3. Send get request for leader partitions. Non-leader partitions will store b0,b1,b2,b3 and b4 in
           remoteReplicaInfo.exchangeMetadataResponse.missingMessages
        4. Send metadata request to remoteNodeInLocalDC with batch size of 4.
        5. We should get only b0,b1,b2,b3.
        6. Send get request for all partitions since it is intra-dc replication.
        7. After the messages are written b0,b1,b2,b3 to local store, we should see them removed from
           remoteReplicaInfo.exchangeMetadataResponse.missingMessages for replica in remoteNodeInRemoteDC.
        8. Fetch remaining messages b4 and TTLUpdate of b0 from remoteNodeInLocalDC.
        9. remoteReplicaInfo.exchangeMetadataResponse.missingMessages for replica in remoteNodeInRemoteDC should be completely empty.

    */

    // set mock local stores on all remoteReplicaInfos which will used during replication.
    for (PartitionId partitionId : replicationManager.partitionToPartitionInfo.keySet()) {
      localHost.addStore(partitionId, null);
      Store localStore = localHost.getStore(partitionId);
      localStore.start();
      List<RemoteReplicaInfo> remoteReplicaInfos =
          replicationManager.partitionToPartitionInfo.get(partitionId).getRemoteReplicaInfos();
      remoteReplicaInfos.forEach(remoteReplicaInfo -> remoteReplicaInfo.setLocalStore(localStore));
    }

    // get remote replicas and replica thread for remote host on local datacenter
    ReplicaThread intraColoReplicaThread = replicationManager.dataNodeIdToReplicaThread.get(remoteNodeInLocalDC);
    List<RemoteReplicaInfo> remoteReplicaInfosForLocalDC =
        intraColoReplicaThread.getRemoteReplicaInfos().get(remoteNodeInLocalDC);

    // get remote replicas and replica thread for remote host on remote datacenter
    ReplicaThread crossColoReplicaThread = replicationManager.dataNodeIdToReplicaThread.get(remoteNodeInRemoteDC);
    List<RemoteReplicaInfo> remoteReplicaInfosForRemoteDC =
        crossColoReplicaThread.getRemoteReplicaInfos().get(remoteNodeInRemoteDC);

    // mock helix transition state from standby to leader for local leader partitions
    List<? extends ReplicaId> replicaIds = clusterMap.getReplicaIds(replicationManager.dataNodeId);
    for (ReplicaId replicaId : replicaIds) {
      MockReplicaId mockReplicaId = (MockReplicaId) replicaId;
      if (mockReplicaId.getReplicaState() == ReplicaState.LEADER) {
        MockPartitionId mockPartitionId = (MockPartitionId) replicaId.getPartitionId();
        mockHelixParticipant.onPartitionBecomeLeaderFromStandby(mockPartitionId.toPathString());
      }
    }

    List<PartitionId> partitionIds = clusterMap.getWritablePartitionIds(null);
    Map<PartitionId, List<StoreKey>> idsToBeIgnoredByPartition = new HashMap<>();
    Map<PartitionId, List<StoreKey>> idsToBeTtlUpdatedByPartition = new HashMap<>();

    int numOfMessages = 5;
    for (PartitionId id : partitionIds) {
      List<StoreKey> toBeTtlUpdated = new ArrayList<>();
      // Adding 5 put messages to remote hosts
      List<StoreKey> ids =
          addPutMessagesToReplicasOfPartition(id, Arrays.asList(remoteHostInLocalDC, remoteHostInRemoteDC),
              numOfMessages);
      //update the TTL of first put message
      StoreKey toTtlUpdateId = ids.get(0);
      addTtlUpdateMessagesToReplicasOfPartition(id, toTtlUpdateId,
          Arrays.asList(remoteHostInLocalDC, remoteHostInRemoteDC), UPDATED_EXPIRY_TIME_MS);
      toBeTtlUpdated.add(toTtlUpdateId);
      idsToBeTtlUpdatedByPartition.put(id, toBeTtlUpdated);
    }

    // Choose partitions that are leaders on both local and remote nodes
    Set<ReplicaId> remoteLeaderReplicasWithLeaderPartitionsOnLocalNode =
        getRemoteLeaderReplicasWithLeaderPartitionsOnLocalNode(clusterMap, replicationManager.dataNodeId,
            remoteNodeInRemoteDC);

    // fetch all messages in the metadata exchange from remote replica on remote datacenter by setting batchSize
    // in the connection to 6
    // There are 6 records in the remote host. 5 put records + 1 ttl update record. We will receive 5 messages in the
    // metadata exchange as ttl update will be merged with its put record and sent as single record.
    ((MockNetworkClient) crossColoReplicaThread.getNetworkClient()).setBatchSize(numOfMessages + 1);
    crossColoReplicaThread.replicate();
    List<ReplicaThread.ExchangeMetadataResponse> responseForRemoteNodeInRemoteDC =
        crossColoReplicaThread.getExchangeMetadataResponsesInEachCycle().get(remoteNodeInRemoteDC);
    assertEquals("Response should contain a response for each replica", remoteReplicaInfosForRemoteDC.size(),
        responseForRemoteNodeInRemoteDC.size());
    for (ReplicaThread.ExchangeMetadataResponse exchangeMetadataResponse : responseForRemoteNodeInRemoteDC) {
      // we should have received 5 messages bo, b1, b2, b3, b4 (ttl_update for b0 will be merged with b0 put message)
      assertEquals("mismatch in number of messages received from remote replica", numOfMessages,
          exchangeMetadataResponse.missingStoreMessages.size());
    }

    // Filter leader replicas to fetch missing keys
    List<RemoteReplicaInfo> leaderReplicas = new ArrayList<>();
    List<ReplicaThread.ExchangeMetadataResponse> exchangeMetadataResponseListForLeaderReplicas = new ArrayList<>();
    crossColoReplicaThread.getLeaderReplicaList(remoteReplicaInfosForRemoteDC, responseForRemoteNodeInRemoteDC,
        leaderReplicas, exchangeMetadataResponseListForLeaderReplicas);

    // verify that only leader replicas in remoteHost2 are chosen for fetching missing messages
    Set<ReplicaId> remoteReplicasToFetchInReplicaThread =
        leaderReplicas.stream().map(RemoteReplicaInfo::getReplicaId).collect(Collectors.toSet());
    assertThat("mismatch in leader remote replicas to fetch missing keys",
        remoteLeaderReplicasWithLeaderPartitionsOnLocalNode, is(remoteReplicasToFetchInReplicaThread));

    // verify that the remote token will be moved for leader replicas and will remain 0 for standby replicas as
    // missing messages are not fetched yet.
    for (int i = 0; i < remoteReplicaInfosForRemoteDC.size(); i++) {
      if (remoteLeaderReplicasWithLeaderPartitionsOnLocalNode.contains(
          remoteReplicaInfosForRemoteDC.get(i).getReplicaId())) {
        assertThat("remote Token should be updated for leader replica", remoteReplicaInfosForRemoteDC.get(i).getToken(),
            is(responseForRemoteNodeInRemoteDC.get(i).remoteToken));
      } else {
        assertThat("remote Token should not be updated for standby replica",
            remoteReplicaInfosForRemoteDC.get(i).getToken(), not(responseForRemoteNodeInRemoteDC.get(i).remoteToken));
      }
    }

    int numOfMessagesToBeFetchedFromRemoteHost1 = numOfMessages - 1;

    // Use blocking channel to make sure we fetch less than batch size messages.
    // Fetch only first 4 messages from remote host in local datacenter
    ((MockNetworkClient) intraColoReplicaThread.getNetworkClient()).setBatchSize(
        numOfMessagesToBeFetchedFromRemoteHost1);
    intraColoReplicaThread.replicate();
    List<ReplicaThread.ExchangeMetadataResponse> responseForRemoteNodeInLocalDC =
        intraColoReplicaThread.getExchangeMetadataResponsesInEachCycle().get(remoteHostInLocalDC.dataNodeId);

    // verify that remote token has moved forward for all partitions since it is intra-dc replication.
    for (int i = 0; i < responseForRemoteNodeInLocalDC.size(); i++) {
      assertEquals("token mismatch for intra-dc replication", remoteReplicaInfosForLocalDC.get(i).getToken(),
          responseForRemoteNodeInLocalDC.get(i).remoteToken);
    }

    // process metadata response for cross-colo replicas. Missing keys b0, b1, b2 and b3 must have come now via intra-dc replication
    for (RemoteReplicaInfo remoteReplicaInfo : remoteReplicaInfosForRemoteDC) {
      crossColoReplicaThread.processMissingKeysFromPreviousMetadataResponse(remoteReplicaInfo);
    }

    // verify that the missingStoreMessages size in standby replicas is 1 as b0, b1, b2 and b3 are fetched and b4 is pending
    for (RemoteReplicaInfo remoteReplicaInfo : remoteReplicaInfosForRemoteDC) {
      if (!remoteLeaderReplicasWithLeaderPartitionsOnLocalNode.contains(remoteReplicaInfo.getReplicaId())) {
        assertEquals("mismatch in number of missing messages for remote standby replicas after intra-dc replication",
            remoteReplicaInfo.getExchangeMetadataResponse().missingStoreMessages.size(), 1);
      }
    }

    //Fetch remaining one message (b4) from remote host in local datacenter
    intraColoReplicaThread.replicate();
    responseForRemoteNodeInLocalDC =
        intraColoReplicaThread.getExchangeMetadataResponsesInEachCycle().get(remoteHostInLocalDC.dataNodeId);
    assertEquals("Response should contain a response for each replica", remoteReplicaInfosForLocalDC.size(),
        responseForRemoteNodeInLocalDC.size());

    // verify remote token for intra-dc replicas
    for (int i = 0; i < responseForRemoteNodeInLocalDC.size(); i++) {
      assertThat("mismatch in remote token after all missing messages are received",
          remoteReplicaInfosForLocalDC.get(i).getToken(), is(responseForRemoteNodeInLocalDC.get(i).remoteToken));
    }

    // process metadata response for cross-colo replicas. Missing keys b0, b1, b2 and b3 must have come now via intra-dc replication
    for (RemoteReplicaInfo remoteReplicaInfo : remoteReplicaInfosForRemoteDC) {
      crossColoReplicaThread.processMissingKeysFromPreviousMetadataResponse(remoteReplicaInfo);
    }

    // Since all the missing messages (b0,b1,b2,b3,b4) for all partitions are fetched now via intra-dc replication,
    // remote token for all replicas (leader and standby) should move forward.
    for (int i = 0; i < responseForRemoteNodeInRemoteDC.size(); i++) {
      assertThat("mismatch in remote token after all missing messages are received",
          remoteReplicaInfosForRemoteDC.get(i).getToken(), is(responseForRemoteNodeInRemoteDC.get(i).remoteToken));
    }

    // compare the final state of all messages and buffers are in sync at local node and remoteNodeInLocalDC
    checkBlobMessagesAreEqualInLocalAndRemoteHosts(localHost, remoteHostInLocalDC, idsToBeIgnoredByPartition,
        idsToBeTtlUpdatedByPartition);
    // compare the final state of all messages and buffers are in sync at local node and remoteNodeInRemoteDC
    checkBlobMessagesAreEqualInLocalAndRemoteHosts(localHost, remoteHostInRemoteDC, idsToBeIgnoredByPartition,
        idsToBeTtlUpdatedByPartition);

    storageManager.shutdown();
  }
Test leader based replication to ensure token is advanced correctly for standby replicas when missing messages
are fetched via intra-dc replication in multiple cycles.
@throws Exception
ambry-replication/src/test/java/com/github/ambry/replication/LeaderBasedReplicationTest.java