+public class ThreePhaseCommitCohortProxy extends AbstractThreePhaseCommitCohort<ActorSelection> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ThreePhaseCommitCohortProxy.class);
+
+ private static final MessageSupplier COMMIT_MESSAGE_SUPPLIER = new MessageSupplier() {
+ @Override
+ public Object newMessage(TransactionIdentifier transactionId, short version) {
+ return new CommitTransaction(transactionId, version).toSerializable();
+ }
+
+ @Override
+ public boolean isSerializedReplyType(Object reply) {
+ return CommitTransactionReply.isSerializedType(reply);
+ }
+ };
+
+ private static final MessageSupplier ABORT_MESSAGE_SUPPLIER = new MessageSupplier() {
+ @Override
+ public Object newMessage(TransactionIdentifier transactionId, short version) {
+ return new AbortTransaction(transactionId, version).toSerializable();
+ }
+
+ @Override
+ public boolean isSerializedReplyType(Object reply) {
+ return AbortTransactionReply.isSerializedType(reply);
+ }
+ };
+
+ private final ActorContext actorContext;
+ private final List<CohortInfo> cohorts;
+ private final SettableFuture<Void> cohortsResolvedFuture = SettableFuture.create();
+ private final TransactionIdentifier transactionId;
+ private volatile OperationCallback commitOperationCallback;
+
+ public ThreePhaseCommitCohortProxy(ActorContext actorContext, List<CohortInfo> cohorts,
+ TransactionIdentifier transactionId) {
+ this.actorContext = actorContext;
+ this.cohorts = cohorts;
+ this.transactionId = Preconditions.checkNotNull(transactionId);
+
+ if (cohorts.isEmpty()) {
+ cohortsResolvedFuture.set(null);
+ }
+ }
+
+ private ListenableFuture<Void> resolveCohorts() {
+ if (cohortsResolvedFuture.isDone()) {
+ return cohortsResolvedFuture;
+ }
+
+ final AtomicInteger completed = new AtomicInteger(cohorts.size());
+ for (final CohortInfo info: cohorts) {
+ info.getActorFuture().onComplete(new OnComplete<ActorSelection>() {
+ @Override
+ public void onComplete(Throwable failure, ActorSelection actor) {
+ synchronized (completed) {
+ boolean done = completed.decrementAndGet() == 0;
+ if (failure != null) {
+ LOG.debug("Tx {}: a cohort Future failed", transactionId, failure);
+ cohortsResolvedFuture.setException(failure);
+ } else if (!cohortsResolvedFuture.isDone()) {
+ LOG.debug("Tx {}: cohort actor {} resolved", transactionId, actor);
+
+ info.setResolvedActor(actor);
+ if (done) {
+ LOG.debug("Tx {}: successfully resolved all cohort actors", transactionId);
+ cohortsResolvedFuture.set(null);
+ }
+ }
+ }
+ }
+ }, actorContext.getClientDispatcher());
+ }
+
+ return cohortsResolvedFuture;
+ }
+
+ @Override
+ public ListenableFuture<Boolean> canCommit() {
+ LOG.debug("Tx {} canCommit", transactionId);
+
+ final SettableFuture<Boolean> returnFuture = SettableFuture.create();
+
+ // The first phase of canCommit is to gather the list of cohort actor paths that will
+ // participate in the commit. buildCohortPathsList combines the cohort path Futures into
+ // one Future which we wait on asynchronously here. The cohort actor paths are
+ // extracted from ReadyTransactionReply messages by the Futures that were obtained earlier
+ // and passed to us from upstream processing. If any one fails then we'll fail canCommit.
+
+ Futures.addCallback(resolveCohorts(), new FutureCallback<Void>() {
+ @Override
+ public void onSuccess(Void notUsed) {
+ finishCanCommit(returnFuture);
+ }
+
+ @Override
+ public void onFailure(Throwable failure) {
+ returnFuture.setException(failure);
+ }
+ });
+
+ return returnFuture;
+ }
+
+ private void finishCanCommit(final SettableFuture<Boolean> returnFuture) {
+ LOG.debug("Tx {} finishCanCommit", transactionId);
+
+ // For empty transactions return immediately
+ if (cohorts.size() == 0) {
+ LOG.debug("Tx {}: canCommit returning result true", transactionId);
+ returnFuture.set(Boolean.TRUE);
+ return;
+ }
+
+ commitOperationCallback = new TransactionRateLimitingCallback(actorContext);
+ commitOperationCallback.run();
+
+ final Iterator<CohortInfo> iterator = cohorts.iterator();
+
+ final OnComplete<Object> onComplete = new OnComplete<Object>() {
+ @Override
+ public void onComplete(Throwable failure, Object response) {
+ if (failure != null) {
+ LOG.debug("Tx {}: a canCommit cohort Future failed", transactionId, failure);
+
+ returnFuture.setException(failure);
+ commitOperationCallback.failure();
+ return;
+ }
+
+ // Only the first call to pause takes effect - subsequent calls before resume are no-ops. So
+ // this means we'll only time the first transaction canCommit which should be fine.
+ commitOperationCallback.pause();
+
+ boolean result = true;
+ if (CanCommitTransactionReply.isSerializedType(response)) {
+ CanCommitTransactionReply reply = CanCommitTransactionReply.fromSerializable(response);
+
+ LOG.debug("Tx {}: received {}", transactionId, response);
+
+ if (!reply.getCanCommit()) {
+ result = false;
+ }
+ } else {
+ LOG.error("Unexpected response type {}", response.getClass());
+ returnFuture.setException(new IllegalArgumentException(
+ String.format("Unexpected response type %s", response.getClass())));
+ return;
+ }
+
+ if (iterator.hasNext() && result) {
+ sendCanCommitTransaction(iterator.next(), this);
+ } else {
+ LOG.debug("Tx {}: canCommit returning result: {}", transactionId, result);
+ returnFuture.set(Boolean.valueOf(result));
+ }
+
+ }
+ };
+
+ sendCanCommitTransaction(iterator.next(), onComplete);
+ }