/**
 * Creates a new recovery target object that represents a recovery to the provided shard.
 *
 * @param indexShard                        local shard where we want to recover to
 * @param sourceNode                        source node of the recovery where we recover from
 * @param listener                          called when recovery is completed/failed
 * @param ensureClusterStateVersionCallback callback to ensure that the current node is at least on a cluster state with the provided
 *                                          version; necessary for primary relocation so that new primary knows about all other ongoing
 *                                          replica recoveries when replicating documents (see {@link RecoverySourceHandler})
 */
public RecoveryTarget(final IndexShard indexShard,
               final DiscoveryNode sourceNode,
               final PeerRecoveryTargetService.RecoveryListener listener,
               final LongConsumer ensureClusterStateVersionCallback) {
    super("recovery_status");
    this.cancellableThreads = new CancellableThreads();
    this.recoveryId = ID_GENERATOR.incrementAndGet();
    this.listener = listener;
    this.logger = Loggers.getLogger(getClass(), indexShard.shardId());
    this.indexShard = indexShard;
    this.sourceNode = sourceNode;
    this.shardId = indexShard.shardId();
    this.tempFilePrefix = RECOVERY_PREFIX + UUIDs.randomBase64UUID() + ".";
    this.store = indexShard.store();
    this.ensureClusterStateVersionCallback = ensureClusterStateVersionCallback;
    // make sure the store is not released until we are done.
    store.incRef();
    indexShard.recoveryStats().incCurrentAsTarget();
}
 

@Test
public void testCancellationsDoesNotLeakPrimaryPermits() throws Exception {
    final CancellableThreads cancellableThreads = new CancellableThreads();
    final IndexShard shard = mock(IndexShard.class);
    final AtomicBoolean freed = new AtomicBoolean(true);
    when(shard.isRelocatedPrimary()).thenReturn(false);
    doAnswer(invocation -> {
        freed.set(false);
        ((ActionListener<Releasable>)invocation.getArguments()[0]).onResponse(() -> freed.set(true));
        return null;
    }).when(shard).acquirePrimaryOperationPermit(any(), anyString(), anyObject());

    Thread cancelingThread = new Thread(() -> cancellableThreads.cancel("test"));
    cancelingThread.start();
    try {
        RecoverySourceHandler.runUnderPrimaryPermit(() -> {}, "test", shard, cancellableThreads, logger);
    } catch (CancellableThreads.ExecutionCancelledException e) {
        // expected.
    }
    cancelingThread.join();
    // we have to use assert busy as we may be interrupted while acquiring the permit, if so we want to check
    // that the permit is released.
    assertBusy(() -> assertTrue(freed.get()));
}
 

private void validateMappingUpdate(final String indexName, final String type, Mapping update) {
    final CountDownLatch latch = new CountDownLatch(1);
    final AtomicReference<Throwable> error = new AtomicReference<>();
    mappingUpdatedAction.updateMappingOnMaster(indexName, type, update, waitForMappingUpdatePostRecovery, new MappingUpdatedAction.MappingUpdateListener() {
        @Override
        public void onMappingUpdate() {
            latch.countDown();
        }

        @Override
        public void onFailure(Throwable t) {
            latch.countDown();
            error.set(t);
        }
    });
    cancellableThreads.execute(new CancellableThreads.Interruptable() {
        @Override
        public void run() throws InterruptedException {
            try {
                if (latch.await(waitForMappingUpdatePostRecovery.millis(), TimeUnit.MILLISECONDS) == false) {
                    logger.debug("waited for mapping update on master for [{}], yet timed out", type);
                } else {
                    if (error.get() != null) {
                        throw new IndexShardRecoveryException(shardId, "Failed to propagate mappings on master post recovery", error.get());
                    }
                }
            } catch (InterruptedException e) {
                logger.debug("interrupted while waiting for mapping update");
                throw e;
            }
        }
    });
}
 

static void runUnderPrimaryPermit(CancellableThreads.Interruptable runnable, String reason,
                                  IndexShard primary, CancellableThreads cancellableThreads, Logger logger) {
    cancellableThreads.execute(() -> {
        CompletableFuture<Releasable> permit = new CompletableFuture<>();
        final ActionListener<Releasable> onAcquired = new ActionListener<>() {
            @Override
            public void onResponse(Releasable releasable) {
                if (permit.complete(releasable) == false) {
                    releasable.close();
                }
            }

            @Override
            public void onFailure(Exception e) {
                permit.completeExceptionally(e);
            }
        };
        primary.acquirePrimaryOperationPermit(onAcquired, ThreadPool.Names.SAME, reason);
        try (Releasable ignored = FutureUtils.get(permit)) {
            // check that the IndexShard still has the primary authority. This needs to be checked under operation permit to prevent
            // races, as IndexShard will switch its authority only when it holds all operation permits, see IndexShard.relocated()
            if (primary.isRelocatedPrimary()) {
                throw new IndexShardRelocatedException(primary.shardId());
            }
            runnable.run();
        } finally {
            // just in case we got an exception (likely interrupted) while waiting for the get
            permit.whenComplete((r, e) -> {
                if (r != null) {
                    r.close();
                }
                if (e != null) {
                    logger.trace("suppressing exception on completion (it was already bubbled up or the operation was aborted)", e);
                }
            });
        }
    });
}
 

/**
 * Closes the current recovery target and waits up to a certain timeout for resources to be freed.
 * Returns true if resetting the recovery was successful, false if the recovery target is already cancelled / failed or marked as done.
 */
boolean resetRecovery(CancellableThreads newTargetCancellableThreads) throws IOException {
    if (finished.compareAndSet(false, true)) {
        try {
            logger.debug("reset of recovery with shard {} and id [{}]", shardId, recoveryId);
        } finally {
            // release the initial reference. recovery files will be cleaned as soon as ref count goes to zero, potentially now.
            decRef();
        }
        try {
            newTargetCancellableThreads.execute(closedLatch::await);
        } catch (CancellableThreads.ExecutionCancelledException e) {
            logger.trace("new recovery target cancelled for shard {} while waiting on old recovery target with id [{}] to close",
                shardId, recoveryId);
            return false;
        }
        RecoveryState.Stage stage = indexShard.recoveryState().getStage();
        if (indexShard.recoveryState().getPrimary() && (stage == RecoveryState.Stage.FINALIZE || stage == RecoveryState.Stage.DONE)) {
            // once primary relocation has moved past the finalization step, the relocation source can put the target into primary mode
            // and start indexing as primary into the target shard (see TransportReplicationAction). Resetting the target shard in this
            // state could mean that indexing is halted until the recovery retry attempt is completed and could also destroy existing
            // documents indexed and acknowledged before the reset.
            assert stage != RecoveryState.Stage.DONE : "recovery should not have completed when it's being reset";
            throw new IllegalStateException("cannot reset recovery as previous attempt made it past finalization step");
        }
        indexShard.performRecoveryRestart();
        return true;
    }
    return false;
}
 

public CancellableThreads CancellableThreads() {
    return cancellableThreads;
}
 

public CancellableThreads cancellableThreads() {
    return cancellableThreads;
}