Fail primary term and generation listeners on a closed shard

If a shard has been closed, we should quickly bail out and fail
all waiting primary term and generation listeners. Otherwise,
the engine implementation may try to successfully to complete the
provided listeners and perform operations on an already closed
shard and cause some unexpected errors.

Author: arteam

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

@@ -4494,11 +4494,13 @@ public void waitForEngineOrClosedShard(ActionListener<Void> listener) {
      * Registers a listener for an event when the shard advances to the provided primary term and segment generation
      */
     public void waitForPrimaryTermAndGeneration(long primaryTerm, long segmentGeneration, ActionListener<Long> listener) {
-        waitForEngineOrClosedShard(
-            listener.delegateFailureAndWrap(
-                (l, ignored) -> getEngine().addPrimaryTermAndGenerationListener(primaryTerm, segmentGeneration, l)
-            )
-        );
+        waitForEngineOrClosedShard(listener.delegateFailureAndWrap((l, ignored) -> {
+            if (state == IndexShardState.CLOSED) {
+                l.onFailure(new IndexShardClosedException(shardId));
+            } else {
+                getEngine().addPrimaryTermAndGenerationListener(primaryTerm, segmentGeneration, l);
+            }
+        }));
     }
 
     /**

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

@@ -3334,6 +3334,22 @@ public void testWaitForClosedListener() throws IOException {
         assertThat("listener should have been called", called.get(), equalTo(true));
     }
 
+    public void testWaitForPrimaryTermAndGenerationFailsForClosedShard() throws IOException {
+        Settings settings = indexSettings(IndexVersion.current(), 1, 1).build();
+        IndexMetadata metadata = IndexMetadata.builder("test").putMapping("""
+            { "properties": { "foo":  { "type": "text"}}}""").settings(settings).primaryTerm(0, 1).build();
+        IndexShard primary = newShard(new ShardId(metadata.getIndex(), 0), true, "n1", metadata, null);
+
+        var exception = new AtomicReference<Exception>();
+        ActionListener<Long> listener = ActionListener.wrap(l -> { assert false : l; }, e -> exception.set(e));
+        primary.waitForPrimaryTermAndGeneration(0L, 0L, listener);
+
+        assertNull("waitForPrimaryTermAndGeneration should be waiting", exception.get());
+        closeShards(primary);
+        // Should bail out earlier without calling the engine
+        assertThat(exception.get(), instanceOf(IndexShardClosedException.class));
+    }
+
     public void testRecoverFromLocalShard() throws IOException {
         Settings settings = indexSettings(IndexVersion.current(), 1, 1).build();
         IndexMetadata metadata = IndexMetadata.builder("source")