/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0; you may not use this file except in compliance with the Elastic License
 * 2.0.
 */

package org.elasticsearch.xpack.ccr;

import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest;
import org.elasticsearch.action.admin.cluster.snapshots.restore.RestoreSnapshotRequest;
import org.elasticsearch.action.admin.cluster.state.ClusterStateResponse;
import org.elasticsearch.action.admin.indices.close.CloseIndexRequest;
import org.elasticsearch.action.admin.indices.stats.IndicesStatsRequest;
import org.elasticsearch.action.admin.indices.stats.IndicesStatsResponse;
import org.elasticsearch.action.admin.indices.stats.ShardStats;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.support.IndicesOptions;
import org.elasticsearch.action.support.PlainActionFuture;
import org.elasticsearch.action.support.broadcast.BroadcastResponse;
import org.elasticsearch.cluster.metadata.IndexMetadata;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.cluster.routing.ShardRouting;
import org.elasticsearch.cluster.service.ClusterService;
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.settings.Setting;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.unit.ByteSizeUnit;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.Index;
import org.elasticsearch.index.IndexService;
import org.elasticsearch.index.seqno.RetentionLease;
import org.elasticsearch.index.seqno.RetentionLeaseActions;
import org.elasticsearch.index.seqno.RetentionLeaseNotFoundException;
import org.elasticsearch.index.seqno.RetentionLeaseUtils;
import org.elasticsearch.index.shard.IndexShard;
import org.elasticsearch.index.shard.IndexShardClosedException;
import org.elasticsearch.indices.IndicesService;
import org.elasticsearch.plugins.Plugin;
import org.elasticsearch.snapshots.RestoreInfo;
import org.elasticsearch.snapshots.RestoreService;
import org.elasticsearch.test.transport.MockTransportService;
import org.elasticsearch.transport.ConnectTransportException;
import org.elasticsearch.transport.RemoteTransportException;
import org.elasticsearch.transport.Transport;
import org.elasticsearch.transport.TransportActionProxy;
import org.elasticsearch.transport.TransportMessageListener;
import org.elasticsearch.transport.TransportService;
import org.elasticsearch.xpack.CcrIntegTestCase;
import org.elasticsearch.xpack.ccr.action.repositories.ClearCcrRestoreSessionAction;
import org.elasticsearch.xpack.ccr.repository.CcrRepository;
import org.elasticsearch.xpack.core.ccr.action.ForgetFollowerAction;
import org.elasticsearch.xpack.core.ccr.action.PutFollowAction;
import org.elasticsearch.xpack.core.ccr.action.ResumeFollowAction;
import org.elasticsearch.xpack.core.ccr.action.UnfollowAction;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
import static org.elasticsearch.xpack.ccr.CcrRetentionLeases.retentionLeaseId;
import static org.hamcrest.Matchers.arrayWithSize;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.empty;
import static org.hamcrest.Matchers.emptyArray;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.greaterThan;
import static org.hamcrest.Matchers.hasSize;

public class CcrRetentionLeaseIT extends CcrIntegTestCase {

    public static final class RetentionLeaseRenewIntervalSettingPlugin extends Plugin {

        @Override
        public List<Setting<?>> getSettings() {
            return Collections.singletonList(CcrRetentionLeases.RETENTION_LEASE_RENEW_INTERVAL_SETTING);
        }

    }

    @Override
    protected Collection<Class<? extends Plugin>> nodePlugins() {
        return Stream.concat(
                super.nodePlugins().stream(),
                Stream.of(RetentionLeaseRenewIntervalSettingPlugin.class))
                .collect(Collectors.toList());
    }

    @Override
    protected Settings followerClusterSettings() {
        return Settings.builder()
                .put(super.followerClusterSettings())
                .put(CcrRetentionLeases.RETENTION_LEASE_RENEW_INTERVAL_SETTING.getKey(), TimeValue.timeValueMillis(200))
                .build();
    }

    private final IndicesOptions indicesOptions = IndicesOptions.strictSingleIndexNoExpandForbidClosed();

    private RestoreSnapshotRequest setUpRestoreSnapshotRequest(
            final String leaderIndex,
            final int numberOfShards,
            final int numberOfReplicas,
            final String followerIndex,
            final int numberOfDocuments) throws IOException {
        final ClusterUpdateSettingsRequest settingsRequest = new ClusterUpdateSettingsRequest().masterNodeTimeout(TimeValue.MAX_VALUE);
        final String chunkSize = new ByteSizeValue(randomFrom(4, 128, 1024), ByteSizeUnit.KB).getStringRep();
        settingsRequest.persistentSettings(Settings.builder().put(CcrSettings.RECOVERY_CHUNK_SIZE.getKey(), chunkSize));
        assertAcked(followerClient().admin().cluster().updateSettings(settingsRequest).actionGet());

        final String leaderClusterRepoName = CcrRepository.NAME_PREFIX + "leader_cluster";

        final Map<String, String> additionalSettings = new HashMap<>();
        additionalSettings.put(IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.getKey(), TimeValue.timeValueMillis(200).getStringRep());
        final String leaderIndexSettings = getIndexSettings(numberOfShards, numberOfReplicas, additionalSettings);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex)
            .setMasterNodeTimeout(TimeValue.MAX_VALUE).setSource(leaderIndexSettings, XContentType.JSON));
        ensureLeaderGreen(leaderIndex);

        logger.info("indexing [{}] docs", numberOfDocuments);
        for (int i = 0; i < numberOfDocuments; i++) {
            final String source = String.format(Locale.ROOT, "{\"f\":%d}", i);
            leaderClient().prepareIndex(leaderIndex, "doc", Integer.toString(i)).setSource(source, XContentType.JSON).get();
            if (rarely()) {
                leaderClient().admin().indices().prepareFlush(leaderIndex).setForce(true).setWaitIfOngoing(true).get();
            }
        }

        leaderClient().admin().indices().prepareFlush(leaderIndex).setForce(true).setWaitIfOngoing(true).get();

        final Settings.Builder settingsBuilder = Settings.builder()
                .put(IndexMetadata.SETTING_INDEX_PROVIDED_NAME, followerIndex)
                .put(CcrSettings.CCR_FOLLOWING_INDEX_SETTING.getKey(), true);
        return new RestoreSnapshotRequest(leaderClusterRepoName, CcrRepository.LATEST)
                .indexSettings(settingsBuilder)
                .indices(leaderIndex)
                .indicesOptions(indicesOptions)
                .renamePattern("^(.*)$")
                .renameReplacement(followerIndex)
                .masterNodeTimeout(TimeValue.MAX_VALUE);
    }

    public void testRetentionLeaseIsTakenAtTheStartOfRecovery() throws Exception {
        final String leaderIndex = "leader";
        final int numberOfShards = randomIntBetween(1, 3);
        final int numberOfReplicas = between(0, 1);
        final String followerIndex = "follower";
        final int numberOfDocuments = scaledRandomIntBetween(1, 8192);
        final RestoreSnapshotRequest restoreRequest =
                setUpRestoreSnapshotRequest(leaderIndex, numberOfShards, numberOfReplicas, followerIndex, numberOfDocuments);
        final RestoreService restoreService = getFollowerCluster().getCurrentMasterNodeInstance(RestoreService.class);
        final ClusterService clusterService = getFollowerCluster().getCurrentMasterNodeInstance(ClusterService.class);

        final PlainActionFuture<RestoreInfo> future = PlainActionFuture.newFuture();
        restoreService.restoreSnapshot(restoreRequest, waitForRestore(clusterService, future));

        // ensure that a retention lease has been put in place on each shard
        assertBusy(() -> {
            final IndicesStatsResponse stats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
            assertNotNull(stats.getShards());
            assertThat(stats.getShards(), arrayWithSize(numberOfShards * (1 + numberOfReplicas)));
            final List<ShardStats> shardsStats = getShardsStats(stats);
            for (int i = 0; i < numberOfShards * (1 + numberOfReplicas); i++) {
                assertNotNull(shardsStats.get(i).getRetentionLeaseStats());
                final Map<String, RetentionLease> currentRetentionLeases = RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                    shardsStats.get(i).getRetentionLeaseStats().retentionLeases());
                assertThat(Strings.toString(shardsStats.get(i)), currentRetentionLeases.values(), hasSize(1));
                final RetentionLease retentionLease = currentRetentionLeases.values().iterator().next();
                assertThat(retentionLease.id(), equalTo(getRetentionLeaseId(followerIndex, leaderIndex)));
            }
        });

        final RestoreInfo restoreInfo = future.actionGet();

        assertEquals(restoreInfo.totalShards(), restoreInfo.successfulShards());
        assertEquals(0, restoreInfo.failedShards());
        for (int i = 0; i < numberOfDocuments; ++i) {
            assertExpectedDocument(followerIndex, i);
        }

    }

    public void testRetentionLeaseIsRenewedDuringRecovery() throws Exception {
        final String leaderIndex = "leader";
        final int numberOfShards = randomIntBetween(1, 3);
        final int numberOfReplicas = between(0, 1);
        final String followerIndex = "follower";
        final int numberOfDocuments = scaledRandomIntBetween(1, 8192);
        final RestoreSnapshotRequest restoreRequest =
                setUpRestoreSnapshotRequest(leaderIndex, numberOfShards, numberOfReplicas, followerIndex, numberOfDocuments);
        final RestoreService restoreService = getFollowerCluster().getCurrentMasterNodeInstance(RestoreService.class);
        final ClusterService clusterService = getFollowerCluster().getCurrentMasterNodeInstance(ClusterService.class);

        final CountDownLatch latch = new CountDownLatch(1);

        // block the recovery from completing; this ensures the background sync is still running
        final ClusterStateResponse followerClusterState = followerClient().admin().cluster().prepareState().clear().setNodes(true).get();
        for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
            final MockTransportService senderTransportService =
                    (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
            senderTransportService.addSendBehavior(
                    (connection, requestId, action, request, options) -> {
                        if (ClearCcrRestoreSessionAction.NAME.equals(action)
                                || TransportActionProxy.getProxyAction(ClearCcrRestoreSessionAction.NAME).equals(action)) {
                            try {
                                latch.await();
                            } catch (final InterruptedException e) {
                                fail(e.toString());
                            }
                        }
                        connection.sendRequest(requestId, action, request, options);
                    });
        }

        final PlainActionFuture<RestoreInfo> future = PlainActionFuture.newFuture();
        restoreService.restoreSnapshot(restoreRequest, waitForRestore(clusterService, future));

        try {
            assertRetentionLeaseRenewal(numberOfShards, numberOfReplicas, followerIndex, leaderIndex);
            latch.countDown();
        } finally {
            for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
                final MockTransportService senderTransportService =
                        (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
                senderTransportService.clearAllRules();
            }
        }

        final RestoreInfo restoreInfo = future.actionGet();

        assertEquals(restoreInfo.totalShards(), restoreInfo.

                successfulShards());

        assertEquals(0, restoreInfo.failedShards());
        for (int i = 0; i < numberOfDocuments; i++) {
            assertExpectedDocument(followerIndex, i);
        }

    }

    public void testRetentionLeasesAreNotBeingRenewedAfterRecoveryCompletes() throws Exception {
        final String leaderIndex = "leader";
        final int numberOfShards = randomIntBetween(1, 3);
        final int numberOfReplicas = between(0, 1);
        final String followerIndex = "follower";
        final int numberOfDocuments = scaledRandomIntBetween(1, 8192);
        final RestoreSnapshotRequest restoreRequest =
                setUpRestoreSnapshotRequest(leaderIndex, numberOfShards, numberOfReplicas, followerIndex, numberOfDocuments);
        final RestoreService restoreService = getFollowerCluster().getCurrentMasterNodeInstance(RestoreService.class);
        final ClusterService clusterService = getFollowerCluster().getCurrentMasterNodeInstance(ClusterService.class);

        final PlainActionFuture<RestoreInfo> future = PlainActionFuture.newFuture();
        restoreService.restoreSnapshot(restoreRequest, waitForRestore(clusterService, future));

        final RestoreInfo restoreInfo = future.actionGet();
        final long start = System.nanoTime();

        /*
         * We want to ensure that the retention leases have been synced to all shard copies, as otherwise they might sync between the two
         * times that we sample the retention leases, which would cause our check to fail.
         */
        final TimeValue syncIntervalSetting = IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.get(
                leaderClient()
                        .admin()
                        .indices()
                        .prepareGetSettings(leaderIndex)
                        .get()
                        .getIndexToSettings()
                        .get(leaderIndex));
        final long syncEnd = System.nanoTime();
        Thread.sleep(Math.max(0, randomIntBetween(2, 4) * syncIntervalSetting.millis() - TimeUnit.NANOSECONDS.toMillis(syncEnd - start)));

        final ClusterStateResponse leaderIndexClusterState =
                leaderClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(leaderIndex).get();
        final String leaderUUID = leaderIndexClusterState.getState().metadata().index(leaderIndex).getIndexUUID();

        /*
         * We want to ensure that the background renewal is cancelled at the end of recovery. To do this, we will sleep a small multiple
         * of the renew interval. If the renews are not cancelled, we expect that a renewal would have been sent while we were sleeping.
         * After we wake up, it should be the case that the retention leases are the same (same timestamp) as that indicates that they were
         * not renewed while we were sleeping.
         */
        assertBusy(() -> {
            // sample the leases after recovery
            final List< Map<String, RetentionLease>> retentionLeases = new ArrayList<>();
            assertBusy(() -> {
                retentionLeases.clear();
                final IndicesStatsResponse stats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
                assertNotNull(stats.getShards());
                assertThat(stats.getShards(), arrayWithSize(numberOfShards * (1 + numberOfReplicas)));
                final List<ShardStats> shardsStats = getShardsStats(stats);
                for (int i = 0; i < numberOfShards * (1 + numberOfReplicas); i++) {
                    assertNotNull(shardsStats.get(i).getRetentionLeaseStats());
                    final Map<String, RetentionLease> currentRetentionLeases
                        = RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                            shardsStats.get(i).getRetentionLeaseStats().retentionLeases());
                    assertThat(Strings.toString(shardsStats.get(i)), currentRetentionLeases.values(), hasSize(1));
                    final ClusterStateResponse followerIndexClusterState =
                        followerClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(followerIndex).get();
                    final String followerUUID = followerIndexClusterState.getState().metadata().index(followerIndex).getIndexUUID();
                    final RetentionLease retentionLease =
                        currentRetentionLeases.values().iterator().next();
                    final String expectedRetentionLeaseId = retentionLeaseId(
                        getFollowerCluster().getClusterName(),
                        new Index(followerIndex, followerUUID),
                        getLeaderCluster().getClusterName(),
                        new Index(leaderIndex, leaderUUID));
                    assertThat(retentionLease.id(), equalTo(expectedRetentionLeaseId));
                    retentionLeases.add(currentRetentionLeases);
                }
            });
            // sleep a small multiple of the renew interval
            final TimeValue renewIntervalSetting = CcrRetentionLeases.RETENTION_LEASE_RENEW_INTERVAL_SETTING.get(followerClusterSettings());
            final long renewEnd = System.nanoTime();
            Thread.sleep(
                Math.max(0, randomIntBetween(2, 4) * renewIntervalSetting.millis() - TimeUnit.NANOSECONDS.toMillis(renewEnd - start)));

            // now ensure that the retention leases are the same
            final IndicesStatsResponse stats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
            assertNotNull(stats.getShards());
            assertThat(stats.getShards(), arrayWithSize(numberOfShards * (1 + numberOfReplicas)));
            final List<ShardStats> shardsStats = getShardsStats(stats);
            for (int i = 0; i < numberOfShards * (1 + numberOfReplicas); i++) {
                if (shardsStats.get(i).getShardRouting().primary() == false) {
                    continue;
                }
                assertNotNull(shardsStats.get(i).getRetentionLeaseStats());
                final Map<String, RetentionLease> currentRetentionLeases = RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                    shardsStats.get(i).getRetentionLeaseStats().retentionLeases());
                assertThat(Strings.toString(shardsStats.get(i)), currentRetentionLeases.values(), hasSize(1));
                final ClusterStateResponse followerIndexClusterState =
                        followerClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(followerIndex).get();
                final String followerUUID = followerIndexClusterState.getState().metadata().index(followerIndex).getIndexUUID();
                final RetentionLease retentionLease =
                        currentRetentionLeases.values().iterator().next();
                assertThat(retentionLease.id(), equalTo(getRetentionLeaseId(followerIndex, followerUUID, leaderIndex, leaderUUID)));
                // we assert that retention leases are being renewed by an increase in the timestamp
                assertThat(retentionLease.timestamp(), equalTo(retentionLeases.get(i).values().iterator().next().timestamp()));
            }
        });

        assertEquals(restoreInfo.totalShards(), restoreInfo.successfulShards());
        assertEquals(0, restoreInfo.failedShards());
        for (int i = 0; i < numberOfDocuments; ++i) {
            assertExpectedDocument(followerIndex, i);
        }
    }

    public void testUnfollowRemovesRetentionLeases() throws Exception {
        final String leaderIndex = "leader";
        final String followerIndex = "follower";
        final int numberOfShards = randomIntBetween(1, 4);
        final String leaderIndexSettings = getIndexSettings(numberOfShards, 0);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex).setSource(leaderIndexSettings, XContentType.JSON).get());
        ensureLeaderYellow(leaderIndex);
        final PutFollowAction.Request followRequest = putFollow(leaderIndex, followerIndex);
        followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();

        ensureFollowerGreen(true, followerIndex);

        final String retentionLeaseId = getRetentionLeaseId(followerIndex, leaderIndex);

        final IndicesStatsResponse stats =
                leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
        final List<ShardStats> shardsStats = getShardsStats(stats);
        for (final ShardStats shardStats : shardsStats) {
            final Map<String, RetentionLease> retentionLeases = RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                shardStats.getRetentionLeaseStats().retentionLeases());
            assertThat(Strings.toString(shardStats), retentionLeases.values(), hasSize(1));
            assertThat(retentionLeases.values().iterator().next().id(), equalTo(retentionLeaseId));
        }

        // we will sometimes fake that some of the retention leases are already removed on the leader shard
        final Set<Integer> shardIds =
                new HashSet<>(randomSubsetOf(
                        randomIntBetween(0, numberOfShards),
                        IntStream.range(0, numberOfShards).boxed().collect(Collectors.toSet())));

        final ClusterStateResponse followerClusterState = followerClient().admin().cluster().prepareState().clear().setNodes(true).get();
        try {
            for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
                final MockTransportService senderTransportService =
                        (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
                senderTransportService.addSendBehavior(
                        (connection, requestId, action, request, options) -> {
                            if (RetentionLeaseActions.Remove.ACTION_NAME.equals(action)
                                    || TransportActionProxy.getProxyAction(RetentionLeaseActions.Remove.ACTION_NAME).equals(action)) {
                                final RetentionLeaseActions.RemoveRequest removeRequest = (RetentionLeaseActions.RemoveRequest) request;
                                if (shardIds.contains(removeRequest.getShardId().id())) {
                                    final String primaryShardNodeId =
                                            getLeaderCluster()
                                                    .clusterService()
                                                    .state()
                                                    .routingTable()
                                                    .index(leaderIndex)
                                                    .shard(removeRequest.getShardId().id())
                                                    .primaryShard()
                                                    .currentNodeId();
                                    final String primaryShardNodeName =
                                            getLeaderCluster().clusterService().state().nodes().get(primaryShardNodeId).getName();
                                    final IndexShard primary =
                                            getLeaderCluster()
                                                    .getInstance(IndicesService.class, primaryShardNodeName)
                                                    .getShardOrNull(removeRequest.getShardId());
                                    final CountDownLatch latch = new CountDownLatch(1);
                                    primary.removeRetentionLease(
                                            retentionLeaseId,
                                            ActionListener.wrap(r -> latch.countDown(), e -> fail(e.toString())));
                                    try {
                                        latch.await();
                                    } catch (final InterruptedException e) {
                                        Thread.currentThread().interrupt();
                                        fail(e.toString());
                                    }
                                }
                            }
                            connection.sendRequest(requestId, action, request, options);
                        });
            }

            pauseFollow(followerIndex);
            assertAcked(followerClient().admin().indices().close(new CloseIndexRequest(followerIndex)).actionGet());
            assertAcked(followerClient().execute(UnfollowAction.INSTANCE, new UnfollowAction.Request(followerIndex)).actionGet());

            final IndicesStatsResponse afterUnfollowStats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
            final List<ShardStats> afterUnfollowShardsStats = getShardsStats(afterUnfollowStats);
            for (final ShardStats shardStats : afterUnfollowShardsStats) {
                assertThat(Strings.toString(shardStats), RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                    shardStats.getRetentionLeaseStats().retentionLeases()).values(), empty());
            }
        } finally {
            for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
                final MockTransportService senderTransportService =
                        (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
                senderTransportService.clearAllRules();
            }
        }
    }

    public void testUnfollowFailsToRemoveRetentionLeases() throws Exception {
        final String leaderIndex = "leader";
        final String followerIndex = "follower";
        final int numberOfShards = randomIntBetween(1, 4);
        final String leaderIndexSettings = getIndexSettings(numberOfShards, 0);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex).setSource(leaderIndexSettings, XContentType.JSON).get());
        ensureLeaderYellow(leaderIndex);
        final PutFollowAction.Request followRequest = putFollow(leaderIndex, followerIndex);
        followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();

        ensureFollowerGreen(true, followerIndex);

        pauseFollow(followerIndex);
        followerClient().admin().indices().close(new CloseIndexRequest(followerIndex).masterNodeTimeout(TimeValue.MAX_VALUE)).actionGet();

        // we will disrupt requests to remove retention leases for these random shards
        final Set<Integer> shardIds =
                new HashSet<>(randomSubsetOf(
                                randomIntBetween(1, numberOfShards),
                                IntStream.range(0, numberOfShards).boxed().collect(Collectors.toSet())));

        final ClusterStateResponse followerClusterState = followerClient().admin().cluster().prepareState().clear().setNodes(true).get();
        try {
            for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
                final MockTransportService senderTransportService =
                        (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
                senderTransportService.addSendBehavior(
                        (connection, requestId, action, request, options) -> {
                            if (RetentionLeaseActions.Remove.ACTION_NAME.equals(action)
                                    || TransportActionProxy.getProxyAction(RetentionLeaseActions.Remove.ACTION_NAME).equals(action)) {
                                final RetentionLeaseActions.RemoveRequest removeRequest = (RetentionLeaseActions.RemoveRequest) request;
                                if (shardIds.contains(removeRequest.getShardId().id())) {
                                    throw randomBoolean()
                                            ? new ConnectTransportException(connection.getNode(), "connection failed")
                                            : new IndexShardClosedException(removeRequest.getShardId());
                                }
                            }
                            connection.sendRequest(requestId, action, request, options);
                        });
            }

            final ElasticsearchException e = expectThrows(
                    ElasticsearchException.class,
                    () -> followerClient().execute(UnfollowAction.INSTANCE, new UnfollowAction.Request(followerIndex)).actionGet());

            final ClusterStateResponse followerIndexClusterState =
                    followerClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(followerIndex).get();
            final String followerUUID = followerIndexClusterState.getState().metadata().index(followerIndex).getIndexUUID();

            final ClusterStateResponse leaderIndexClusterState =
                    leaderClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(leaderIndex).get();
            final String leaderUUID = leaderIndexClusterState.getState().metadata().index(leaderIndex).getIndexUUID();

            assertThat(
                    e.getMetadata("es.failed_to_remove_retention_leases"),
                    contains(retentionLeaseId(
                            getFollowerCluster().getClusterName(),
                            new Index(followerIndex, followerUUID),
                            getLeaderCluster().getClusterName(),
                            new Index(leaderIndex, leaderUUID))));
        } finally {
            for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
                final MockTransportService senderTransportService =
                        (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
                senderTransportService.clearAllRules();
            }
        }
    }

    public void testRetentionLeaseRenewedWhileFollowing() throws Exception {
        final String leaderIndex = "leader";
        final String followerIndex = "follower";
        final int numberOfShards = randomIntBetween(1, 4);
        final int numberOfReplicas = randomIntBetween(0, 1);
        final Map<String, String> additionalIndexSettings = new HashMap<>();
        additionalIndexSettings.put(
                IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.getKey(),
                TimeValue.timeValueMillis(200).getStringRep());
        final String leaderIndexSettings = getIndexSettings(numberOfShards, numberOfReplicas, additionalIndexSettings);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex).setSource(leaderIndexSettings, XContentType.JSON).get());
        ensureLeaderYellow(leaderIndex);
        final PutFollowAction.Request followRequest = putFollow(leaderIndex, followerIndex);
        followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();

        ensureFollowerGreen(true, followerIndex);
        assertRetentionLeaseRenewal(numberOfShards, numberOfReplicas, followerIndex, leaderIndex);
    }

    public void testRetentionLeaseAdvancesWhileFollowing() throws Exception {
        final String leaderIndex = "leader";
        final String followerIndex = "follower";
        final int numberOfShards = randomIntBetween(1, 4);
        final int numberOfReplicas = randomIntBetween(0, 1);
        final Map<String, String> additionalIndexSettings = new HashMap<>();
        additionalIndexSettings.put(
                IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.getKey(),
                TimeValue.timeValueMillis(200).getStringRep());
        final String leaderIndexSettings = getIndexSettings(numberOfShards, numberOfReplicas, additionalIndexSettings);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex).setSource(leaderIndexSettings, XContentType.JSON).get());
        ensureLeaderYellow(leaderIndex);
        final PutFollowAction.Request followRequest = putFollow(leaderIndex, followerIndex);
        followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();

        ensureFollowerGreen(true, followerIndex);

        final int numberOfDocuments = randomIntBetween(128, 2048);
        logger.debug("indexing [{}] docs", numberOfDocuments);
        for (int i = 0; i < numberOfDocuments; i++) {
            final String source = String.format(Locale.ROOT, "{\"f\":%d}", i);
            leaderClient().prepareIndex(leaderIndex, "doc", Integer.toString(i)).setSource(source, XContentType.JSON).get();
            if (rarely()) {
                leaderClient().admin().indices().prepareFlush(leaderIndex).setForce(true).setWaitIfOngoing(true).get();
            }
        }

        // wait until the follower global checkpoints have caught up to the leader
        assertIndexFullyReplicatedToFollower(leaderIndex, followerIndex);

        final List<ShardStats> leaderShardsStats = getShardsStats(leaderClient().admin().indices().prepareStats(leaderIndex).get());
        final Map<Integer, Long> leaderGlobalCheckpoints = new HashMap<>();
        for (final ShardStats leaderShardStats : leaderShardsStats) {
            final ShardRouting routing = leaderShardStats.getShardRouting();
            if (routing.primary() == false) {
                continue;
            }
            leaderGlobalCheckpoints.put(routing.id(), leaderShardStats.getSeqNoStats().getGlobalCheckpoint());
        }

        // now assert that the retention leases have advanced to the global checkpoints
        assertBusy(() -> {
            final IndicesStatsResponse stats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
            assertNotNull(stats.getShards());
            assertThat(stats.getShards(), arrayWithSize(numberOfShards * (1 + numberOfReplicas)));
            final List<ShardStats> shardsStats = getShardsStats(stats);
            for (int i = 0; i < numberOfShards * (1 + numberOfReplicas); i++) {
                assertNotNull(shardsStats.get(i).getRetentionLeaseStats());
                final Map<String, RetentionLease> currentRetentionLeases = RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                    shardsStats.get(i).getRetentionLeaseStats().retentionLeases());
                assertThat(Strings.toString(shardsStats.get(i)), currentRetentionLeases.values(), hasSize(1));
                final RetentionLease retentionLease =
                        currentRetentionLeases.values().iterator().next();
                assertThat(retentionLease.id(), equalTo(getRetentionLeaseId(followerIndex, leaderIndex)));
                // we assert that retention leases are being advanced
                assertThat(
                        retentionLease.retainingSequenceNumber(),
                        equalTo(leaderGlobalCheckpoints.get(shardsStats.get(i).getShardRouting().id()) + 1));
            }
        });
    }

    public void testRetentionLeaseRenewalIsCancelledWhenFollowingIsPaused() throws Exception {
        final String leaderIndex = "leader";
        final String followerIndex = "follower";
        final int numberOfShards = randomIntBetween(1, 4);
        final int numberOfReplicas = randomIntBetween(0, 1);
        final Map<String, String> additionalIndexSettings = new HashMap<>();
        additionalIndexSettings.put(
                IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.getKey(),
                TimeValue.timeValueMillis(200).getStringRep());
        final String leaderIndexSettings = getIndexSettings(numberOfShards, numberOfReplicas, additionalIndexSettings);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex).setSource(leaderIndexSettings, XContentType.JSON).get());
        ensureLeaderYellow(leaderIndex);
        final PutFollowAction.Request followRequest = putFollow(leaderIndex, followerIndex);
        followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();

        ensureFollowerGreen(true, followerIndex);

        final long start = System.nanoTime();
        pauseFollow(followerIndex);

        /*
         * We want to ensure that the retention leases have been synced to all shard copies, as otherwise they might sync between the two
         * times that we sample the retention leases, which would cause our check to fail.
         */
        final TimeValue syncIntervalSetting = IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.get(
                leaderClient()
                        .admin()
                        .indices()
                        .prepareGetSettings(leaderIndex)
                        .get()
                        .getIndexToSettings()
                        .get(leaderIndex));
        final long syncEnd = System.nanoTime();
        Thread.sleep(Math.max(0, randomIntBetween(2, 4) * syncIntervalSetting.millis() - TimeUnit.NANOSECONDS.toMillis(syncEnd - start)));

        final ClusterStateResponse leaderIndexClusterState =
                leaderClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(leaderIndex).get();
        final String leaderUUID = leaderIndexClusterState.getState().metadata().index(leaderIndex).getIndexUUID();
        /*
         * We want to ensure that the background renewal is cancelled after pausing. To do this, we will sleep a small multiple of the renew
         * interval. If the renews are not cancelled, we expect that a renewal would have been sent while we were sleeping. After we wake
         * up, it should be the case that the retention leases are the same (same timestamp) as that indicates that they were not renewed
         * while we were sleeping.
         */
        assertBusy(() -> {
            // sample the leases after pausing
            final List<Map<String, RetentionLease>> retentionLeases = new ArrayList<>();
            assertBusy(() -> {
                retentionLeases.clear();
                final IndicesStatsResponse stats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
                assertNotNull(stats.getShards());
                assertThat(stats.getShards(), arrayWithSize(numberOfShards * (1 + numberOfReplicas)));
                final List<ShardStats> shardsStats = getShardsStats(stats);
                for (int i = 0; i < numberOfShards * (1 + numberOfReplicas); i++) {
                    assertNotNull(shardsStats.get(i).getRetentionLeaseStats());
                    final Map<String, RetentionLease> currentRetentionLeases
                        = RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                            shardsStats.get(i).getRetentionLeaseStats().retentionLeases());
                    assertThat(Strings.toString(shardsStats.get(i)), currentRetentionLeases.values(), hasSize(1));
                    final ClusterStateResponse followerIndexClusterState =
                        followerClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(followerIndex).get();
                    final String followerUUID = followerIndexClusterState.getState().metadata().index(followerIndex).getIndexUUID();
                    final RetentionLease retentionLease =
                        currentRetentionLeases.values().iterator().next();
                    final String expectedRetentionLeaseId = retentionLeaseId(
                        getFollowerCluster().getClusterName(),
                        new Index(followerIndex, followerUUID),
                        getLeaderCluster().getClusterName(),
                        new Index(leaderIndex, leaderUUID));
                    assertThat(retentionLease.id(), equalTo(expectedRetentionLeaseId));
                    retentionLeases.add(currentRetentionLeases);
                }
            });
            // sleep a small multiple of the renew interval
            final TimeValue renewIntervalSetting = CcrRetentionLeases.RETENTION_LEASE_RENEW_INTERVAL_SETTING.get(followerClusterSettings());
            final long renewEnd = System.nanoTime();
            Thread.sleep(
                Math.max(0, randomIntBetween(2, 4) * renewIntervalSetting.millis() - TimeUnit.NANOSECONDS.toMillis(renewEnd - start)));

            // now ensure that the retention leases are the same
            final IndicesStatsResponse stats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
            assertNotNull(stats.getShards());
            assertThat(stats.getShards(), arrayWithSize(numberOfShards * (1 + numberOfReplicas)));
            final List<ShardStats> shardsStats = getShardsStats(stats);
            for (int i = 0; i < numberOfShards * (1 + numberOfReplicas); i++) {
                if (shardsStats.get(i).getShardRouting().primary() == false) {
                    continue;
                }
                assertNotNull(shardsStats.get(i).getRetentionLeaseStats());
                final Map<String, RetentionLease> currentRetentionLeases = RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                    shardsStats.get(i).getRetentionLeaseStats().retentionLeases());
                assertThat(Strings.toString(shardsStats.get(i)), currentRetentionLeases.values(), hasSize(1));
                final ClusterStateResponse followerIndexClusterState =
                        followerClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(followerIndex).get();
                final String followerUUID = followerIndexClusterState.getState().metadata().index(followerIndex).getIndexUUID();
                final RetentionLease retentionLease =
                        currentRetentionLeases.values().iterator().next();
                assertThat(retentionLease.id(), equalTo(getRetentionLeaseId(followerIndex, followerUUID, leaderIndex, leaderUUID)));
                // we assert that retention leases are not being renewed by an unchanged timestamp
                assertThat(retentionLease.timestamp(), equalTo(retentionLeases.get(i).values().iterator().next().timestamp()));
            }
        });
    }

    public void testRetentionLeaseRenewalIsResumedWhenFollowingIsResumed() throws Exception {
        final String leaderIndex = "leader";
        final String followerIndex = "follower";
        final int numberOfShards = randomIntBetween(1, 4);
        final int numberOfReplicas = randomIntBetween(0, 1);
        final Map<String, String> additionalIndexSettings = new HashMap<>();
        additionalIndexSettings.put(
                IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.getKey(),
                TimeValue.timeValueMillis(200).getStringRep());
        final String leaderIndexSettings = getIndexSettings(numberOfShards, numberOfReplicas, additionalIndexSettings);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex).setSource(leaderIndexSettings, XContentType.JSON).get());
        ensureLeaderYellow(leaderIndex);
        final PutFollowAction.Request followRequest = putFollow(leaderIndex, followerIndex);
        followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();

        ensureFollowerGreen(true, followerIndex);

        pauseFollow(followerIndex);

        followerClient().execute(ResumeFollowAction.INSTANCE, resumeFollow(followerIndex)).actionGet();

        ensureFollowerGreen(true, followerIndex);

        assertRetentionLeaseRenewal(numberOfShards, numberOfReplicas, followerIndex, leaderIndex);
    }

    public void testRetentionLeaseIsAddedIfItDisappearsWhileFollowing() throws Exception {
        final String leaderIndex = "leader";
        final String followerIndex = "follower";
        final int numberOfShards = 1;
        final int numberOfReplicas = 1;
        final Map<String, String> additionalIndexSettings = new HashMap<>();
        additionalIndexSettings.put(
                IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.getKey(),
                TimeValue.timeValueMillis(200).getStringRep());
        final String leaderIndexSettings = getIndexSettings(numberOfShards, numberOfReplicas, additionalIndexSettings);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex).setSource(leaderIndexSettings, XContentType.JSON).get());
        ensureLeaderYellow(leaderIndex);
        final PutFollowAction.Request followRequest = putFollow(leaderIndex, followerIndex);
        followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();

        ensureFollowerGreen(true, followerIndex);

        final CountDownLatch latch = new CountDownLatch(1);

        final ClusterStateResponse followerClusterState = followerClient().admin().cluster().prepareState().clear().setNodes(true).get();
        try {
            for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
                final MockTransportService senderTransportService =
                    (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
                senderTransportService.addSendBehavior(
                    (connection, requestId, action, request, options) -> {
                        if (RetentionLeaseActions.Renew.ACTION_NAME.equals(action)
                            || TransportActionProxy.getProxyAction(RetentionLeaseActions.Renew.ACTION_NAME).equals(action)) {
                            final RetentionLeaseActions.RenewRequest renewRequest = (RetentionLeaseActions.RenewRequest) request;
                            final String retentionLeaseId = getRetentionLeaseId(followerIndex, leaderIndex);
                            if (retentionLeaseId.equals(renewRequest.getId())) {
                                logger.info("--> intercepting renewal request for retention lease [{}]", retentionLeaseId);
                                senderTransportService.clearAllRules();
                                final String primaryShardNodeId =
                                    getLeaderCluster()
                                        .clusterService()
                                        .state()
                                        .routingTable()
                                        .index(leaderIndex)
                                        .shard(renewRequest.getShardId().id())
                                        .primaryShard()
                                        .currentNodeId();
                                final String primaryShardNodeName =
                                    getLeaderCluster().clusterService().state().nodes().get(primaryShardNodeId).getName();
                                final IndexShard primary =
                                    getLeaderCluster()
                                        .getInstance(IndicesService.class, primaryShardNodeName)
                                        .getShardOrNull(renewRequest.getShardId());
                                final CountDownLatch innerLatch = new CountDownLatch(1);
                                try {
                                    // this forces the background renewal from following to face a retention lease not found exception
                                    logger.info("--> removing retention lease [{}] on the leader", retentionLeaseId);
                                    primary.removeRetentionLease(retentionLeaseId,
                                        ActionListener.wrap(r -> innerLatch.countDown(), e -> fail(e.toString())));
                                    logger.info("--> waiting for the removed retention lease [{}] to be synced on the leader",
                                        retentionLeaseId);
                                    innerLatch.await();
                                    logger.info("--> removed retention lease [{}] on the leader", retentionLeaseId);
                                } catch (final Exception e) {
                                    throw new AssertionError("failed to remove retention lease [" + retentionLeaseId + "] on the leader");
                                } finally {
                                    latch.countDown();
                                }
                            }
                        }
                        connection.sendRequest(requestId, action, request, options);
                    });
            }

            latch.await();

            assertRetentionLeaseRenewal(numberOfShards, numberOfReplicas, followerIndex, leaderIndex);
        } finally {
            for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
                final MockTransportService senderTransportService =
                    (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
                senderTransportService.clearAllRules();
            }
        }
    }

    /**
     * This test is fairly evil. This test is to ensure that we are protected against a race condition when unfollowing and a background
     * renewal fires. The action of unfollowing will remove retention leases from the leader. If a background renewal is firing at that
     * time, it means that we will be met with a retention lease not found exception. That will in turn trigger behavior to attempt to
     * re-add the retention lease, which means we are left in a situation where we have unfollowed, but the retention lease still remains
     * on the leader. However, we have a guard against this in the callback after the retention lease not found exception is thrown, which
     * checks if the shard follow node task is cancelled or completed.
     *
     * To test this this behavior is correct, we capture the call to renew the retention lease. Then, we will step in between and execute
     * an unfollow request. This will remove the retention lease on the leader. At this point, we can unlatch the renew call, which will
     * now be met with a retention lease not found exception. We will cheat and wait for that response to come back, and then synchronously
     * trigger the listener which will check to see if the shard follow node task is cancelled or completed, and if not, add the retention
     * lease back. After that listener returns, we can check to see if a retention lease exists on the leader.
     *
     * Note, this done mean that listener will fire twice, once in our onResponseReceived hook, and once after our onResponseReceived
     * callback returns. 🤷‍♀️
     *
     * @throws Exception if an exception occurs in the main test thread
     */
    public void testPeriodicRenewalDoesNotAddRetentionLeaseAfterUnfollow() throws Exception {
        final String leaderIndex = "leader";
        final String followerIndex = "follower";
        final int numberOfShards = 1;
        final int numberOfReplicas = 1;
        final Map<String, String> additionalIndexSettings = new HashMap<>();
        additionalIndexSettings.put(
                IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.getKey(),
                TimeValue.timeValueMillis(200).getStringRep());
        final String leaderIndexSettings = getIndexSettings(numberOfShards, numberOfReplicas, additionalIndexSettings);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex).setSource(leaderIndexSettings, XContentType.JSON).get());
        ensureLeaderYellow(leaderIndex);
        final PutFollowAction.Request followRequest = putFollow(leaderIndex, followerIndex);
        followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();

        ensureFollowerGreen(true, followerIndex);

        final CountDownLatch removeLeaseLatch = new CountDownLatch(1);
        final CountDownLatch unfollowLatch = new CountDownLatch(1);
        final CountDownLatch responseLatch = new CountDownLatch(1);

        final ClusterStateResponse followerClusterState = followerClient().admin().cluster().prepareState().clear().setNodes(true).get();

        try {
            for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
                final MockTransportService senderTransportService =
                        (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
                senderTransportService.addSendBehavior(
                        (connection, requestId, action, request, options) -> {
                            if (RetentionLeaseActions.Renew.ACTION_NAME.equals(action)
                                    || TransportActionProxy.getProxyAction(RetentionLeaseActions.Renew.ACTION_NAME).equals(action)) {
                                final String retentionLeaseId = getRetentionLeaseId(followerIndex, leaderIndex);
                                logger.info("--> blocking renewal request for retention lease [{}] until unfollowed", retentionLeaseId);
                                try {
                                    removeLeaseLatch.countDown();
                                    unfollowLatch.await();

                                    senderTransportService.addMessageListener(new TransportMessageListener() {

                                        @SuppressWarnings("rawtypes")
                                        @Override
                                        public void onResponseReceived(
                                                final long responseRequestId,
                                                final Transport.ResponseContext context) {
                                            if (requestId == responseRequestId) {
                                                final RetentionLeaseNotFoundException e =
                                                        new RetentionLeaseNotFoundException(retentionLeaseId);
                                                context.handler().handleException(new RemoteTransportException(e.getMessage(), e));
                                                responseLatch.countDown();
                                                senderTransportService.removeMessageListener(this);
                                            }
                                        }

                                    });

                                } catch (final InterruptedException e) {
                                    Thread.currentThread().interrupt();
                                    fail(e.toString());
                                }
                            }
                            connection.sendRequest(requestId, action, request, options);
                        });
            }

            removeLeaseLatch.await();

            pauseFollow(followerIndex);
            assertAcked(followerClient().admin().indices().close(new CloseIndexRequest(followerIndex)).actionGet());
            assertAcked(followerClient().execute(UnfollowAction.INSTANCE, new UnfollowAction.Request(followerIndex)).actionGet());

            unfollowLatch.countDown();

            responseLatch.await();

            final IndicesStatsResponse afterUnfollowStats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
            final List<ShardStats> afterUnfollowShardsStats = getShardsStats(afterUnfollowStats);
            for (final ShardStats shardStats : afterUnfollowShardsStats) {
                assertNotNull(shardStats.getRetentionLeaseStats());
                assertThat(Strings.toString(shardStats), RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                    shardStats.getRetentionLeaseStats().retentionLeases()).values(), empty());
            }
        } finally {
            for (final DiscoveryNode senderNode : followerClusterState.getState().nodes()) {
                final MockTransportService senderTransportService =
                    (MockTransportService) getFollowerCluster().getInstance(TransportService.class, senderNode.getName());
                senderTransportService.clearAllRules();
            }
        }
    }

    public void testForgetFollower() throws Exception {
        final String leaderIndex = "leader";
        final String followerIndex = "follower";
        final int numberOfShards = randomIntBetween(1, 4);
        final String leaderIndexSettings = getIndexSettings(numberOfShards, 0);
        assertAcked(leaderClient().admin().indices().prepareCreate(leaderIndex).setSource(leaderIndexSettings, XContentType.JSON).get());
        ensureLeaderYellow(leaderIndex);
        final PutFollowAction.Request followRequest = putFollow(leaderIndex, followerIndex);
        followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();

        ensureFollowerGreen(true, followerIndex);

        pauseFollow(followerIndex);
        followerClient().admin().indices().close(new CloseIndexRequest(followerIndex).masterNodeTimeout(TimeValue.MAX_VALUE)).actionGet();

        final ClusterStateResponse followerIndexClusterState =
                followerClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(followerIndex).get();
        final String followerUUID = followerIndexClusterState.getState().metadata().index(followerIndex).getIndexUUID();

        final BroadcastResponse forgetFollowerResponse = leaderClient().execute(
                ForgetFollowerAction.INSTANCE,
                new ForgetFollowerAction.Request(
                        getFollowerCluster().getClusterName(),
                        followerIndex,
                        followerUUID,
                        "leader_cluster",
                        leaderIndex)).actionGet();

        logger.info(Strings.toString(forgetFollowerResponse));
        assertThat(forgetFollowerResponse.getTotalShards(), equalTo(numberOfShards));
        assertThat(forgetFollowerResponse.getSuccessfulShards(), equalTo(numberOfShards));
        assertThat(forgetFollowerResponse.getFailedShards(), equalTo(0));
        assertThat(forgetFollowerResponse.getShardFailures(), emptyArray());

        final IndicesStatsResponse afterForgetFollowerStats =
                leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
        final List<ShardStats> afterForgetFollowerShardsStats = getShardsStats(afterForgetFollowerStats);
        for (final ShardStats shardStats : afterForgetFollowerShardsStats) {
            assertNotNull(shardStats.getRetentionLeaseStats());
            assertThat(Strings.toString(shardStats), RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                shardStats.getRetentionLeaseStats().retentionLeases()).values(), empty());
        }
    }

    private void assertRetentionLeaseRenewal(
            final int numberOfShards,
            final int numberOfReplicas,
            final String followerIndex,
            final String leaderIndex) throws Exception {
        // ensure that a retention lease has been put in place on each shard, and grab a copy of them
        final List<Map<String, RetentionLease>> retentionLeases = new ArrayList<>();
        assertBusy(() -> {
            retentionLeases.clear();
            final IndicesStatsResponse stats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
            assertNotNull(stats.getShards());
            assertThat(stats.getShards(), arrayWithSize(numberOfShards * (1 + numberOfReplicas)));
            final List<ShardStats> shardsStats = getShardsStats(stats);
            for (int i = 0; i < numberOfShards * (1 + numberOfReplicas); i++) {
                assertNotNull(shardsStats.get(i).getRetentionLeaseStats());
                final Map<String, RetentionLease> currentRetentionLeases = RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                        shardsStats.get(i).getRetentionLeaseStats().retentionLeases());
                assertThat(Strings.toString(shardsStats.get(i)), currentRetentionLeases.values(), hasSize(1));
                final RetentionLease retentionLease =
                        currentRetentionLeases.values().iterator().next();
                assertThat(retentionLease.id(), equalTo(getRetentionLeaseId(followerIndex, leaderIndex)));
                retentionLeases.add(currentRetentionLeases);
            }
        });

        // now ensure that the retention leases are being renewed
        assertBusy(() -> {
            final IndicesStatsResponse stats =
                    leaderClient().admin().indices().stats(new IndicesStatsRequest().clear().indices(leaderIndex)).actionGet();
            assertNotNull(stats.getShards());
            assertThat(stats.getShards(), arrayWithSize(numberOfShards * (1 + numberOfReplicas)));
            final List<ShardStats> shardsStats = getShardsStats(stats);
            for (int i = 0; i < numberOfShards * (1 + numberOfReplicas); i++) {
                assertNotNull(shardsStats.get(i).getRetentionLeaseStats());
                final Map<String, RetentionLease> currentRetentionLeases = RetentionLeaseUtils.toMapExcludingPeerRecoveryRetentionLeases(
                    shardsStats.get(i).getRetentionLeaseStats().retentionLeases());
                assertThat(Strings.toString(shardsStats.get(i)), currentRetentionLeases.values(), hasSize(1));
                final RetentionLease retentionLease =
                        currentRetentionLeases.values().iterator().next();
                assertThat(retentionLease.id(), equalTo(getRetentionLeaseId(followerIndex, leaderIndex)));
                // we assert that retention leases are being renewed by an increase in the timestamp
                assertThat(retentionLease.timestamp(), greaterThan(retentionLeases.get(i).values().iterator().next().timestamp()));
            }
        });
    }

    /**
     * Extract the shard stats from an indices stats response, with the stats ordered by shard ID with primaries first. This is to have a
     * consistent ordering when comparing two responses.
     *
     * @param stats the indices stats
     * @return the shard stats in sorted order with (shard ID, primary) as the sort key
     */
    private List<ShardStats> getShardsStats(final IndicesStatsResponse stats) {
        return Arrays.stream(stats.getShards())
                .sorted((s, t) -> {
                    if (s.getShardRouting().shardId().id() == t.getShardRouting().shardId().id()) {
                        return -Boolean.compare(s.getShardRouting().primary(), t.getShardRouting().primary());
                    } else {
                        return Integer.compare(s.getShardRouting().shardId().id(), t.getShardRouting().shardId().id());
                    }
                })
                .collect(Collectors.toList());
    }

    private String getRetentionLeaseId(final String followerIndex, final String leaderIndex) {
        final ClusterStateResponse followerIndexClusterState =
                followerClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(followerIndex).get();
        final String followerUUID = followerIndexClusterState.getState().metadata().index(followerIndex).getIndexUUID();

        final ClusterStateResponse leaderIndexClusterState =
                leaderClient().admin().cluster().prepareState().clear().setMetadata(true).setIndices(leaderIndex).get();
        final String leaderUUID = leaderIndexClusterState.getState().metadata().index(leaderIndex).getIndexUUID();

        return getRetentionLeaseId(followerIndex, followerUUID, leaderIndex, leaderUUID);
    }

    private String getRetentionLeaseId(String followerIndex, String followerUUID, String leaderIndex, String leaderUUID) {
        return retentionLeaseId(
                            getFollowerCluster().getClusterName(),
                            new Index(followerIndex, followerUUID),
                            getLeaderCluster().getClusterName(),
                            new Index(leaderIndex, leaderUUID));
    }

    private void assertExpectedDocument(final String followerIndex, final int value) {
        final GetResponse getResponse = followerClient().prepareGet(followerIndex, "doc", Integer.toString(value)).get();
        assertTrue("doc with id [" + value + "] is missing", getResponse.isExists());
        if (sourceEnabled) {
            assertTrue((getResponse.getSource().containsKey("f")));
            assertThat(getResponse.getSource().get("f"), equalTo(value));
        }
    }

}
