ReplicationOperation should fail gracefully

Problem:
finishAsFailed could be called asynchronously in the
middle of operations like runPostReplicationActions which try to
sync the translog. finishAsFailed immediately triggers the failure
of the resultListener which releases the index shard primary
operation permit. This means that runPostReplicationActions may
try to sync the translog without an operation permit.

Solution:
We refactor the infrastructure of ReplicationOperation regarding
pendingActions and the resultListener, by replacing them with a
RefCountingListener. This way, if there are async failures, they
are aggregated, and the result listener is called once, after all
mid-way operations are done.

For the specific error we got in issue elastic#97183, this means that
a call to onNoLongerPrimary (which can happen if we fail to fail
a replica shard or mark it as stale) will not immediately release
the primary operation permit and the assertion in the translog sync
will be honored.

Fixes elastic#97183

Author: kingherc

server/src/main/java/org/elasticsearch/action/support/replication/ReplicationOperation.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.UnavailableShardsException;
 import org.elasticsearch.action.support.ActiveShardCount;
+import org.elasticsearch.action.support.RefCountingListener;
 import org.elasticsearch.action.support.RetryableAction;
 import org.elasticsearch.action.support.TransportActions;
 import org.elasticsearch.cluster.action.shard.ShardStateAction;
@@ -56,16 +57,16 @@ public class ReplicationOperation<
     private final String opType;
     private final AtomicInteger totalShards = new AtomicInteger();
     /**
-     * The number of pending sub-operations in this operation. This is incremented when the following operations start and decremented when
-     * they complete:
+     * A {@link RefCountingListener} that encapsulates the pending sub-operations in this operation. The references are incremented when
+     * the following operations start and decremented when they complete:
      * <ul>
      * <li>The operation on the primary</li>
      * <li>The operation on each replica</li>
      * <li>Coordination of the operation as a whole. This prevents the operation from terminating early if we haven't started any replica
      * operations and the primary finishes.</li>
      * </ul>
      */
-    private final AtomicInteger pendingActions = new AtomicInteger();
+    private final RefCountingListener pendingActionsListener;
     private final AtomicInteger successfulShards = new AtomicInteger();
     private final Primary<Request, ReplicaRequest, PrimaryResultT> primary;
     private final Replicas<ReplicaRequest> replicasProxy;
@@ -74,9 +75,6 @@ public class ReplicationOperation<
     private final TimeValue retryTimeout;
     private final long primaryTerm;
 
-    // exposed for tests
-    private final ActionListener<PrimaryResultT> resultListener;
-
     private volatile PrimaryResultT primaryResult = null;
 
     private final List<ReplicationResponse.ShardInfo.Failure> shardReplicaFailures = Collections.synchronizedList(new ArrayList<>());
@@ -95,7 +93,16 @@ public ReplicationOperation(
     ) {
         this.replicasProxy = replicas;
         this.primary = primary;
-        this.resultListener = listener;
+        this.pendingActionsListener = new RefCountingListener(ActionListener.wrap((ignored) -> {
+            primaryResult.setShardInfo(
+                ReplicationResponse.ShardInfo.of(
+                    totalShards.get(),
+                    successfulShards.get(),
+                    shardReplicaFailures.toArray(ReplicationResponse.NO_FAILURES)
+                )
+            );
+            listener.onResponse(primaryResult);
+        }, listener::onFailure));
         this.logger = logger;
         this.threadPool = threadPool;
         this.request = request;
@@ -106,28 +113,37 @@ public ReplicationOperation(
     }
 
     public void execute() throws Exception {
-        final String activeShardCountFailure = checkActiveShardCount();
-        final ShardRouting primaryRouting = primary.routingEntry();
-        final ShardId primaryId = primaryRouting.shardId();
-        if (activeShardCountFailure != null) {
-            finishAsFailed(
-                new UnavailableShardsException(
-                    primaryId,
-                    "{} Timeout: [{}], request: [{}]",
-                    activeShardCountFailure,
-                    request.timeout(),
-                    request
+        try (pendingActionsListener) {
+            final String activeShardCountFailure = checkActiveShardCount();
+            final ShardRouting primaryRouting = primary.routingEntry();
+            final ShardId primaryId = primaryRouting.shardId();
+            if (activeShardCountFailure != null) {
+                pendingActionsListener.acquire()
+                    .onFailure(
+                        new UnavailableShardsException(
+                            primaryId,
+                            "{} Timeout: [{}], request: [{}]",
+                            activeShardCountFailure,
+                            request.timeout(),
+                            request
+                        )
+                    );
+                return;
+            }
+
+            totalShards.incrementAndGet();
+            var primaryCoordinationListener = pendingActionsListener.acquire(); // triggered when we finish all primary coordination
+            primary.perform(
+                request,
+                ActionListener.wrap(
+                    primaryResult -> { handlePrimaryResult(primaryResult, primaryCoordinationListener); },
+                    primaryCoordinationListener::onFailure
                 )
             );
-            return;
         }
-
-        totalShards.incrementAndGet();
-        pendingActions.incrementAndGet(); // increase by 1 until we finish all primary coordination
-        primary.perform(request, ActionListener.wrap(this::handlePrimaryResult, this::finishAsFailed));
     }
 
-    private void handlePrimaryResult(final PrimaryResultT primaryResult) {
+    private void handlePrimaryResult(final PrimaryResultT primaryResult, ActionListener<Void> primaryCoordinationListener) {
         this.primaryResult = primaryResult;
         final ReplicaRequest replicaRequest = primaryResult.replicaRequest();
         if (replicaRequest != null) {
@@ -136,11 +152,11 @@ private void handlePrimaryResult(final PrimaryResultT primaryResult) {
             }
             final ReplicationGroup replicationGroup = primary.getReplicationGroup();
 
-            pendingActions.incrementAndGet();
+            var primaryOperationCompleteListener = pendingActionsListener.acquire();
             replicasProxy.onPrimaryOperationComplete(
                 replicaRequest,
                 replicationGroup.getRoutingTable(),
-                ActionListener.wrap(ignored -> decPendingAndFinishIfNeeded(), exception -> {
+                ActionListener.wrap(ignored -> primaryOperationCompleteListener.onResponse(null), exception -> {
                     totalShards.incrementAndGet();
                     shardReplicaFailures.add(
                         new ReplicationResponse.ShardInfo.Failure(
@@ -151,7 +167,7 @@ private void handlePrimaryResult(final PrimaryResultT primaryResult) {
                             false
                         )
                     );
-                    decPendingAndFinishIfNeeded();
+                    primaryOperationCompleteListener.onResponse(null);
                 })
             );
 
@@ -181,7 +197,7 @@ public void onResponse(Void aVoid) {
                     primary.routingEntry(),
                     primary::localCheckpoint,
                     primary::globalCheckpoint,
-                    () -> decPendingAndFinishIfNeeded()
+                    () -> primaryCoordinationListener.onResponse(null)
                 );
             }
 
@@ -193,20 +209,25 @@ public void onFailure(Exception e) {
                 // We update the checkpoints since a refresh might fail but the operations could be safely persisted, in the case that the
                 // fsync failed the local checkpoint won't advance and the engine will be marked as failed when the next indexing operation
                 // is appended into the translog.
-                updateCheckPoints(primary.routingEntry(), primary::localCheckpoint, primary::globalCheckpoint, () -> finishAsFailed(e));
+                updateCheckPoints(
+                    primary.routingEntry(),
+                    primary::localCheckpoint,
+                    primary::globalCheckpoint,
+                    () -> primaryCoordinationListener.onFailure(e)
+                );
             }
         });
     }
 
     private void markUnavailableShardsAsStale(ReplicaRequest replicaRequest, ReplicationGroup replicationGroup) {
         // if inSyncAllocationIds contains allocation ids of shards that don't exist in RoutingTable, mark copies as stale
         for (String allocationId : replicationGroup.getUnavailableInSyncShards()) {
-            pendingActions.incrementAndGet();
+            var staleCopyListener = pendingActionsListener.acquire();
             replicasProxy.markShardCopyAsStaleIfNeeded(
                 replicaRequest.shardId(),
                 allocationId,
                 primaryTerm,
-                ActionListener.wrap(r -> decPendingAndFinishIfNeeded(), ReplicationOperation.this::onNoLongerPrimary)
+                ActionListener.wrap(r -> staleCopyListener.onResponse(null), e -> onNoLongerPrimary(e, staleCopyListener))
             );
         }
     }
@@ -243,12 +264,17 @@ private void performOnReplica(
             logger.trace("[{}] sending op [{}] to replica {} for request [{}]", shard.shardId(), opType, shard, replicaRequest);
         }
         totalShards.incrementAndGet();
-        pendingActions.incrementAndGet();
+        var replicationPendingActionListener = pendingActionsListener.acquire();
         final ActionListener<ReplicaResponse> replicationListener = new ActionListener<>() {
             @Override
             public void onResponse(ReplicaResponse response) {
                 successfulShards.incrementAndGet();
-                updateCheckPoints(shard, response::localCheckpoint, response::globalCheckpoint, () -> decPendingAndFinishIfNeeded());
+                updateCheckPoints(
+                    shard,
+                    response::localCheckpoint,
+                    response::globalCheckpoint,
+                    () -> replicationPendingActionListener.onResponse(null)
+                );
             }
 
             @Override
@@ -282,7 +308,10 @@ public void onFailure(Exception replicaException) {
                     primaryTerm,
                     message,
                     replicaException,
-                    ActionListener.wrap(r -> decPendingAndFinishIfNeeded(), ReplicationOperation.this::onNoLongerPrimary)
+                    ActionListener.wrap(
+                        r -> replicationPendingActionListener.onResponse(null),
+                        e -> onNoLongerPrimary(e, replicationPendingActionListener)
+                    )
                 );
             }
 
@@ -369,13 +398,13 @@ public void onAfter() {
         }
     }
 
-    private void onNoLongerPrimary(Exception failure) {
+    private void onNoLongerPrimary(Exception failure, ActionListener<Void> listener) {
         final Throwable cause = ExceptionsHelper.unwrapCause(failure);
         final boolean nodeIsClosing = cause instanceof NodeClosedException;
         if (nodeIsClosing) {
             // We prefer not to fail the primary to avoid unnecessary warning log
             // when the node with the primary shard is gracefully shutting down.
-            finishAsFailed(
+            listener.onFailure(
                 new RetryOnPrimaryException(
                     primary.routingEntry().shardId(),
                     String.format(
@@ -398,7 +427,7 @@ protected void doRun() {
                         primary.routingEntry()
                     );
                     primary.failShard(message, failure);
-                    finishAsFailed(new RetryOnPrimaryException(primary.routingEntry().shardId(), message, failure));
+                    listener.onFailure(new RetryOnPrimaryException(primary.routingEntry().shardId(), message, failure));
                 }
 
                 @Override
@@ -411,7 +440,7 @@ public void onFailure(Exception e) {
                     e.addSuppressed(failure);
                     assert false : e;
                     logger.error(() -> "unexpected failure while failing primary [" + primary.routingEntry() + "]", e);
-                    finishAsFailed(
+                    listener.onFailure(
                         new RetryOnPrimaryException(
                             primary.routingEntry().shardId(),
                             String.format(Locale.ROOT, "unexpected failure while failing primary [%s]", primary.routingEntry()),
@@ -461,32 +490,6 @@ protected String checkActiveShardCount() {
         }
     }
 
-    private void decPendingAndFinishIfNeeded() {
-        assert pendingActions.get() > 0 : "pending action count goes below 0 for request [" + request + "]";
-        if (pendingActions.decrementAndGet() == 0) {
-            finish();
-        }
-    }
-
-    private void finish() {
-        if (finished.compareAndSet(false, true)) {
-            primaryResult.setShardInfo(
-                ReplicationResponse.ShardInfo.of(
-                    totalShards.get(),
-                    successfulShards.get(),
-                    shardReplicaFailures.toArray(ReplicationResponse.NO_FAILURES)
-                )
-            );
-            resultListener.onResponse(primaryResult);
-        }
-    }
-
-    private void finishAsFailed(Exception exception) {
-        if (finished.compareAndSet(false, true)) {
-            resultListener.onFailure(exception);
-        }
-    }
-
     /**
      * An encapsulation of an operation that is to be performed on the primary shard
      */
@@ -693,7 +696,7 @@ public interface PrimaryResult<RequestT extends ReplicationRequest<RequestT>> {
 
         /**
          * Run actions to be triggered post replication
-         * @param listener calllback that is invoked after post replication actions have completed
+         * @param listener callback that is invoked after post replication actions have completed
          * */
         void runPostReplicationActions(ActionListener<Void> listener);
     }

server/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -3860,8 +3860,7 @@ public int getActiveOperationsCount() {
      * listener handles all exception cases internally.
      */
     public final void syncAfterWrite(Translog.Location location, Consumer<Exception> syncListener) {
-        // TODO AwaitsFix https://github.com/elastic/elasticsearch/issues/97183
-        // assert indexShardOperationPermits.getActiveOperationsCount() != 0;
+        assert indexShardOperationPermits.getActiveOperationsCount() != 0;
         verifyNotClosed();
         getEngine().asyncEnsureTranslogSynced(location, syncListener);
     }