Add possibility to acquire permits on primary shards with different checks (#119794)

Since #42241 we check that the shard must be in a primary mode for acquiring a
primary permit on it. We would like customize this check and an option to
perform different checks before running the `onPermitAcquired` listener.
For example, we would to skip the primary mode check when we acquire primary permits
during recovering of a hollow indexing shard.

See ES-10487

Author: arteam

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

@@ -189,6 +189,7 @@
 import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.index.seqno.RetentionLeaseActions.RETAIN_ALL;
 import static org.elasticsearch.index.seqno.SequenceNumbers.UNASSIGNED_SEQ_NO;
+import static org.elasticsearch.index.shard.IndexShard.PrimaryPermitCheck.CHECK_PRIMARY_MODE;
 
 public class IndexShard extends AbstractIndexShardComponent implements IndicesClusterStateService.Shard {
 
@@ -3568,58 +3569,100 @@ private EngineConfig newEngineConfig(LongSupplier globalCheckpointSupplier) {
         );
     }
 
+    /**
+     * Check to run before running the primary permit operation
+     */
+    public enum PrimaryPermitCheck {
+        CHECK_PRIMARY_MODE,
+        /**
+         * IMPORTANT: Currently intented to be used only for acquiring primary permits during the recovery of hollow shards.
+         * Don't disable primary mode checks unless you're really sure.
+         */
+        NONE
+    }
+
     /**
      * Acquire a primary operation permit whenever the shard is ready for indexing. If a permit is directly available, the provided
      * ActionListener will be called on the calling thread. During relocation hand-off, permit acquisition can be delayed. The provided
      * ActionListener will then be called using the provided executor.
-     *
      */
     public void acquirePrimaryOperationPermit(ActionListener<Releasable> onPermitAcquired, Executor executorOnDelay) {
-        acquirePrimaryOperationPermit(onPermitAcquired, executorOnDelay, false);
+        acquirePrimaryOperationPermit(onPermitAcquired, executorOnDelay, false, CHECK_PRIMARY_MODE);
     }
 
     public void acquirePrimaryOperationPermit(
         ActionListener<Releasable> onPermitAcquired,
         Executor executorOnDelay,
         boolean forceExecution
+    ) {
+        acquirePrimaryOperationPermit(onPermitAcquired, executorOnDelay, forceExecution, CHECK_PRIMARY_MODE);
+    }
+
+    public void acquirePrimaryOperationPermit(
+        ActionListener<Releasable> onPermitAcquired,
+        Executor executorOnDelay,
+        boolean forceExecution,
+        PrimaryPermitCheck primaryPermitCheck
     ) {
         verifyNotClosed();
         assert shardRouting.primary() : "acquirePrimaryOperationPermit should only be called on primary shard: " + shardRouting;
-        indexShardOperationPermits.acquire(wrapPrimaryOperationPermitListener(onPermitAcquired), executorOnDelay, forceExecution);
+        indexShardOperationPermits.acquire(
+            wrapPrimaryOperationPermitListener(primaryPermitCheck, onPermitAcquired),
+            executorOnDelay,
+            forceExecution
+        );
     }
 
     public boolean isPrimaryMode() {
         assert indexShardOperationPermits.getActiveOperationsCount() != 0 : "must hold permit to check primary mode";
         return replicationTracker.isPrimaryMode();
     }
 
+    public void acquireAllPrimaryOperationsPermits(final ActionListener<Releasable> onPermitAcquired, final TimeValue timeout) {
+        acquireAllPrimaryOperationsPermits(onPermitAcquired, timeout, CHECK_PRIMARY_MODE);
+    }
+
     /**
      * Acquire all primary operation permits. Once all permits are acquired, the provided ActionListener is called.
      * It is the responsibility of the caller to close the {@link Releasable}.
      */
-    public void acquireAllPrimaryOperationsPermits(final ActionListener<Releasable> onPermitAcquired, final TimeValue timeout) {
+    public void acquireAllPrimaryOperationsPermits(
+        final ActionListener<Releasable> onPermitAcquired,
+        final TimeValue timeout,
+        final PrimaryPermitCheck primaryPermitCheck
+    ) {
         verifyNotClosed();
         assert shardRouting.primary() : "acquireAllPrimaryOperationsPermits should only be called on primary shard: " + shardRouting;
 
-        asyncBlockOperations(wrapPrimaryOperationPermitListener(onPermitAcquired), timeout.duration(), timeout.timeUnit());
+        asyncBlockOperations(
+            wrapPrimaryOperationPermitListener(primaryPermitCheck, onPermitAcquired),
+            timeout.duration(),
+            timeout.timeUnit()
+        );
     }
 
     /**
-     * Wraps the action to run on a primary after acquiring permit. This wrapping is used to check if the shard is in primary mode before
-     * executing the action.
+     * Wraps the action to run on a primary after acquiring permit.
      *
+     * @param primaryPermitCheck check to run before the primary mode operation
      * @param listener the listener to wrap
      * @return the wrapped listener
      */
-    private ActionListener<Releasable> wrapPrimaryOperationPermitListener(final ActionListener<Releasable> listener) {
-        return listener.delegateFailure((l, r) -> {
-            if (isPrimaryMode()) {
-                l.onResponse(r);
-            } else {
-                r.close();
-                l.onFailure(new ShardNotInPrimaryModeException(shardId, state));
-            }
-        });
+    private ActionListener<Releasable> wrapPrimaryOperationPermitListener(
+        final PrimaryPermitCheck primaryPermitCheck,
+        final ActionListener<Releasable> listener
+    ) {
+        return switch (primaryPermitCheck) {
+            case CHECK_PRIMARY_MODE -> listener.delegateFailure((l, r) -> {
+                if (isPrimaryMode()) {
+                    l.onResponse(r);
+                } else {
+                    r.close();
+                    l.onFailure(new ShardNotInPrimaryModeException(shardId, state));
+                }
+            });
+            case NONE -> listener;
+        };
     }
 
     private void asyncBlockOperations(ActionListener<Releasable> onPermitAcquired, long timeout, TimeUnit timeUnit) {
@@ -3657,7 +3700,7 @@ public void runUnderPrimaryPermit(final Runnable runnable, final Consumer<Except
                 runnable.run();
             }
         }, onFailure);
-        acquirePrimaryOperationPermit(onPermitAcquired, executorOnDelay);
+        acquirePrimaryOperationPermit(onPermitAcquired, executorOnDelay, false, CHECK_PRIMARY_MODE);
     }
 
     private <E extends Exception> void bumpPrimaryTerm(

server/src/test/java/org/elasticsearch/index/shard/IndexShardTests.java

@@ -790,6 +790,21 @@ public void onFailure(final Exception e) {
                 }
             }, TimeValue.timeValueSeconds(30));
             latch.await();
+
+            // It's possible to acquire permits if we skip the primary mode check
+            var permitAcquiredLatch = new CountDownLatch(1);
+            indexShard.acquirePrimaryOperationPermit(ActionListener.wrap(r -> {
+                r.close();
+                permitAcquiredLatch.countDown();
+            }, Assert::assertNotNull), EsExecutors.DIRECT_EXECUTOR_SERVICE, false, IndexShard.PrimaryPermitCheck.NONE);
+            safeAwait(permitAcquiredLatch);
+
+            var allPermitsAcquiredLatch = new CountDownLatch(1);
+            indexShard.acquireAllPrimaryOperationsPermits(ActionListener.wrap(r -> {
+                r.close();
+                allPermitsAcquiredLatch.countDown();
+            }, Assert::assertNotNull), TimeValue.timeValueSeconds(30), IndexShard.PrimaryPermitCheck.NONE);
+            safeAwait(allPermitsAcquiredLatch);
         }
 
         if (Assertions.ENABLED && indexShard.routingEntry().isRelocationTarget() == false) {