From 69ce93ef11aa25c467f510dcf2073306619c2a9e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Francisco=20Fern=C3=A1ndez=20Casta=C3=B1o?=
 <francisco.fernandez.castano@gmail.com>
Date: Thu, 27 Feb 2025 11:52:14 +0100
Subject: [PATCH] Add ActionListener.releaseBefore

This commit introduces a new helper function, ActionListener.releaseBefore, which ensures that a resource is released before the listener is notified.
---
 .../elasticsearch/action/ActionListener.java  | 19 +++++++++++++++++
 .../action/ActionListenerTests.java           | 21 ++++++++++++++-----
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/action/ActionListener.java b/server/src/main/java/org/elasticsearch/action/ActionListener.java
index a158669d936fe..a53b084bb4bf1 100644
--- a/server/src/main/java/org/elasticsearch/action/ActionListener.java
+++ b/server/src/main/java/org/elasticsearch/action/ActionListener.java
@@ -335,6 +335,25 @@ static <Response> ActionListener<Response> runBefore(ActionListener<Response> de
         return assertOnce(new ActionListenerImplementations.RunBeforeActionListener<>(delegate, runBefore));
     }
 
+    /**
+     * Wraps a given listener and returns a new listener which releases the provided {@code releaseBefore}
+     * resource before the listener is notified via either {@code #onResponse} or {@code #onFailure}.
+     */
+    static <Response> ActionListener<Response> releaseBefore(ActionListener<Response> delegate, Releasable releaseBefore) {
+        Runnable releasableRunnable = runnableFromReleasable(releaseBefore);
+        return assertOnce(new ActionListenerImplementations.RunBeforeActionListener<>(delegate, new CheckedRunnable<>() {
+            @Override
+            public void run() {
+                releasableRunnable.run();
+            }
+
+            @Override
+            public String toString() {
+                return releasableRunnable.toString();
+            }
+        }));
+    }
+
     /**
      * Wraps a given listener and returns a new listener which makes sure {@link #onResponse(Object)}
      * and {@link #onFailure(Exception)} of the provided listener will be called at most once.
diff --git a/server/src/test/java/org/elasticsearch/action/ActionListenerTests.java b/server/src/test/java/org/elasticsearch/action/ActionListenerTests.java
index bb785b2b6093a..6ca1959b319bd 100644
--- a/server/src/test/java/org/elasticsearch/action/ActionListenerTests.java
+++ b/server/src/test/java/org/elasticsearch/action/ActionListenerTests.java
@@ -28,6 +28,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 
 import static org.hamcrest.Matchers.containsString;
@@ -610,16 +611,26 @@ public String toString() {
         );
     }
 
+    public void testReleaseBefore() {
+        runReleaseListenerTest(true, false, ActionListener::releaseBefore);
+        runReleaseListenerTest(true, true, ActionListener::releaseBefore);
+        runReleaseListenerTest(false, false, ActionListener::releaseBefore);
+    }
+
     public void testReleaseAfter() {
-        runReleaseAfterTest(true, false);
-        runReleaseAfterTest(true, true);
-        runReleaseAfterTest(false, false);
+        runReleaseListenerTest(true, false, ActionListener::releaseAfter);
+        runReleaseListenerTest(true, true, ActionListener::releaseAfter);
+        runReleaseListenerTest(false, false, ActionListener::releaseAfter);
     }
 
-    private static void runReleaseAfterTest(boolean successResponse, final boolean throwFromOnResponse) {
+    private static void runReleaseListenerTest(
+        boolean successResponse,
+        final boolean throwFromOnResponse,
+        BiFunction<ActionListener<Void>, Releasable, ActionListener<Void>> releaseListenerProvider
+    ) {
         final AtomicBoolean released = new AtomicBoolean();
         final String description = randomAlphaOfLength(10);
-        final ActionListener<Void> l = ActionListener.releaseAfter(new ActionListener<>() {
+        final ActionListener<Void> l = releaseListenerProvider.apply(new ActionListener<>() {
             @Override
             public void onResponse(Void unused) {
                 if (throwFromOnResponse) {