package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.client.MeshClient;
import com.atlassian.bitbucket.mesh.client.ReplicateParameters;
import com.atlassian.bitbucket.mesh.config.ApplicationConfig;
import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.coordination.WriteLock;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenRegistry;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.management.SidebandManager;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.network.Partition;
import com.atlassian.bitbucket.mesh.network.PartitionReplica;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.util.SatiableConsumer;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static com.atlassian.bitbucket.mesh.MeshConstants.VERSION_DELETED;
import static com.atlassian.bitbucket.mesh.repository.RefHasher.NULL_HASH;
import static com.atlassian.bitbucket.mesh.transaction.PrepareResult.*;
import static com.atlassian.bitbucket.mesh.util.MoreCollectors.toImmutableList;
import static com.atlassian.bitbucket.mesh.util.MoreCollectors.toImmutableSet;
import static com.google.common.base.Preconditions.checkArgument;
import static java.util.Objects.requireNonNull;
import static java.util.Optional.of;

/**
 * {@link TransactionLeader} implementation where the local node is the transaction coordinator and responsible for
 * forwarding and coordinating {@link #onPrepared(Repository, RpcTransaction) prepared} transaction with remote
 * replicas.
 */
public class LocalTransactionLeader implements TransactionLeader, ReplicaTransactionCallback {

    private static final RepositoryVisitor<String> GET_PARTITION = new RepositoryVisitor<String>() {

        @Override
        public String visit(@Nonnull ManagedRepository repository) {
            return repository.getPartition();
        }
    };
    private static final Logger log = LoggerFactory.getLogger(LocalTransactionLeader.class);

    private final ApplicationConfig config;
    private final EventPublisher eventPublisher;
    private final CompletableFuture<PrepareResult> localPrepareResult;
    private final LogicalClock logicalClock;
    private final MeshClient meshClient;
    private final MeshManager meshManager;
    private final ConcurrentMap<MeshNode, ReplicaVersion> preparedVersions;
    private final QuorumMode quorumMode;
    private final Set<MeshNode> reportedInconsistentReplicas;
    private final RepositoryAttributeManager repositoryAttributeManager;
    private final RepositoryHashService repositoryHashService;
    private final SidebandManager sidebandManager;
    private final String transactionId;
    private final WriteTokenRegistry writeTokenRegistry;

    private volatile ReplicaVersion baselineVersion;
    private volatile Boolean committed;
    private volatile TeePackWriter packWriter;
    private volatile Set<Set<MeshNode>> quorumSets;
    private volatile List<ReplicaTransaction> replicaTransactions;
    private volatile Repository repository;
    private volatile boolean repositoryDelete;
    private volatile WriteLock writeLock;

    LocalTransactionLeader(@Nonnull String transactionId, @Nonnull ApplicationConfig config,
                           @Nonnull EventPublisher eventPublisher, @Nonnull LogicalClock logicalClock,
                           @Nonnull MeshClient meshClient, @Nonnull MeshManager meshManager,
                           @Nonnull RepositoryAttributeManager repositoryAttributeManager,
                           @Nonnull RepositoryHashService repositoryHashService,
                           @Nonnull SidebandManager sidebandManager,
                           @Nonnull WriteTokenRegistry writeTokenRegistry) {
        this.config = requireNonNull(config, "config");
        this.eventPublisher = requireNonNull(eventPublisher, "eventPublisher");
        this.logicalClock = requireNonNull(logicalClock, "logicalClock");
        this.meshClient = requireNonNull(meshClient, "meshClient");
        this.meshManager = requireNonNull(meshManager, "meshManager");
        this.repositoryAttributeManager = requireNonNull(repositoryAttributeManager, "repositoryAttributeManager");
        this.repositoryHashService = requireNonNull(repositoryHashService, "repositoryHashService");
        this.sidebandManager = requireNonNull(sidebandManager, "sidebandManager");
        this.transactionId = requireNonNull(transactionId, "transactionId");
        this.writeTokenRegistry = requireNonNull(writeTokenRegistry, "writeTokenRegistry");

        localPrepareResult = new CompletableFuture<>();
        preparedVersions = new ConcurrentHashMap<>();
        quorumMode = config.getReplication().getQuorumMode();
        reportedInconsistentReplicas = new HashSet<>();
    }

    @Nonnull
    @Override
    public MeshNode getNode() {
        return meshManager.getLocalNode();
    }

    @Nonnull
    @Override
    public Optional<SatiableConsumer<RpcStreamChunk>> newPackWriter(@Nonnull Repository repository) {
        if (replicaTransactions == null) {
            this.repository = repository;
            maybeStartReplicaTransactions();
            if (replicaTransactions == null) {
                return Optional.empty();
            }
        }

        if (packWriter == null || packWriter.isClosed()) {
            //If replicaTransactions isn't null, wrapUp assumes a baseline version is available. It's possible
            //for the transaction to fail after the packWriter is created but before onPrepared, so we need to
            //set an initial baselineVersion
            if (baselineVersion == null) {
                baselineVersion = getLocalVersion();
            }

            packWriter = new TeePackWriter(replicaTransactions.stream()
                    .collect(Collectors.toMap(
                            Function.identity(),
                            replicaTransaction -> replicaTransaction.newPackWriter(repository))));
            return of(packWriter);
        }
        throw new IllegalStateException("Cannot create a new pack writer because a pack writer is already open");
    }

    @Override
    public void onCommitted(Throwable thrown) {
        log.debug("[{}] Transaction {} has been committed", repository, transactionId, thrown);
        wrapUp(true, thrown);
    }

    @Nonnull
    @Override
    public Future<PrepareResult> onPrepared(@Nonnull Repository repository, @Nonnull RpcTransaction transaction) {
        if (this.repository != null && !this.repository.equals(requireNonNull(repository, "repository"))) {
            log.warn("[{}] A pack was staged for another repository ({}) in transaction {}!", repository,
                    this.repository, transactionId);
        }
        this.repository = repository;
        checkArgument(transactionId.equals(requireNonNull(transaction, "transaction").getId()),
                "onPrepared called for transaction %s on transaction leader for %s", transaction.getId(), transactionId);

        log.debug("[{}] Transaction {} ({}) is prepared", repository, transactionId, transaction.getDetailsOneofCase());

        //There's special handling for deletes; we don't check YES votes for inconsistencies for deletes and
        //use a special VERSION_DELETED as the resulting version to make them easily recognizable when repairing.
        repositoryDelete = transaction.hasDeleteRepository();
        if (baselineVersion == null) { //May be set in newPackWriter before onPrepared is called
            //Record the baseline version here to have a non-null version in wrapUp, even for non-managed repositories.
            //For managed repositories, the baselineVersion is determined (again) when holding the write lock in
            //discardVotesFromInconsistentReplicasAndMaybeCommit.
            baselineVersion = getLocalVersion();
        }
        preparedVersions.put(getNode(), baselineVersion);

        maybeStartReplicaTransactions();
        if (replicaTransactions == null) {
            return CompletableFuture.completedFuture(SUCCESS);
        }

        //With the local changes already prepared, and replica transactions started on all followers, start trying
        //to acquire the local write token and prepare to negotiate tokens on the followers
        writeLock = writeTokenRegistry.startLock(repository, transactionId, replicaTransactions.stream()
                .map(ReplicaTransaction::getNode)
                .collect(toImmutableSet()));

        //Use thenCompose to add an _optional_ stage that waits for the WriteLock to be acquired. This stage should
        //be skipped if the transaction fails to prepare. Note that this uses a CompletionStage to allow processing
        //to continue asynchronously while the leader is waiting for the lock to be acquired; if we block on a gRPC
        //thread we won't be able to receive RpcWriteTokenGranted messages, resulting in hung processing.
        CompletableFuture<Integer> yesCountResult = prepareReplicaTransactions(maybeSplitTransaction(transaction));
        CompletableFuture<Boolean> commitResult = yesCountResult.thenCompose(yesCount -> {
            log.debug("[{}] All replicas have voted on the transaction ({} prepared)", transactionId, yesCount);
            if (shouldCommit(false)) {
                return writeLock.asCompletionStage().thenApply(ignored -> {
                    log.debug("[{}] Write lock has been acquired", transactionId);
                    return discardVotesFromInconsistentReplicasAndMaybeCommit();
                });
            }
            log.debug("[{}] A quorum to commit could not be reached ({} of {} prepared successfully)",
                    transactionId, yesCount, replicaTransactions.size() + 1);
            return CompletableFuture.completedFuture(false);
        });

        //Register a consumer that takes the summed YES votes (after discarding YES votes from inconsistent replicas)
        //and invokes rollback on each ReplicaTransaction if any replica failed to prepare.
        commitResult.whenComplete((result, thrown) -> {
            if (result == null || !result) {
                if (thrown instanceof CancellationException) {
                    log.debug("[{}] Canceled while waiting for replicas to prepare the transaction", transactionId);
                } else if (thrown != null) {
                    log.debug("[{}] At least one replica failed to prepare the transaction", transactionId, thrown);
                }
                rollBack();
            }
        });

        localPrepareResult.whenComplete((ignored, thrown) -> {
            if (thrown != null) {
                //If the local transaction times out, it'll cancel the Future<PrepareResult> the transaction is
                //blocked on. When this happens, we need to abort and roll back all replica transactions if they
                //have not yet completed.
                commitResult.completeExceptionally(thrown);
            }
        });

        return localPrepareResult;
    }

    @Override
    public void onReplicaError(@Nonnull MeshNode node) {
        if (writeLock != null) {
            if (preparedVersions.remove(node) != null && !localPrepareResult.isDone()) {
                log.debug("[{}] Removed YES vote from {} because the replica reported an error", repository, node);
            }
            writeLock.remove(node);
        }
    }

    @Override
    public void onRolledBack(Throwable thrown) {
        log.debug("[{}] Transaction {} has been rolled back", repository, transactionId, thrown);
        wrapUp(false, thrown);
    }

    //ReplicaTransactionCallback method
    @Override
    public boolean onTokenGranted(@Nonnull MeshNode replica, @Nonnull ReplicaVersion replicaVersion) {
        //The replica granted us their write token. The replica will remain at the provided replica version
        //until we release their token, or we instruct them to commit. Record the replica's version so we can
        //compare them after all replicas that have voted YES have granted their tokens.
        preparedVersions.put(replica, replicaVersion);
        return writeLock.onTokenGranted(replica);
    }

    //ReplicaTransactionCallback method
    @Override
    public boolean yieldToken(@Nonnull MeshNode replica) {
        //The replica has received a higher priority write request and wants its write token back. Propagate
        //the request to the writeLock
        if (writeLock.yieldToken(replica)) {
            //The token has been yielded; remove the now stale replica version from the prepared versions.
            //We'll get a new replica version when the token is re-granted to us.
            preparedVersions.remove(replica);
            return true;
        }
        return false;
    }

    /**
     * Returns the {@link ReplicaVersion} with the highest {@code version}. If there are different
     * {@link ReplicaVersion}s at that same version, but with different content hashes, we select
     * the content hash that occurs most frequently.
     */
    private static ReplicaVersion getConsensusVersion(Collection<ReplicaVersion> replicaVersions) {
        long maxVersion = 0L;
        String bestContentHash = null;
        int bestContentHashCount = 0;
        Map<String, Integer> countByContentHash = new HashMap<>();
        for (ReplicaVersion replicaVersion : replicaVersions) {
            if (replicaVersion.getVersion() > maxVersion) {
                maxVersion = replicaVersion.getVersion();
                countByContentHash.clear();
                countByContentHash.put(replicaVersion.getContentHash(), 1);
                bestContentHash = replicaVersion.getContentHash();
                bestContentHashCount = 1;
            } else if (replicaVersion.getVersion() == maxVersion) {
                String contentHash = replicaVersion.getContentHash();
                int count = countByContentHash.compute(contentHash,
                        (ignored, current) -> current == null ? 1 : current + 1);
                if (contentHash.equals(bestContentHash)) {
                    bestContentHashCount = count;
                } else if (count > bestContentHashCount) {
                    bestContentHash = contentHash;
                    bestContentHashCount = count;
                }
            }
        }

        return bestContentHash == null ? null : new ReplicaVersion(bestContentHash, maxVersion);
    }

    private Set<MeshNode> awaitReplicas(boolean expectedResult) {
        Set<MeshNode> inconsistentReplicas = new HashSet<>();
        replicaTransactions.forEach(transaction -> {
            try {
                boolean replicaResult = transaction.await(config.getReplication().getFinalizeInterval());
                if (expectedResult == replicaResult) {
                    log.debug("[{}] {} as expected", transaction, expectedResult ? "Committed" : "Rolled back");
                    return;
                }
                //Any replica that is detected to be inconsistent while voting YES is instructed to roll back by
                //discardVotesFromInconsistentReplicasAndMaybeCommit. For these inconsistent replicas a roll back
                //is expected even if the transaction as a whole commits.
                if (!expectedResult || isPreparedConsistently(transaction.getNode())) {
                    log.warn("[{}] Unexpectedly {}", transaction, replicaResult ? "committed" : "rolled back");
                }
            } catch (Exception e) {
                if (transaction.getPrepareResult().isPresent()) {
                    //Don't log a warning if the replica failed to prepare (e.g. because it is offline);
                    //a warning will already have been logged in prepareTransactions.
                    log.warn("[{}] Failed to {}", transaction, expectedResult ? "commit" : "roll back", e);
                }
                if (!expectedResult) {
                    //Don't mark the replica as inconsistent if the replica failed to confirm that it rolled back.
                    //At worst, the replica will have some stale state but it will not have committed the transaction.
                    return;
                }
            }
            inconsistentReplicas.add(transaction.getNode());
        });

        return inconsistentReplicas;
    }

    private boolean discardVotesFromInconsistentReplicasAndMaybeCommit() {
        //Record the version just before deciding on the outcome of the transaction. This version is used
        //in wrapUp to decide if any replica (including the local replica) arrived at an unexpected
        //version.
        repositoryHashService.maybeRecalculateContentHash(repository);
        ReplicaVersion localVersion = getLocalVersion();

        //Determine whether enough of the replicas agree on the repository version and content hash
        //to commit. If inconsistent replicas are detected, their YES vote will be discarded and they
        //will be instructed to roll back.
        preparedVersions.put(getNode(), localVersion);
        baselineVersion = getPreparedVersion();

        Set<MeshNode> inconsistentNodes = new HashSet<>();

        //If the repository was deleted, we don't care whether all the nodes arrived at the same version.
        //For any other operation, though, we discard YES votes from any nodes at an unexpected version
        if (!repositoryDelete) {
            Iterator<Map.Entry<MeshNode, ReplicaVersion>> it = preparedVersions.entrySet().iterator();
            while (it.hasNext()) {
                Map.Entry<MeshNode, ReplicaVersion> entry = it.next();
                ReplicaVersion replicaVersion = entry.getValue();
                if (!replicaVersion.equals(baselineVersion)) {
                    MeshNode node = entry.getKey();
                    log.info("[{}] Discarding YES vote by {} because it is not consistent ({} instead of {})",
                            transactionId, node, replicaVersion, baselineVersion);
                    it.remove();
                    inconsistentNodes.add(node);
                }
            }
        }

        boolean leaderCanCommit = isPreparedConsistently(getNode());
        boolean commit = leaderCanCommit && shouldCommit(true);
        long observedVersion = baselineVersion.getVersion();
        if (commit) {
            //We're about to commit, so we'll end up at observedVersion + 1 so that's the version all
            //inconsistentNodes are going to be inconsistent with
            observedVersion++;

            //Any node that voted NO, or didn't arrive at the expected version, is about to become inconsistent
            replicaTransactions.stream()
                    .map(ReplicaTransaction::getNode)
                    .filter(node -> !isPreparedConsistently(node))
                    .forEach(inconsistentNodes::add);
        }
        //Inform the control plane before committing to allow it to route around inconsistent nodes
        //as soon as possible. It is also possible that replicas voted YES, were detected to be inconsistent and the
        //system discarded their YES vote. We need to report these nodes as inconsistent, even if we don't end up
        //committing.
        reportInconsistentReplicas(inconsistentNodes, observedVersion);

        if (commit) {
            long timestamp = logicalClock.tick();
            //`committed` is a tri-state, only when it's `false` this transaction has been rolled back, and due
            //to autoboxing we can't just use `== false`.
            boolean alreadyRolledBack = Boolean.FALSE.equals(committed);
            localPrepareResult.complete(
                    (repositoryDelete || localVersion.equals(baselineVersion) && !alreadyRolledBack) ?
                            SUCCESS : FAILURE);
            replicaTransactions.forEach(tx -> {
                //Even though the replica may have prepared correctly the transaction may have been rolled back on the
                //transaction leader. So instead of committing and then becoming inconsistent with the transaction leader
                //which will trigger a repair and roll back we will be rolling back the transaction on the follower
                //replicas
                if (isPreparedConsistently(tx.getNode()) && !alreadyRolledBack) {
                    tx.commit(timestamp);
                } else {
                    tx.rollBack(timestamp);
                }
            });
        }

        //If the leader can't commit, we always want to roll back the whole transaction because it'll result in an error
        //message (possibly from git) bubbling up to the client. That client will assume that the write failed and it
        //would be strange if under the covers the write succeeded. Return whether we are voting to commit.
        return commit;
    }

    private ReplicaVersion getLocalVersion() {
        return new ReplicaVersion(
                repositoryHashService.getContentHash(repository),
                repositoryAttributeManager.getVersion(repository.getId()));
    }

    private Partition getPartition(Repository repository) {
        String partitionId = repository.accept(GET_PARTITION);

        return partitionId == null ? null : meshManager.getPartition(partitionId);
    }

    /**
     * Selects the "consensus" replica version from the versions each of the replicas (including the local node)
     * have prepared. The consensus replica version is the highest version.
     * <p>
     * It shouldn't happen, but if there are nodes reported different content hashes for that highest version,
     * the most frequently reported content hash is returned.
     */
    private ReplicaVersion getPreparedVersion() {
        if (quorumMode == QuorumMode.ALWAYS) {
            //In ALWAYS mode, the transaction leader dictates what should be committed
            return preparedVersions.get(getNode());
        }

        return getConsensusVersion(preparedVersions.values());
    }

    private ReplicaVersion getSummaryVersion(ReplicaVersion localVersion) {
        if (quorumMode == QuorumMode.ALWAYS && localVersion != null) {
            //If the leader successfully committed or rolled back in ALWAYS mode, its version should be used.
            return localVersion;
        }
        Stream<ReplicaVersion> replicaVersions = replicaTransactions.stream()
                .flatMap(tx -> tx.getReplicaVersion()
                        .map(Stream::of)
                        .orElseGet(Stream::empty));
        if (localVersion != null) {
            replicaVersions = Stream.concat(Stream.of(localVersion), replicaVersions);
        }

        return getConsensusVersion(replicaVersions.collect(Collectors.toList()));
    }

    private boolean isPreparedConsistently(MeshNode node) {
        //Any replicas that prepared, but were detected to be at an inconsistent version/content-hash are told
        //to roll back and removed from preparedVersions (see discardVotesFromInconsistentReplicasAndMaybeCommit).
        //Only nodes that voted YES _and_ landed on the consensus ReplicaVersion have an entry in preparedVersions.
        return preparedVersions.containsKey(node);
    }

    private ReplicaVersion maybeIncrementLocalVersion(boolean committed, Throwable thrown) {
        if (repository == null || !committed || thrown != null) {
            //Transaction never prepared, the transaction rolled back, or the local commit failed. In all these cases,
            //The leader's local version should not be used to determine the transaction's summary version.
            return null;
        }
        //The local replica committed successfully; increment the version
        long localVersion = repositoryDelete ? VERSION_DELETED :
                repositoryAttributeManager.incrementAndGetVersion(repository.getId());
        long expectedVersion = repositoryDelete ? VERSION_DELETED : baselineVersion.getVersion() + 1;
        if (localVersion == expectedVersion) {
            log.debug("[{}] Committed transaction {}; t = {}, version = {}",
                    repository, transactionId, logicalClock.get(), localVersion);
        } else {
            log.error("[{}] Committed transaction {}; t = {}, version incremented to {} while expecting {}",
                    repository, transactionId, logicalClock.get(), localVersion, expectedVersion);
        }
        String contentHash = repositoryDelete ? NULL_HASH : repositoryHashService.getContentHash(repository);
        return new ReplicaVersion(contentHash, localVersion);
    }

    private Iterable<RpcTransaction> maybeSplitTransaction(RpcTransaction transaction) {
        long maxSize = config.getGrpc().getMaxFragmentSize();

        //We only know how to split up update-refs transactions (for now)
        if (!transaction.hasUpdateRefs() || transaction.getSerializedSize() < maxSize) {
            return ImmutableList.of(transaction);
        }

        //Calculate how large each list of ref updates can be for the fragment to not exceed maxSize
        List<RpcProposedRefUpdate> proposedUpdates = transaction.getUpdateRefs().getRefUpdatesList();
        long fragmentOverheadSize = transaction.getSerializedSize() -
                proposedUpdates.stream().mapToInt(RpcProposedRefUpdate::getSerializedSize).sum();
        maxSize -= fragmentOverheadSize;

        //Split the big transaction in smaller fragment, each of which should be under maxSize
        ImmutableList.Builder<RpcTransaction> fragments = ImmutableList.builder();
        RpcUpdateRefsRequest.Builder builder = transaction.getUpdateRefs().toBuilder().clearRefUpdates();
        long size = 0L;
        for (RpcProposedRefUpdate refUpdate : proposedUpdates) {
            int refUpdateSize = refUpdate.getSerializedSize();
            if (size + refUpdateSize > maxSize) {
                fragments.add(transaction.toBuilder()
                        .setMoreDetails(true)
                        .setUpdateRefs(builder)
                        .build());
                builder.clearRefUpdates();
                size = 0L;
            }
            builder.addRefUpdates(refUpdate);
            size += refUpdateSize;
        }
        fragments.add(transaction.toBuilder()
                .setUpdateRefs(builder)
                .build());
        List<RpcTransaction> transactions = fragments.build();
        log.debug("[{}] Split the update-refs transaction into {} fragments", repository, transactions.size());
        return transactions;
    }

    private void maybeStartReplicaTransactions() {
        if (replicaTransactions == null) {
            Partition partition = getPartition(repository);
            if (partition == null) {
                //The repository is not managed or (somehow) we don't know about its partition.
                log.debug("[{}] Not replicating {} transaction for non-managed repository",
                        repository, transactionId);
                return;
            }

            //Start the replica transactions and capture the quorum sets that need to be acquired for the write
            //to succeed.
            replicaTransactions = startReplicaTransactions(partition.getReplicas());
            quorumSets = new WriteQuorumHelper(meshManager).getQuorumSets(repository, replicaTransactions);
        }
    }

    private CompletableFuture<Integer> prepareReplicaTransactions(Iterable<RpcTransaction> txFragments) {
        //The local transaction has prepared, so count that as a YES vote (and start with 1).
        CompletableFuture<Integer> yesCountFuture = CompletableFuture.completedFuture(1);
        if (!replicaTransactions.isEmpty()) {
            log.debug("[{}] Preparing transaction on {} replica(s)", transactionId, replicaTransactions.size());
            long timestamp = writeLock.getLogicalTimestamp();
            for (ReplicaTransaction remoteTransaction : replicaTransactions) {
                //Convert each ReplicaTransaction's vote into a 0 or 1, then combine the results and sum up the YES
                //votes. When all replicas have completed, this produces a single CompletableFuture which gives the
                //total number of YES votes across all replicas.
                CompletionStage<Integer> replicaStage = remoteTransaction.prepare(txFragments, timestamp)
                        .handle((prepareResult, thrown) -> {
                            if (prepareResult != null && prepareResult.isSuccess()) {
                                return 1;
                            }
                            //The replica did not prepare successfully, so it's not going to write; remove it from
                            //as a participant in the writeLock
                            MeshNode replica = remoteTransaction.getNode();
                            writeLock.remove(replica);
                            if (thrown == null) {
                                log.debug("[{}] Removed {} from the write lock because it {}",
                                        transactionId, replica, prepareResult == ABSTAIN ? "abstained" : "voted NO");
                            } else {
                                log.warn("[{}] Removed {} from the write lock because it failed to prepare",
                                        transactionId, replica, thrown);
                            }
                            return 0;
                        });
                yesCountFuture = yesCountFuture.thenCombine(replicaStage, Integer::sum);
            }
        }

        return yesCountFuture;
    }

    private void reportInconsistentReplicas(Collection<MeshNode> nodes, long version) {
        RpcRepositoryReplicaObservation.Builder builder = RpcRepositoryReplicaObservation.newBuilder();
        //Only report inconsistent nodes that have not yet been reported
        nodes.stream()
                .filter(reportedInconsistentReplicas::add)
                .map(MeshNode::getId)
                .forEach(builder::addNodes);

        if (builder.getNodesCount() > 0) {
            sidebandManager.send(RpcSidebandFollowerFragment.newBuilder()
                    .setRepositoryReplicaObservation(builder.setRepository(repository.getId())
                            .setState(ReplicaState.INCONSISTENT.toRpc())
                            .setVersion(version))
                    .build());
        }
    }

    private void rollBack() {
        long rollbackTimestamp = logicalClock.tick();
        localPrepareResult.complete(FAILURE);
        replicaTransactions.forEach(tx -> tx.rollBack(rollbackTimestamp));
    }

    private boolean shouldCommit(boolean writeLockHeld) {
        MeshNode localNode = getNode();
        if (!isPreparedConsistently(localNode)) {
            //Transaction leader did not prepare successfully. As a result, git will likely report errors when we
            //roll back the local transaction. The client will believe the operation failed, so let's just roll back
            //the whole transaction, even if we could theoretically commit if we can still acquire a write quorum.
            return false;
        }

        for (Set<MeshNode> quorumSet : quorumSets) {
            int yesVotes = 0;
            int replicaCount = 0;
            if (quorumSet.contains(localNode)) {
                //The local transaction has prepared successfully and is part of the quorum set we're evaluating
                yesVotes = replicaCount = 1;
            }
            for (ReplicaTransaction tx : replicaTransactions) {
                if (quorumSet.contains(tx.getNode())) {
                    //Replicas which abstain are treated as if they were never part of the transaction. This allows
                    //a smaller quorum to potentially still form and commit the transaction
                    PrepareResult result = tx.getPrepareResult().orElse(FAILURE);
                    if (result.isBinding()) {
                        ++replicaCount;
                        //If we're not holding the write lock, just count the success responses and ignore the prepared
                        //versions. The prepared versions are only updated once the follower sends us their write token
                        //and this check would always fail if we don't hold the write token
                        if (result.isSuccess() && (!writeLockHeld || isPreparedConsistently(tx.getNode()))) {
                            //The replica's YES vote only counts if it has successfully prepared at the consensus
                            //replica version.
                            ++yesVotes;
                        }
                    }
                }
            }
            if (yesVotes < quorumMode.getRequiredVotes(replicaCount)) {
                if (log.isDebugEnabled()) {
                    log.debug("[{}] A quorum to commit could not be reached across replicas {} " +
                                    "(got {} out of {} required YES votes)",
                            transactionId, quorumSet.stream()
                                    .map(MeshNode::toString)
                                    .collect(Collectors.joining(", ", "[", "]")),
                            yesVotes, quorumMode.getRequiredVotes(replicaCount));
                }
                return false;
            }
        }

        return true;
    }

    private List<ReplicaTransaction> startReplicaTransactions(Collection<PartitionReplica> replicas) {
        ReplicateParameters parameters = new ReplicateParameters.Builder(meshManager.getLocalNode(), transactionId)
                .callback(this)
                .build();
        return replicas.stream()
                .map(PartitionReplica::getNode)
                .filter(MeshNode::isRemote)
                .map(node -> meshClient.replicate(node, parameters))
                .collect(toImmutableList());
    }

    private Stream<MeshNode> streamNodesNotAtSummaryVersion(ReplicaVersion summaryVersion,
                                                            ReplicaVersion leaderVersion) {
        Stream<MeshNode> inconsistentReplicas = replicaTransactions.stream()
                .filter(transaction -> transaction.getReplicaVersion()
                        .map(replicaVersion -> !summaryVersion.equals(replicaVersion))
                        .orElse(true)) //the replica rolled back or did not commit at all
                .map(ReplicaTransaction::getNode);

        if (leaderVersion == null || !leaderVersion.equals(summaryVersion)) {
            //local node is inconsistent
            return Stream.concat(Stream.of(getNode()), inconsistentReplicas);
        }
        return inconsistentReplicas;
    }

    private void wrapUp(boolean committed, Throwable thrown) {
        try {
            synchronized (this) {
                if (this.committed == null) {
                    this.committed = committed;
                } else {
                    log.warn("[{}] Transaction leader's wrap up attempted to modify whether the transaction should be" +
                                    " committed to {}. Ignoring the new value as it has already been set to {}.",
                            transactionId, committed, this.committed);
                }
            }

            ReplicaVersion localVersion = maybeIncrementLocalVersion(committed, thrown);

            if (replicaTransactions != null) {
                //If we never started the write lock, it means no prepare (or commit or roll back) message
                //was ever sent to the followers. We skip calling await, in that case, because the replica
                //transactions will never complete, and just move forward with sending out the summary
                Set<MeshNode> inconsistentNodes = writeLock == null ?
                        new HashSet<>() :
                        awaitReplicas(committed);
                MeshNode localNode = meshManager.getLocalNode();
                ReplicaVersion summaryVersion = null;
                RpcTransactionSummary.Builder builder = RpcTransactionSummary.newBuilder();
                if (committed) {
                    if (thrown != null) {
                        //If an exception was thrown locally, it means this node failed to commit, so it needs to
                        //mark itself as inconsistent in addition to any followers that didn't reach the state
                        inconsistentNodes.add(localNode);
                    }

                    //Only send the content hash and version if the transaction successfully committed; if the
                    //transaction rolled back, no WriteLock lock has been acquired and therefore it's possible
                    //that a concurrent write has been processed on some, but not all replica nodes. As a result,
                    //any follower that compared its local version and content hash with the versions on the summary
                    //could incorrectly conclude that it is inconsistent.
                    summaryVersion = getSummaryVersion(localVersion);
                    //Mark any node not at the summary version as inconsistent.
                    if (summaryVersion != null) {
                        streamNodesNotAtSummaryVersion(summaryVersion, localVersion)
                                .forEach(inconsistentNodes::add);

                        builder.setContentHash(summaryVersion.getContentHash())
                                .setVersion(summaryVersion.getVersion());
                    } else {
                        //This is all super unlikely... but if none of the nodes successfully committed, we could end
                        //up marking all replicas as inconsistent with the control plane with no way to repair.
                        //Instead, we'll treat this as a rollback and mark none of the replicas as inconsistent.
                        log.warn("[{}] All of the replicas failed to commit transaction {}. Treating it as a rollback.",
                                repository, transactionId);
                        committed = false;
                        inconsistentNodes.clear();
                    }
                }
                //Any nodes that have already been marked as inconsistent in the PREPARE phase are still inconsistent.
                inconsistentNodes.addAll(reportedInconsistentReplicas);
                if (committed && inconsistentNodes.isEmpty()) {
                    log.debug("[{}] Transaction {} committed cleanly.", repository, transactionId);
                } else {
                    log.info("[{}] Transaction {} {}. Inconsistent nodes: {}",
                            repository, transactionId, committed ? "committed" : "rolled back", inconsistentNodes);
                }
                long version = summaryVersion == null ? baselineVersion.getVersion() : summaryVersion.getVersion();
                reportInconsistentReplicas(inconsistentNodes, version);
                builder.addAllInconsistentReplicas(Iterables.transform(inconsistentNodes, MeshNode::getId))
                        .setCommitted(committed);
                long summaryTimestamp = logicalClock.tick();
                replicaTransactions.forEach(transaction -> {
                    try {
                        transaction.summarize(builder.build(), summaryTimestamp);
                    } catch (Exception e) {
                        log.warn("[{}] Failed to send summary", transaction, e);
                    }
                });

                if (inconsistentNodes.contains(localNode)) {
                    String repositoryId = repository.getId();
                    eventPublisher.publish(new ReplicaStateObservedEvent(repositoryId,
                            new ReplicaStateObservation(ReplicaState.INCONSISTENT, summaryTimestamp, version),
                            true));
                }
            }
        } finally {
            if (writeLock != null) {
                writeLock.close();
                writeLock = null;
            }
        }
    }
}
