package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.AbstractStatusException;
import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.coordination.WriteToken;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenHandler;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenRegistry;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.git.builder.GitProcessBuilderFactory;
import com.atlassian.bitbucket.mesh.git.builder.indexpack.IndexPackStdioHandler;
import com.atlassian.bitbucket.mesh.git.exception.AbortedException;
import com.atlassian.bitbucket.mesh.git.exception.FailedPreconditionException;
import com.atlassian.bitbucket.mesh.git.exception.InvalidArgumentException;
import com.atlassian.bitbucket.mesh.git.exception.InvalidPackException;
import com.atlassian.bitbucket.mesh.gprc.GrpcSupport;
import com.atlassian.bitbucket.mesh.gprc.SynchronizedStreamObserver;
import com.atlassian.bitbucket.mesh.io.MoreFiles;
import com.atlassian.bitbucket.mesh.io.SyncFileTreeCallback;
import com.atlassian.bitbucket.mesh.io.SyncFileTreeParameters;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.transaction.*;
import com.atlassian.util.profiling.Ticker;
import com.atlassian.util.profiling.Timer;
import com.atlassian.util.profiling.Timers;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableSet;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.Closeable;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.nio.file.attribute.BasicFileAttributes;
import java.time.Duration;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

import static com.atlassian.bitbucket.mesh.MeshConstants.VERSION_DELETED;
import static com.atlassian.bitbucket.mesh.git.GitConstants.*;
import static com.atlassian.bitbucket.mesh.repository.RefHasher.NULL_HASH;
import static com.atlassian.bitbucket.mesh.rpc.v1.RpcPrepareTransactionResponse.Vote.*;
import static com.google.common.base.Preconditions.checkState;
import static java.lang.String.format;
import static java.util.Optional.empty;
import static java.util.Optional.of;
import static org.apache.commons.lang3.StringUtils.defaultString;

class RemoteTransactionLeader implements StreamObserver<RpcTransactLeaderFragment>, TransactionLeader {

    private static final Timer TIMER = Timers.timerWithMetric(RemoteTransactionLeader.class.getName() + ".prepare",
            "replication.Prepare");
    private static final String UNKNOWN_TX = "unknown";
    private static final Logger log = LoggerFactory.getLogger(RemoteTransactionLeader.class);

    private final GitProcessBuilderFactory builderFactory;
    private final AtomicBoolean completed;
    private final EventPublisher eventPublisher;
    private final ExecutionManager executionManager;
    private final ExecutorService executorService;
    private final RepositoryVisitor<Boolean> isMigrationTargetVisitor;
    private final LogicalClock logicalClock;
    private final MeshManager meshManager;
    private final CompletableFuture<PrepareResult> prepareResult;
    private final ReplicaStateRegistry replicaStateRegistry;
    private final RepositoryAttributeManager repositoryAttributeManager;
    private final RepositoryHashService repositoryHashService;
    private final RepositoryManager repositoryManager;
    private final StreamObserver<RpcTransactFollowerFragment> responseObserver;
    private final AtomicReference<TransactionState> state;
    private final RpcTransactionExecutor transactionExecutor;
    private final TransactionManager transactionManager;
    private final TokenHandler writeTokenHandler;
    private final WriteTokenRegistry writeTokenRegistry;

    private volatile IncomingPack incomingPack;
    private volatile boolean inconsistentAtStart;
    private volatile MeshNode leader;
    private volatile RpcPrepareTransactionResponse.Vote leaderVote;
    private volatile RpcPrepareTransactionResponse.Vote localVote;
    private volatile long logicalTimestamp;
    private volatile RpcTransaction operation;
    private volatile Future<?> replayFuture;
    private volatile Repository repository;
    private volatile long tokenTimestamp;
    private volatile String transactionId;
    private volatile WriteToken.Promise writeTokenPromise;

    RemoteTransactionLeader(StreamObserver<RpcTransactFollowerFragment> responseObserver,
                            GitProcessBuilderFactory builderFactory,
                            EventPublisher eventPublisher, ExecutionManager executionManager,
                            ExecutorService executorService, LogicalClock logicalClock,
                            MeshManager meshManager, ReplicaStateRegistry replicaStateRegistry,
                            RepositoryAttributeManager repositoryAttributeManager,
                            RepositoryHashService repositoryHashService, RepositoryManager repositoryManager,
                            RpcTransactionExecutor transactionExecutor, TransactionManager transactionManager,
                            WriteTokenRegistry writeTokenRegistry) {
        this.builderFactory = builderFactory;
        this.eventPublisher = eventPublisher;
        this.executionManager = executionManager;
        this.executorService = executorService;
        this.logicalClock = logicalClock;
        this.meshManager = meshManager;
        this.replicaStateRegistry = replicaStateRegistry;
        this.repositoryAttributeManager = repositoryAttributeManager;
        this.repositoryHashService = repositoryHashService;
        this.repositoryManager = repositoryManager;
        this.responseObserver = new SynchronizedStreamObserver<>(responseObserver);
        this.transactionExecutor = transactionExecutor;
        this.transactionManager = transactionManager;
        this.writeTokenRegistry = writeTokenRegistry;

        completed = new AtomicBoolean();
        prepareResult = new CompletableFuture<>();
        state = new AtomicReference<>(TransactionState.INITIAL);
        transactionId = UNKNOWN_TX; // a non-null value for logging purposes
        writeTokenHandler = new TokenHandler();
        isMigrationTargetVisitor = new RepositoryVisitor<Boolean>() {

            @Override
            public Boolean defaultValue(Repository repository) {
                return Boolean.FALSE;
            }

            @Override
            public Boolean visit(@Nonnull ManagedRepository repository) {
                MeshNode localNode = meshManager.getLocalNode();
                return meshManager.getPartitionMigrations(repository.getPartition()).stream()
                        .anyMatch(migration -> migration.getTargetNode().equals(localNode));
            }
        };
    }

    @Nonnull
    @Override
    public Optional<Path> getIncomingDir() {
        IncomingPack incomingPack = this.incomingPack;
        return incomingPack == null ? empty() : incomingPack.getIncomingDir();
    }

    @Nonnull
    @Override
    public MeshNode getNode() {
        return leader;
    }

    @Override
    public void onCommitted(Throwable throwable) {
        RpcCommitTransactionResponse.Builder builder = RpcCommitTransactionResponse.newBuilder();
        if (throwable == null) {
            long version = isDelete() ? VERSION_DELETED :
                    repositoryAttributeManager.incrementAndGetVersion(repository.getId());
            log.debug("[{}] Committed transaction {}; t = {}, version = {}",
                    repository, transactionId, logicalClock.get(), version);
            builder.setReplicaVersion(RpcReplicaVersion.newBuilder()
                    .setContentHash(isDelete() ? NULL_HASH : repositoryHashService.getContentHash(repository))
                    .setVersion(version));
            transitionTo(TransactionState.COMMITTED);
        } else {
            log.warn("[{}] Failed to commit {} transaction {}",
                    repository, operation.getDetailsOneofCase(), transactionId);
            builder.setCommitFailed(toCommitFailed(throwable));
            transitionTo(TransactionState.ROLLED_BACK);
        }

        //Inform the remote leader of the outcome of our attempt to commit the transaction
        sendFragmentToLeader(RpcTransactFollowerFragment.newBuilder()
                .setCommitResponse(builder)
                .setLogicalTimestamp(logicalClock.tick()));
    }

    @Override
    public void onCompleted() {
        log.debug("[{}] Leader for transaction {} completed", repository, transactionId);
    }

    @Override
    public void onError(Throwable t) {
        //If the error was sent by us, completed is already set to true. If the error was sent by the leader, we
        //still need to complete the RPC from our end
        maybeComplete();

        logicalTimestamp = logicalClock.get();
        TransactionState state = this.state.updateAndGet(current -> {
            if (current.canTransitionTo(TransactionState.ROLLING_BACK)) {
                return TransactionState.ROLLING_BACK;
            }
            return current;
        });
        if (!state.isFinished() && operation != null) {
            log.info("[{}] Marking replica as potentially inconsistent because the leader reported an error",
                    operation.getRepository());
            //Mark the repository at an unknown state for the current version because we cannot be sure that a next
            //version has been / is going to be committed.
            publishReplicaState(ReplicaState.UNKNOWN, repositoryAttributeManager.getVersion(repository.getId()));
        }
        if (state == TransactionState.ROLLING_BACK) {
            //We either transitioned to ROLLING_BACK, or that was already our state
            log.info("[{}] Rolling back after leader reported an error",
                    transactionId, log.isDebugEnabled() ? t : null);
        } else if (state.isFinished() && Status.fromThrowable(t).getCode() == Status.Code.CANCELLED) {
            //TODO: There's some sort of odd race/bug that is resulting in followers getting a Status.CANCELLED
            //      onError callback instead of onCompleted, even though the leader is calling onCompleted
            log.debug("[{}] Unexpected cancellation from leader after {}", transactionId, state);
        } else {
            //The transaction must already be beyond the PREPARED state, which means that we must have already
            //received instructions to commit or roll back
            log.warn("[{}] Remote leader ran into an error (state = {})", transactionId, state, t);
        }
        //Release the onPrepared callback if it's blocked
        prepareResult.complete(PrepareResult.FAILURE);
        //Release the write token if held
        maybeReleaseWriteToken();
        maybeRemoveIncomingPack();
    }

    @Override
    public void onNext(RpcTransactLeaderFragment fragment) {
        logicalTimestamp = logicalClock.onReceived(fragment.getLogicalTimestamp());
        switch (fragment.getFragmentOneofCase()) {
            case COMMIT_REQUEST:
                commit();
                break;
            case PACK_CHUNK:
                onIncomingPackChunk(fragment.getPackChunk());
                break;
            case PREPARE_REQUEST:
                //A prepare request is an implicit request for this node's write token, so record the timestamp
                tokenTimestamp = fragment.getLogicalTimestamp();
                prepare(fragment.getPrepareRequest());
                break;
            case ROLL_BACK_REQUEST:
                rollBack();
                break;
            case STAGE_PACK_REQUEST:
                startIncomingPack(fragment.getStagePackRequest());
                break;
            case SUMMARY:
                wrapUp(fragment.getSummary());
                break;
            case TOKEN_YIELDED:
                //Leader has yielded the write token as requested by us, propagate to the registry
                writeTokenHandler.onYielded();
                break;
            default:
                throw new InvalidArgumentException("Unsupported fragment type. Only COMMIT_REQUEST, PREPARE_REQUEST," +
                        " ROLL_BACK_REQUEST and SUMMARY are supported");
        }
    }

    /**
     * Invoked by the local transaction when the operation has successfully prepared
     *
     * @param repository  the repository the transaction modifies
     * @param transaction describes the operation
     * @return a {@link Future future} that completes when the remote leader informs this replica of the outcome
     *         of coordination
     */
    @Nonnull
    @Override
    public Future<PrepareResult> onPrepared(@Nonnull Repository repository, @Nonnull RpcTransaction transaction) {
        if (!isSameOperation(transaction)) {
            throw new IllegalArgumentException("Prepared operation does not match requested operation");
        }

        this.repository = repository;
        transitionTo(TransactionState.PREPARED);
        localVote = YES;

        writeTokenPromise = writeTokenRegistry.acquireToken(
                repository, transactionId, tokenTimestamp, writeTokenHandler);

        //Ensure the token is released and the promise canceled if the context is closed unexpectedly
        executionManager.getContext().ifPresent(context -> context.addCleanupCallback(writeTokenPromise::close));

        //Delegate sending of the prepare response to the writeTokenHandler so it can synchronize token grants and yield
        //requests internally. If the token is granted at the time the response is sent, the token is included in
        //the prepare response. If not, the leader is notified with a separate message when the token is granted.
        writeTokenHandler.sendPrepareResponse(localVote);

        //prepareResult will be completed when the transaction leader instructs us how to proceed
        return prepareResult;
    }

    @Override
    public void onRolledBack(Throwable throwable) {
        internalRollBack(throwable);
    }

    @VisibleForTesting
    static Path getIncomingDir(@Nonnull Repository repository, @Nonnull String transactionId) {
        return MoreFiles.resolve(repository.getPath(), PATH_OBJECTS, "incoming-" + MoreFiles.pathSafe(transactionId));
    }

    private static boolean isSameUpdateRefs(RpcUpdateRefsRequest first, RpcUpdateRefsRequest second) {
        if (first.getRefUpdatesCount() > 1 && first.getTransactional() != second.getTransactional()) {
            return false;
        }
        //Convert to a set because order does not matter (and has been seen to occasionally differ)
        Set<RpcProposedRefUpdate> firstUpdates = ImmutableSet.copyOf(first.getRefUpdatesList());
        Set<RpcProposedRefUpdate> secondUpdates = ImmutableSet.copyOf(second.getRefUpdatesList());
        return firstUpdates.equals(secondUpdates);
    }

    private static boolean shouldRepositoryExist(RpcTransaction transaction) {
        if (transaction.hasCreateRepository()) {
            //Fork operations are performed as a transaction on the origin repository, but have the fork
            //as the subject. So for repository creation, the repository should exist if it's used as the
            //origin in the RpcCreateRepositoryRequest.
            return transaction.getRepository().equals(transaction.getCreateRepository().getOrigin());
        }
        return true;
    }

    private static RpcTransactionCommitFailed toCommitFailed(Throwable exception) {
        Throwable cause = exception;
        RpcTransactionCommitFailed.Reason reason = RpcTransactionCommitFailed.Reason.UNKNOWN;
        while (reason == RpcTransactionCommitFailed.Reason.UNKNOWN &&
                cause != null && !cause.equals(cause.getCause())) {
            if (cause instanceof TransactionCommitFailedException) {
                return ((TransactionCommitFailedException) cause).toCommitFailed();
            }
            if (cause instanceof TimeoutException) {
                reason = RpcTransactionCommitFailed.Reason.TIMEOUT;
            } else if (cause instanceof InterruptedException || Thread.currentThread().isInterrupted()) {
                reason = RpcTransactionCommitFailed.Reason.INTERRUPTED;
            }
            cause = cause.getCause();
        }

        return RpcTransactionCommitFailed.newBuilder()
                .setReason(reason)
                .build();
    }

    private static RpcTransactionRollbackFailed toRollbackFailed(Throwable exception) {
        Throwable cause = exception;
        RpcTransactionRollbackFailed.Reason reason = RpcTransactionRollbackFailed.Reason.UNKNOWN;
        while (reason == RpcTransactionRollbackFailed.Reason.UNKNOWN &&
                cause != null && !cause.equals(cause.getCause())) {
            if (cause instanceof TransactionRollbackFailedException) {
                return ((TransactionRollbackFailedException) cause).toRollbackFailed();
            }
            if (cause instanceof TimeoutException) {
                reason = RpcTransactionRollbackFailed.Reason.TIMEOUT;
            } else if (cause instanceof InterruptedException || Thread.currentThread().isInterrupted()) {
                reason = RpcTransactionRollbackFailed.Reason.INTERRUPTED;
            }
            cause = cause.getCause();
        }

        return RpcTransactionRollbackFailed.newBuilder()
                .setReason(reason)
                .build();
    }

    private void awaitReplayComplete() {
        Future<?> replayFuture = this.replayFuture;
        if (replayFuture != null) {
            this.replayFuture = null;
            try {
                replayFuture.get(2, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } catch (ExecutionException | RuntimeException e) {
                log.trace("[{}] Replicated transaction failed", repository, e);
            } catch (TimeoutException e) {
                log.debug("[{}] Timed out waiting for replicated transaction", repository);
            }
        }
    }

    /**
     * Checks whether a transition from {@code current} to {@code next} is allowed. If so, {@code next} is
     * returned. If we're transitioning to {@code PREPARED} and we've received an early {@link #commit} or
     * {@link #rollBack roll back} instruction from the leader, {@code current} is returned to skip the
     * transition. Otherwise, an {@link IllegalTransactionStateException} is thrown.
     *
     * @param current the state to transition from
     * @param next    the state to transition to
     * @return {@code current} or {@code next}, depending on the current state
     * @throws IllegalTransactionStateException if the requested transition is not allowed
     */
    private TransactionState checkTransition(TransactionState current, TransactionState next) {
        if (current == next || current.canTransitionTo(next)) {
            return next;
        }
        if (next == TransactionState.PREPARED &&
                (current == TransactionState.COMMITTING || current == TransactionState.ROLLING_BACK)) {
            //Ignore transition to PREPARED if the leader has already notified us of the outcome
            return current;
        }

        throw new IllegalTransactionStateException(transactionId, current, next);
    }

    private void commit() {
        leaderVote = YES;
        transitionTo(TransactionState.COMMITTING);

        //If the remote leader sent us a pack that we've staged, publish it to the repository before allowing the
        //local transaction to commit.
        maybePublishIncomingPack();
        //Instruct the local operation to commit
        prepareResult.complete(PrepareResult.SUCCESS);
    }

    private void internalRollBack(Throwable t) {
        try {
            //First, perform the state update atomically, and return the _previous state_. This update needs
            //to be idempotent because it may run multiple times. We'll use the previous state to know what
            //processing the transition involves and do it after
            //This variable is named "state" intentionally to hide the `state` field
            TransactionState state = this.state.getAndUpdate(current -> {
                if (current.isFinished()) {
                    //If we've already reached an end state (COMMITTED or ROLLED_BACK), don't transition
                    return current;
                }

                //Otherwise, do a _checked_ transition to ROLLED_BACK
                return checkTransition(current, TransactionState.ROLLED_BACK);
            });

            //The transition is complete. Now we use the previous state to apply its processing
            switch (state) {
                case COMMITTING:
                    log.warn("[{}] Failed to commit transaction {}", repository, transactionId, t);
                    //Notify the transaction leader the local transaction failed to commit cleanly
                    sendFragmentToLeader(RpcTransactFollowerFragment.newBuilder()
                            .setCommitResponse(RpcCommitTransactionResponse.newBuilder()
                                    .setCommitFailed(toCommitFailed(t))));
                    break;
                case COMMITTED:
                    log.warn("[{}] Cannot roll back already COMMITTED transaction {}", repository, transactionId, t);
                    sendErrorToLeader(new IllegalTransactionStateException(transactionId,
                            TransactionState.COMMITTED, TransactionState.ROLLED_BACK));
                    break;
                case PREPARING:
                    if (t == null) {
                        log.debug("[{}] Failed to prepare transaction {}", repository, transactionId);
                    } else {
                        log.warn("[{}] Failed to prepare transaction {}", repository, transactionId, t);
                    }
                    //Notify the leader that we won't be able to apply the proposed operation
                    sendVote(NO);
                    break;
                case ROLLED_BACK:
                    //We've already notified the leader, nothing left to do
                    log.debug("[{}] Transaction {} has already been ROLLED_BACK", repository, transactionId, t);
                    break;
                case ROLLING_BACK:
                    RpcRollBackTransactionResponse.Builder builder = RpcRollBackTransactionResponse.newBuilder();
                    if (t != null) {
                        log.warn("[{}] Transaction {} failed to roll back cleanly", repository, transactionId, t);
                        builder.setRollbackFailed(toRollbackFailed(t));
                    }
                    //Notify the leader of our outcome trying to roll back the transaction
                    sendFragmentToLeader(RpcTransactFollowerFragment.newBuilder()
                            .setRollBackResponse(builder));
                    break;
                default:
                    //We're rolling back from INITIAL or PREPARED, which is unexpected.
                    //TODO: return a better error
                    String message = format("[%s] Unexpected rollback for transaction %s from state %s",
                            repository, transactionId, state);
                    log.warn(message, t);
                    sendErrorToLeader(new IllegalTransactionStateException(
                            transactionId, state, TransactionState.ROLLED_BACK));
            }
        } finally {
            //Release the transactional operation if it's latched in the onPrepared() callback
            prepareResult.complete(PrepareResult.FAILURE);
        }
    }

    private boolean isDelete() {
        return operation != null && operation.hasDeleteRepository();
    }

    private boolean isSameOperation(RpcTransaction transaction) {
        RpcTransaction.Builder minimalOperation = operation.toBuilder()
                .clearId()
                .clearPerson()
                .clearTimestamp();
        RpcTransaction.Builder minimalTransaction = transaction.toBuilder()
                .clearId()
                .clearPerson()
                .clearTimestamp();
        if (transaction.hasUpdateRefs()) {
            //ref-updates cannot be compared using a simple equals operations; the ref-update order is not significant
            //and when only one ref is updated, the transactional field also doesn't matter
            if (isSameUpdateRefs(transaction.getUpdateRefs(), operation.getUpdateRefs())) {
                minimalOperation.clearUpdateRefs();
                minimalTransaction.clearUpdateRefs();
            } else {
                return false;
            }
        }
        return minimalOperation.build().equals(minimalTransaction.build());
    }

    private void maybeComplete() {
        if (completed.compareAndSet(false, true)) {
            responseObserver.onCompleted();
        }
    }

    private void maybePublishIncomingPack() {
        IncomingPack incomingPack = this.incomingPack;
        if (incomingPack != null) {
            incomingPack.publish();
            incomingPack.close();
            this.incomingPack = null;
        }
    }

    private void maybeReleaseWriteToken() {
        if (writeTokenPromise != null) {
            writeTokenPromise.close();
            writeTokenPromise = null;
        }
    }

    private void maybeRemoveIncomingPack() {
        IncomingPack incomingPack = this.incomingPack;
        if (incomingPack != null) {
            incomingPack.close();
            this.incomingPack = null;
        }
    }

    private void maybeWaitForIncomingPack() {
        //maybeRemoveIncomingPack, which can be triggered from a few places, can null out incomingPack
        //at any time after this null check. We shadow the class field to get a stable reference
        IncomingPack incomingPack = this.incomingPack;
        if (incomingPack != null) {
            try {
                incomingPack.await();
            } catch (InvalidPackException e) {
                RpcInvalidPack.Reason reason = e.getReason();
                if (inconsistentAtStart && (reason == RpcInvalidPack.Reason.UNRESOLVED_DELTA ||
                        reason == RpcInvalidPack.Reason.MISSING_OBJECT)) {
                    log.info("[{}] Transaction leader {}@{} staged a pack that cannot be used because this " +
                            "replica is inconsistent and is missing some expected objects. " +
                                "Falling back to just-in-time fetch.", repository, transactionId, leader);
                } else {
                    log.warn("[{}] Transaction leader {}@{} staged a pack that cannot be used. Falling back to " +
                                    "just-in-time fetch {}", repository, transactionId, leader, e.getMessage());
                }
            } catch (RuntimeException e) {
                log.warn("[{}] index-pack failed to process staged pack for transaction {}",
                        repository, transactionId, e);
            }
        }
    }

    private void onIncomingPackChunk(RpcStreamChunk chunk) {
        if (incomingPack != null) {
            incomingPack.write(chunk);
        }
    }

    private void prepare(RpcPrepareTransactionRequest request) {
        if (localVote != null) {
            //We've already determined the outcome of this transaction (and sent it to the leader)
            //and do not need to process any subsequent request fragments.
            return;
        }
        if (operation == null) {
            leader = meshManager.getNodeById(request.getLeader());
            operation = request.getTransaction();
            transactionId = operation.getId();

            //Record whether the repository was known to be inconsistent locally. If so, it is somewhat expected that
            //we fail to replicate the transaction, and we can log accordingly. Note that we still try to replicate
            //the transaction because repair may already be running and may complete in time for this transaction to
            //go through.
            String repositoryId = operation.getRepository();
            inconsistentAtStart = replicaStateRegistry.getState(repositoryId) != ReplicaState.CONSISTENT;

            //In some scenarios, we should straight out abstain and not even attempt to replicate the transaction;
            if (isPerformingInitialSync(repositoryId)) {
                //Notify the leader we can't participate in this transaction
                transitionTo(TransactionState.ROLLED_BACK);
                sendVote(ABSTAIN);

                return;
            }

            transitionTo(TransactionState.PREPARING);
        } else {
            //This is a secondary prepare fragment; merge it with the existing operation
            RpcTransaction fragment = request.getTransaction();
            if (!operation.getMoreDetails()) {
                throw new FailedPreconditionException("No additional transaction fragments expected");
            }
            RpcTransaction.DetailsOneofCase fragmentCase = fragment.getDetailsOneofCase();
            if (fragmentCase != operation.getDetailsOneofCase()) {
                throw new FailedPreconditionException(format(
                        "Received an incompatible transaction fragment; expected %s but got %s",
                        operation.getDetailsOneofCase(), fragmentCase));
            }
            if (!operation.hasUpdateRefs()) {
                throw new FailedPreconditionException("Cannot merge " + fragmentCase + " transaction fragments");
            }
            operation = operation.toBuilder()
                    .setMoreDetails(fragment.getMoreDetails())
                    .setUpdateRefs(operation.getUpdateRefs().toBuilder()
                            .addAllRefUpdates(fragment.getUpdateRefs().getRefUpdatesList()))
                    .build();
            log.debug("[{}] Merged {} ref updates from prepare transaction fragment", repository,
                    fragment.getUpdateRefs().getRefUpdatesCount());
        }

        if (!operation.getMoreDetails()) {
            //Execute the provided operation on a separate thread to allow the gRPC processing thread to return, which
            //is necessary to be able to receive subsequent messages.
            replayFuture = executorService.submit(executionManager.maybeWrap(() -> {
                StartTransactionParameters parameters = new StartTransactionParameters.Builder(this)
                        .id(operation.getId())
                        .build();
                try (Transaction transaction = transactionManager.start(parameters);
                     Ticker ignored = TIMER.start(operation.getDetailsOneofCase())) {
                    //If the leader sent us a pack, wait for index-pack to finish before preparing the transaction
                    maybeWaitForIncomingPack();
                    transaction.rollBackIfThrows(() -> transactionExecutor.execute(operation));
                } catch (RuntimeException e) {
                    //An exception is _expected_ here if the leader instructed us to roll back; for instance a failed
                    //ref-update or repository creation will throw an exception.
                    if (leaderVote != NO) {
                        if (inconsistentAtStart) {
                            //This replica was inconsistent when the replicated transaction started, so it's not
                            //unexpected that we failed to replicate the write.
                            Throwable cause = Throwables.getRootCause(e);
                            log.info("[{}] Failed to replicate {} transaction {} because this replica was already " +
                                    "inconsistent (error = {}:{})", repository, operation.getDetailsOneofCase(),
                                    transactionId, cause.getClass().getSimpleName(),
                                    defaultString(cause.getMessage(), "-"));
                        } else {
                            log.warn("[{}] Could not replicate {} transaction {}: \\n{}",
                                    repository, operation.getDetailsOneofCase(), operation.getId(),
                                    StringUtils.abbreviate(operation.toString(), 2048), e);
                        }
                    }
                }
            }));
        }
    }

    private void publishReplicaState(ReplicaState state, long version) {
        eventPublisher.publish(new ReplicaStateObservedEvent(
                operation.getRepository(), new ReplicaStateObservation(state, logicalTimestamp, version)));
    }

    private void rollBack() {
        leaderVote = NO;

        //If we can, we transition to ROLLING_BACK. Otherwise, the state is left unchanged to allow for
        //cases where we've already rolled back before the leader votes (e.g. if preparation fails)
        state.updateAndGet(current ->
                current.canTransitionTo(TransactionState.ROLLING_BACK) ? TransactionState.ROLLING_BACK : current);

        //Instruct the local operation to roll back (if it's still preparing)
        prepareResult.complete(PrepareResult.FAILURE);
        //Release the write token if we hold it
        maybeReleaseWriteToken();
    }

    private void sendErrorToLeader(Throwable throwable) {
        if (completed.compareAndSet(false, true)) {
            if (throwable instanceof AbstractStatusException) {
                throwable = ((AbstractStatusException) throwable).toStatusException();
            }
            //sending an error closes the RPC
            responseObserver.onError(throwable);
        }
    }

    private void sendFragmentToLeader(RpcTransactFollowerFragment.Builder fragment) {
        if (!completed.get()) {
            responseObserver.onNext(fragment.setLogicalTimestamp(logicalClock.tick()).build());
        }
    }

    private void sendVote(RpcPrepareTransactionResponse.Vote vote) {
        localVote = vote;
        sendFragmentToLeader(RpcTransactFollowerFragment.newBuilder()
                .setPrepareResponse(RpcPrepareTransactionResponse.newBuilder()
                        .setVote(localVote)));
    }

    private boolean isPerformingInitialSync(String repositoryId) {
        if (inconsistentAtStart) {
            //If we do not have the repository locally, and the operation is not trying to create the repository,
            //we first need to check if we already know this repository is inconsistent. That implies we've been
            //notified by the control plane of the repository's existence, but we haven't yet been able to repair
            //our own copy. The most likely scenario here is we've recently been assigned a new partition replica
            //and haven't finished populating it yet.
            repository = repositoryManager.findById(repositoryId);
            if (repository == null) {
                return shouldRepositoryExist(operation);
            }
            //If this node's replica of the partition in question is the target of an ongoing partition migration,
            //and the repository is known to be inconsistent, the repository may already exist and be in the process
            //of being repaired. In this case, we can be confident that the control plane won't complete the migration
            //until all repositories have been migrated.
            return repository.accept(isMigrationTargetVisitor);
        }
        return false;
    }

    private void startIncomingPack(RpcStagePackRequest request) {
        if (repository == null && UNKNOWN_TX.equals(transactionId)) {
            repository = repositoryManager.findById(request.getRepository());
            transactionId = request.getTransactionId();
        } else {
            checkState(repository != null, "Cannot stage pack for CreateRepository transaction %s", transactionId);
            checkState(repository.getId().equals(request.getRepository()) &&
                            request.getTransactionId().equals(transactionId),
                    "Cannot stage pack for %s@%s; the current transaction is for %s@%s",
                    request.getRepository(), request.getTransactionId(), repository, transactionId);
        }
        if (repository == null) {
            log.info("[{}] Cannot stage pack for transaction {} because the repository does not exist",
                    request.getRepository(), transactionId);
            sendFragmentToLeader(RpcTransactFollowerFragment.newBuilder()
                    .setAbortStagingPack(RpcAbortStagingPack.getDefaultInstance()));
        } else {
            incomingPack = new IncomingPack(repository, transactionId);
            //Automatically close the incoming pack when the RPC terminates.
            executionManager.getContext().ifPresent(ctx -> ctx.addCleanupCallback(incomingPack::close));
        }
    }

    private void transitionTo(TransactionState nextState) {
        state.accumulateAndGet(nextState, this::checkTransition);
    }

    private void wrapUp(RpcTransactionSummary summary) {
        //Wait for the process that replays the transaction to complete. This ensures that we don't get 'fake'
        //reference-transaction hook callbacks after the RPC operation has already completed, which would trigger
        //warnings being logged
        awaitReplayComplete();
        maybeComplete();
        maybeReleaseWriteToken();

        if (summary.getCommitted()) {
            maybePublishIncomingPack();
        } else {
            maybeRemoveIncomingPack();
        }

        if (operation == null) {
            //If no operation is set, prepare was never called. This means the leader was trying to stage an
            //incoming pack (which will have been cleaned up above) and failed either before staging finished
            //or before it could prepare changes
            log.info("[{}] Transaction {} aborted before preparing any changes", repository, transactionId);
            //Return early because we shouldn't trust any inconsistent nodes the leader may report
            return;
        }

        List<String> inconsistentReplicas = summary.getInconsistentReplicasList();
        if (summary.getCommitted() && inconsistentReplicas.isEmpty()) {
            log.debug("[{}] Transaction {} {} committed cleanly",
                    repository, transactionId, operation.getDetailsOneofCase());
        } else {
            log.info("[{}] Transaction {} {} {}. Inconsistent nodes: {}",
                    repository, transactionId, operation.getDetailsOneofCase(),
                    summary.getCommitted() ? "committed" : "rolled back", inconsistentReplicas);
        }

        if (inconsistentReplicas.contains(meshManager.getLocalNode().getId())) {
            publishReplicaState(ReplicaState.INCONSISTENT, summary.getVersion());
        }
    }

    private class IncomingPack implements Closeable {

        private final AtomicBoolean aborted;
        private final Path incomingDir;
        private final Future<Path> indexPackFuture;
        private final IndexPackStdioHandler indexPackStdioHandler;
        private final AtomicLong size;

        IncomingPack(Repository repository, String transactionId) {
            aborted = new AtomicBoolean();
            incomingDir = MoreFiles.mkdir(RemoteTransactionLeader.getIncomingDir(repository, transactionId));
            indexPackStdioHandler = new IndexPackStdioHandler(repository.getPath().resolve(PATH_OBJECTS));
            indexPackFuture = builderFactory.builder(repository).indexPack()
                    .fixThin(true)
                    .stdin(true)
                    .withEnvironment(ENV_OBJECT_DIR, incomingDir.toString())
                    .withEnvironment(ENV_ALTERNATE_OBJECT_DIRS, repository.getPath().resolve(PATH_OBJECTS).toString())
                    .withTimeout(Duration.ofDays(7)) //timeout is always restricted to gRPC deadline
                    .build(indexPackStdioHandler)
                    .start();
            size = new AtomicLong();
        }

        public void abort() {
            if (aborted.compareAndSet(false, true)) {
                log.info("[{}] Aborting staging pack for transaction {}", repository, transactionId);
                sendFragmentToLeader(RpcTransactFollowerFragment.newBuilder()
                        .setAbortStagingPack(RpcAbortStagingPack.getDefaultInstance()));
                close();
            }
        }

        @Override
        public void close() {
            //Abort index-pack if still running
            indexPackFuture.cancel(false);
            //Clean up the (partial) pack file
            MoreFiles.deleteQuietly(incomingDir);
        }

        void await() {
            try {
                //Close stdin in case the closed flag was not sent
                indexPackStdioHandler.closeStdin();
                //Wait for up to the gRPC deadline for the pack to finish indexing
                indexPackFuture.get(GrpcSupport.getTimeout().toMillis(), TimeUnit.MILLISECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                indexPackFuture.cancel(false);
                throw new AbortedException("Interrupted while waiting for index-pack to finish");
            } catch (ExecutionException e) {
                Throwable cause = e.getCause();
                Throwables.throwIfUnchecked(cause);
                log.warn("[{}] git index-pack failed to process the staged pack sent by leader of transaction {}. " +
                        " Discarding partial pack file and continuing without it", repository, transactionId, e);
                MoreFiles.deleteQuietly(incomingDir);
            } catch (TimeoutException e) {
                indexPackFuture.cancel(false);
                throw new AbortedException("Timed out while waiting for index-pack to finish");
            }
        }

        Optional<Path> getIncomingDir() {
            //incomingDir will never be null but both await and publish can recursively delete it, so we need
            //to verify the directory still exists before we return it
            return Files.isDirectory(incomingDir) ? of(incomingDir) : empty();
        }

        void publish() {
            Path objectDir = repository.getPath().resolve(PATH_OBJECTS);
            try {
                MoreFiles.syncFileTree(new SyncFileTreeParameters.Builder(incomingDir, objectDir)
                        .allowHardLinks(true)
                        .detectMissing(false)
                        .replaceExisting(false)
                        .callback(new SyncFileTreeCallback() {

                            @Override
                            public boolean isSourceIncluded(@Nonnull Path source,
                                                            @Nonnull BasicFileAttributes sourceAttributes) {
                                //If we get a .keep file, don't copy it across
                                return !source.getFileName().endsWith(".keep");
                            }
                        })
                        .build());
                log.debug("[{}] Published {} byte staged pack file", repository, size.get());
                if (!MoreFiles.deleteQuietly(incomingDir)) {
                    log.debug("[{}] Failed to delete staged pack dir for transaction {}", repository, transactionId);
                }
            } catch (FileNotFoundException | NoSuchFileException e) {
                //There's nothing to publish
            } catch (IOException e) {
                throw new TransactionCommitFailedException("Failed to publish staged pack before committing",
                        RpcTransactionCommitFailed.Reason.PUBLISH_FAILED, e);
            }
        }

        void write(RpcStreamChunk chunk) {
            if (!aborted.get()) {
                try {
                    size.addAndGet(chunk.getData().size());
                    indexPackStdioHandler.writeStdin(chunk.getData(), chunk.getClosed());
                } catch (RuntimeException e) {
                    log.warn("[{}] git index-pack failed", repository, e);
                    abort();
                }
            }
        }
    }

    /**
     * This handler serves a bridge between the local WriteTokenRegistry and the remote transaction that requested
     * the write token (implicitly by sending us a prepare request).
     * <p>
     * A note on concurrency: the onGranted and yieldToken methods are called from the local WriteTokenRegistry,
     * which internally guarantees thread safety. In other words, TokenHandler does not need to perform internal
     * synchronization to prevent out-of-order processing.
     */
    private class TokenHandler implements WriteTokenHandler {

        private final Object lock;

        private boolean prepareResponseSent;
        private volatile boolean remoteLeaderHoldsToken;
        private WriteToken token;

        private TokenHandler() {
            lock = new Object();
        }

        @Override
        public void onGranted(@Nonnull WriteToken token) {
            synchronized (lock) {
                this.token = token;
                if (prepareResponseSent) {
                    //We've already sent our prepare response, which implies that the token had not yet been
                    //granted when we sent our initial prepare response. Follow up with an additional message.
                    log.debug("[{}] Granting {} token to tx leader {}", repository, transactionId, leader);
                    sendFragmentToLeader(RpcTransactFollowerFragment.newBuilder()
                            .setTokenGranted(createTokenGrantedMessage()));
                    remoteLeaderHoldsToken = true;
                }
            }
        }

        @Override
        public void yieldToken() {
            //Called by the local WriteTokenRegistry when the token should be yielded.
            boolean remoteLeaderHeldToken;
            synchronized (lock) {
                remoteLeaderHeldToken = remoteLeaderHoldsToken;
                if (remoteLeaderHeldToken) {
                    //Forward the request to the transaction leader.
                    log.debug("[{}] Requesting the {} token back from tx leader {}", repository, transactionId, leader);
                    sendFragmentToLeader(RpcTransactFollowerFragment.newBuilder()
                            .setYieldTokenRequest(RpcYieldWriteTokenRequest.getDefaultInstance())
                            .setLogicalTimestamp(logicalClock.tick()));
                }
            }
            if (!remoteLeaderHeldToken) {
                //We haven't sent the token grant to the leader yet, so we can immediately yield it
                onYielded();
            }
        }

        void onYielded() {
            WriteToken token = null;
            synchronized (lock) {
                //Called when the transaction leader has yielded the token.
                if (this.token != null) {
                    //Release the token back to the registry.
                    log.debug("[{}] Transaction leader {} has yielded the {} token", leader, repository, transactionId);
                    token = this.token;
                    this.token = null;
                }
            }
            if (token != null) {
                token.yieldToken();
                remoteLeaderHoldsToken = false;
            }
        }

        void sendPrepareResponse(RpcPrepareTransactionResponse.Vote localVote) {
            synchronized (lock) {
                //Let the leader know that the local transaction was successfully prepared
                RpcPrepareTransactionResponse.Builder builder = RpcPrepareTransactionResponse.newBuilder()
                        .setVote(localVote);
                if (token != null) {
                    builder.setWriteTokenGranted(createTokenGrantedMessage());
                }
                log.debug("[{}] Sending tx leader {} prepare response (vote = {}) {} {} token",
                        repository, leader, localVote, token == null ? "without" : "with", transactionId);
                sendFragmentToLeader(RpcTransactFollowerFragment.newBuilder().setPrepareResponse(builder));
                prepareResponseSent = true;
                remoteLeaderHoldsToken = true;
            }
        }

        private RpcWriteTokenGranted createTokenGrantedMessage() {
            repositoryHashService.maybeRecalculateContentHash(repository);
            return RpcWriteTokenGranted.newBuilder()
                    .setReplicaVersion(RpcReplicaVersion.newBuilder()
                            .setContentHash(repositoryHashService.getContentHash(repository))
                            .setVersion(repositoryAttributeManager.getVersion(repository.getId())))
                    .build();
        }
    }
}
