Introduce clean transition on primary promotion

This commit introduces a clean transition from the old primary term to
the new primary term when a replica is promoted primary. To accomplish
this, we delay all operations before incrementing the primary term. The
delay is guaranteed to be in place before we increment the term, and
then all operations that are delayed are executed after the delay is
removed which asynchronously happens on another thread. This thread does
not progress until in-flight operations that were executing are
completed, and after these operations drain, the delayed operations
re-acquire permits and are executed.

Relates #24925

Author: jasontedor

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

@@ -131,6 +131,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -332,12 +333,14 @@ public long getPrimaryTerm() {
     }
 
     /**
-     * notifies the shard of an increase in the primary term
+     * Notifies the shard of an increase in the primary term.
+     *
+     * @param newPrimaryTerm the new primary term
      */
-    public void updatePrimaryTerm(final long newTerm) {
+    public void updatePrimaryTerm(final long newPrimaryTerm) {
         assert shardRouting.primary() : "primary term can only be explicitly updated on a primary shard";
         synchronized (mutex) {
-            if (newTerm != primaryTerm) {
+            if (newPrimaryTerm != primaryTerm) {
                 // Note that due to cluster state batching an initializing primary shard term can failed and re-assigned
                 // in one state causing it's term to be incremented. Note that if both current shard state and new
                 // shard state are initializing, we could replace the current shard and reinitialize it. It is however
@@ -354,10 +357,22 @@ public void updatePrimaryTerm(final long newTerm) {
                         "a started primary shard should never update its term; "
                                 + "shard " + shardRouting + ", "
                                 + "current term [" + primaryTerm + "], "
-                                + "new term [" + newTerm + "]";
-                assert newTerm > primaryTerm :
-                        "primary terms can only go up; current term [" + primaryTerm + "], new term [" + newTerm + "]";
-                primaryTerm = newTerm;
+                                + "new term [" + newPrimaryTerm + "]";
+                assert newPrimaryTerm > primaryTerm :
+                        "primary terms can only go up; current term [" + primaryTerm + "], new term [" + newPrimaryTerm + "]";
+                /*
+                 * Before this call returns, we are guaranteed that all future operations are delayed and so this happens before we
+                 * increment the primary term. The latch is needed to ensure that we do not unblock operations before the primary term is
+                 * incremented.
+                 */
+                final CountDownLatch latch = new CountDownLatch(1);
+                indexShardOperationPermits.asyncBlockOperations(
+                        30,
+                        TimeUnit.MINUTES,
+                        latch::await,
+                        e -> failShard("exception during primary term transition", e));
+                primaryTerm = newPrimaryTerm;
+                latch.countDown();
             }
         }
     }

core/src/main/java/org/elasticsearch/index/shard/IndexShardOperationPermits.java

@@ -20,12 +20,14 @@
 package org.elasticsearch.index.shard;
 
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.Assertions;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.action.support.ThreadedActionListener;
 import org.elasticsearch.common.CheckedRunnable;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.lease.Releasable;
+import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.util.concurrent.ThreadContext.StoredContext;
 import org.elasticsearch.threadpool.ThreadPool;
 
@@ -36,20 +38,35 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
+/**
+ * Tracks shard operation permits. Each operation on the shard obtains a permit. When we need to block operations (e.g., to transition
+ * between terms) we immediately delay all operations to a queue, obtain all available permits, and wait for outstanding operations to drain
+ * and return their permits. Delayed operations will acquire permits and be completed after the operation that blocked all operations has
+ * completed.
+ */
 final class IndexShardOperationPermits implements Closeable {
+
     private final ShardId shardId;
     private final Logger logger;
     private final ThreadPool threadPool;
 
     private static final int TOTAL_PERMITS = Integer.MAX_VALUE;
-    // fair semaphore to ensure that blockOperations() does not starve under thread contention
-    final Semaphore semaphore = new Semaphore(TOTAL_PERMITS, true);
-    @Nullable private List<ActionListener<Releasable>> delayedOperations; // operations that are delayed
+    final Semaphore semaphore = new Semaphore(TOTAL_PERMITS, true); // fair to ensure a blocking thread is not starved
+    private final List<ActionListener<Releasable>> delayedOperations = new ArrayList<>(); // operations that are delayed
     private volatile boolean closed;
+    private boolean delayed; // does not need to be volatile as all accesses are done under a lock on this
 
-    IndexShardOperationPermits(ShardId shardId, Logger logger, ThreadPool threadPool) {
+    /**
+     * Construct operation permits for the specified shards.
+     *
+     * @param shardId    the shard
+     * @param logger     the logger for the shard
+     * @param threadPool the thread pool (used to execute delayed operations)
+     */
+    IndexShardOperationPermits(final ShardId shardId, final Logger logger, final ThreadPool threadPool) {
         this.shardId = shardId;
         this.logger = logger;
         this.threadPool = threadPool;
@@ -61,99 +78,170 @@ public void close() {
     }
 
     /**
-     * Wait for in-flight operations to finish and executes onBlocked under the guarantee that no new operations are started. Queues
-     * operations that are occurring in the meanwhile and runs them once onBlocked has executed.
+     * Wait for in-flight operations to finish and executes {@code onBlocked} under the guarantee that no new operations are started. Queues
+     * operations that are occurring in the meanwhile and runs them once {@code onBlocked} has executed.
      *
-     * @param timeout the maximum time to wait for the in-flight operations block
-     * @param timeUnit the time unit of the {@code timeout} argument
+     * @param timeout   the maximum time to wait for the in-flight operations block
+     * @param timeUnit  the time unit of the {@code timeout} argument
      * @param onBlocked the action to run once the block has been acquired
-     * @throws InterruptedException if calling thread is interrupted
-     * @throws TimeoutException if timed out waiting for in-flight operations to finish
+     * @param <E>       the type of checked exception thrown by {@code onBlocked}
+     * @throws InterruptedException      if calling thread is interrupted
+     * @throws TimeoutException          if timed out waiting for in-flight operations to finish
      * @throws IndexShardClosedException if operation permit has been closed
      */
-    public <E extends Exception> void blockOperations(long timeout, TimeUnit timeUnit, CheckedRunnable<E> onBlocked) throws
-        InterruptedException, TimeoutException, E {
+    <E extends Exception> void blockOperations(
+            final long timeout,
+            final TimeUnit timeUnit,
+            final CheckedRunnable<E> onBlocked) throws InterruptedException, TimeoutException, E {
         if (closed) {
             throw new IndexShardClosedException(shardId);
         }
+        delayOperations();
         try {
-            if (semaphore.tryAcquire(TOTAL_PERMITS, timeout, timeUnit)) {
-                assert semaphore.availablePermits() == 0;
-                try {
-                    onBlocked.run();
-                } finally {
-                    semaphore.release(TOTAL_PERMITS);
-                }
+            doBlockOperations(timeout, timeUnit, onBlocked);
+        } finally {
+            releaseDelayedOperations();
+        }
+    }
+
+    /**
+     * Immediately delays operations and on another thread waits for in-flight operations to finish and then executes {@code onBlocked}
+     * under the guarantee that no new operations are started. Delayed operations are run after {@code onBlocked} has executed. After
+     * operations are delayed and the blocking is forked to another thread, returns to the caller. If a failure occurs while blocking
+     * operations or executing {@code onBlocked} then the {@code onFailure} handler will be invoked.
+     *
+     * @param timeout   the maximum time to wait for the in-flight operations block
+     * @param timeUnit  the time unit of the {@code timeout} argument
+     * @param onBlocked the action to run once the block has been acquired
+     * @param onFailure the action to run if a failure occurs while blocking operations
+     * @param <E>       the type of checked exception thrown by {@code onBlocked} (not thrown on the calling thread)
+     */
+    <E extends Exception> void asyncBlockOperations(
+            final long timeout, final TimeUnit timeUnit, final CheckedRunnable<E> onBlocked, final Consumer<Exception> onFailure) {
+        delayOperations();
+        threadPool.executor(ThreadPool.Names.GENERIC).execute(new AbstractRunnable() {
+            @Override
+            public void onFailure(final Exception e) {
+                onFailure.accept(e);
+            }
+
+            @Override
+            protected void doRun() throws Exception {
+                doBlockOperations(timeout, timeUnit, onBlocked);
+            }
+
+            @Override
+            public void onAfter() {
+                releaseDelayedOperations();
+            }
+        });
+    }
+
+    private void delayOperations() {
+        synchronized (this) {
+            if (delayed) {
+                throw new IllegalStateException("operations are already delayed");
             } else {
-                throw new TimeoutException("timed out during blockOperations");
+                assert delayedOperations.isEmpty();
+                delayed = true;
             }
-        } finally {
-            final List<ActionListener<Releasable>> queuedActions;
+        }
+    }
+
+    private <E extends Exception> void doBlockOperations(
+            final long timeout,
+            final TimeUnit timeUnit,
+            final CheckedRunnable<E> onBlocked) throws InterruptedException, TimeoutException, E {
+        if (Assertions.ENABLED) {
+            // since delayed is not volatile, we have to synchronize even here for visibility
             synchronized (this) {
-                queuedActions = delayedOperations;
-                delayedOperations = null;
+                assert delayed;
             }
-            if (queuedActions != null) {
-                // Try acquiring permits on fresh thread (for two reasons):
-                // - blockOperations can be called on recovery thread which can be expected to be interrupted when recovery is cancelled.
-                //   Interruptions are bad here as permit acquisition will throw an InterruptedException which will be swallowed by
-                //   ThreadedActionListener if the queue of the thread pool on which it submits is full.
-                // - if permit is acquired and queue of the thread pool which the ThreadedActionListener uses is full, the onFailure
-                //   handler is executed on the calling thread. This should not be the recovery thread as it would delay the recovery.
-                threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> {
-                    for (ActionListener<Releasable> queuedAction : queuedActions) {
-                        acquire(queuedAction, null, false);
-                    }
-                });
+        }
+        if (semaphore.tryAcquire(TOTAL_PERMITS, timeout, timeUnit)) {
+            assert semaphore.availablePermits() == 0;
+            try {
+                onBlocked.run();
+            } finally {
+                semaphore.release(TOTAL_PERMITS);
             }
+        } else {
+            throw new TimeoutException("timeout while blocking operations");
+        }
+    }
+
+    private void releaseDelayedOperations() {
+        final List<ActionListener<Releasable>> queuedActions;
+        synchronized (this) {
+            assert delayed;
+            queuedActions = new ArrayList<>(delayedOperations);
+            delayedOperations.clear();
+            delayed = false;
+        }
+        if (!queuedActions.isEmpty()) {
+            /*
+             * Try acquiring permits on fresh thread (for two reasons):
+             *   - blockOperations can be called on a recovery thread which can be expected to be interrupted when recovery is cancelled;
+             *     interruptions are bad here as permit acquisition will throw an interrupted exception which will be swallowed by
+             *     the threaded action listener if the queue of the thread pool on which it submits is full
+             *   - if a permit is acquired and the queue of the thread pool which the the threaded action listener uses is full, the
+             *     onFailure handler is executed on the calling thread; this should not be the recovery thread as it would delay the
+             *     recovery
+             */
+            threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> {
+                for (ActionListener<Releasable> queuedAction : queuedActions) {
+                    acquire(queuedAction, null, false);
+                }
+            });
         }
     }
 
     /**
      * Acquires a permit whenever permit acquisition is not blocked. If the permit is directly available, the provided
      * {@link ActionListener} will be called on the calling thread. During calls of
-     * {@link #blockOperations(long, TimeUnit, CheckedRunnable)}, permit acquisition can be delayed. The provided ActionListener will
-     * then be called using the provided executor once operations are no longer blocked.
+     * {@link #blockOperations(long, TimeUnit, CheckedRunnable)}, permit acquisition can be delayed. The provided {@link ActionListener}
+     * will then be called using the provided executor once operations are no longer blocked.
      *
      * @param onAcquired      {@link ActionListener} that is invoked once acquisition is successful or failed
      * @param executorOnDelay executor to use for delayed call
      * @param forceExecution  whether the runnable should force its execution in case it gets rejected
      */
-    public void acquire(ActionListener<Releasable> onAcquired, String executorOnDelay, boolean forceExecution) {
+    public void acquire(final ActionListener<Releasable> onAcquired, final String executorOnDelay, final boolean forceExecution) {
         if (closed) {
             onAcquired.onFailure(new IndexShardClosedException(shardId));
             return;
         }
-        Releasable releasable;
+        final Releasable releasable;
         try {
             synchronized (this) {
-                releasable = tryAcquire();
-                if (releasable == null) {
-                    // blockOperations is executing, this operation will be retried by blockOperations once it finishes
-                    if (delayedOperations == null) {
-                        delayedOperations = new ArrayList<>();
-                    }
+                if (delayed) {
                     final Supplier<StoredContext> contextSupplier = threadPool.getThreadContext().newRestorableContext(false);
                     if (executorOnDelay != null) {
                         delayedOperations.add(
-                            new ThreadedActionListener<>(logger, threadPool, executorOnDelay,
-                                new ContextPreservingActionListener<>(contextSupplier, onAcquired), forceExecution));
+                                new ThreadedActionListener<>(logger, threadPool, executorOnDelay,
+                                        new ContextPreservingActionListener<>(contextSupplier, onAcquired), forceExecution));
                     } else {
                         delayedOperations.add(new ContextPreservingActionListener<>(contextSupplier, onAcquired));
                     }
                     return;
+                } else {
+                    releasable = tryAcquire();
+                    assert releasable != null;
                 }
             }
-        } catch (InterruptedException e) {
+        } catch (final InterruptedException e) {
             onAcquired.onFailure(e);
             return;
         }
+        // execute this outside the synchronized block!
         onAcquired.onResponse(releasable);
     }
 
-    @Nullable private Releasable tryAcquire() throws InterruptedException {
-        if (semaphore.tryAcquire(1, 0, TimeUnit.SECONDS)) { // the untimed tryAcquire methods do not honor the fairness setting
-            AtomicBoolean closed = new AtomicBoolean();
+    @Nullable
+    private Releasable tryAcquire() throws InterruptedException {
+        assert Thread.holdsLock(this);
+        if (semaphore.tryAcquire(1, 0, TimeUnit.SECONDS)) { // the un-timed tryAcquire methods do not honor the fairness setting
+            final AtomicBoolean closed = new AtomicBoolean();
             return () -> {
                 if (closed.compareAndSet(false, true)) {
                     semaphore.release(1);
@@ -163,13 +251,23 @@ public void acquire(ActionListener<Releasable> onAcquired, String executorOnDela
         return null;
     }
 
-    public int getActiveOperationsCount() {
+    /**
+     * Obtain the active operation count, or zero if all permits are held (even if there are outstanding operations in flight).
+     *
+     * @return the active operation count, or zero when all permits ar eheld
+     */
+    int getActiveOperationsCount() {
         int availablePermits = semaphore.availablePermits();
         if (availablePermits == 0) {
-            // when blockOperations is holding all permits
+            /*
+             * This occurs when either doBlockOperations is holding all the permits or there are outstanding operations in flight and the
+             * remainder of the permits are held by doBlockOperations. We do not distinguish between these two cases and simply say that
+             * the active operations count is zero.
+             */
             return 0;
         } else {
             return TOTAL_PERMITS - availablePermits;
         }
     }
+
 }