onFailure() with TaskCancelledException when cancellation is detected

Author: JeremyDahlgren

server/src/main/java/org/elasticsearch/common/util/CancellableSingleObjectCache.java

@@ -100,6 +100,13 @@ protected boolean isFresh(Key currentKey, Key newKey) {
         return currentKey.equals(newKey);
     }
 
+    /**
+     * Sets the currently cached item reference to {@code null}, which will result in a {@code refresh()} on the next {@code get()} call.
+     */
+    protected void clearCurrentCachedItem() {
+        this.currentCachedItemRef.set(null);
+    }
+
     /**
      * Start a retrieval for the value associated with the given {@code input}, and pass it to the given {@code listener}.
      * <p>
@@ -110,7 +117,8 @@ protected boolean isFresh(Key currentKey, Key newKey) {
      *
      * @param input       The input to compute the desired value, converted to a {@link Key} to determine if the value that's currently
      *                    cached or pending is fresh enough.
-     * @param isCancelled Returns {@code true} if the listener no longer requires the value being computed.
+     * @param isCancelled Returns {@code true} if the listener no longer requires the value being computed. The listener is expected to be
+     *                    completed as soon as possible when cancellation is detected.
      * @param listener    The listener to notify when the desired value becomes available.
      */
     public final void get(Input input, BooleanSupplier isCancelled, ActionListener<Value> listener) {
@@ -230,11 +238,15 @@ boolean addListener(ActionListener<Value> listener, BooleanSupplier isCancelled)
                     ActionListener.completeWith(listener, future::actionResult);
                 } else {
                     // Refresh is still pending; it's not cancelled because there are still references.
-                    future.addListener(ContextPreservingActionListener.wrapPreservingContext(listener, threadContext));
+                    final var cancellableListener = ActionListener.notifyOnce(
+                        ContextPreservingActionListener.wrapPreservingContext(listener, threadContext)
+                    );
+                    future.addListener(cancellableListener);
                     final AtomicBoolean released = new AtomicBoolean();
                     cancellationChecks.add(() -> {
                         if (released.get() == false && isCancelled.getAsBoolean() && released.compareAndSet(false, true)) {
                             decRef();
+                            cancellableListener.onFailure(new TaskCancelledException("task cancelled"));
                         }
                     });
                 }

server/src/test/java/org/elasticsearch/common/util/CancellableSingleObjectCacheTests.java

@@ -101,8 +101,35 @@ public void testListenerCompletedByRefreshEvenIfDiscarded() throws ExecutionExce
         testCache.completeNextRefresh("foo", 1);
         assertThat(future2.result(), equalTo(1));
 
-        // ... and the original listener is also completed successfully
-        assertThat(future1.result(), sameInstance(future2.result()));
+        // We expect the first listener to have been completed with a cancellation exception when detected in the ensureNotCancelled() call.
+        assertTrue(future1.isDone());
+        expectThrows(ExecutionException.class, TaskCancelledException.class, future1::result);
+    }
+
+    public void testBothListenersReceiveTaskCancelledExceptionWhenBothSupersededAndNewTasksAreCancelled() {
+        final TestCache testCache = new TestCache();
+
+        // This computation is superseded and then cancelled.
+        final AtomicBoolean isCancelled = new AtomicBoolean();
+        final TestFuture future1 = new TestFuture();
+        testCache.get("foo", isCancelled::get, future1);
+        testCache.assertPendingRefreshes(1);
+
+        // A second get() call that supersedes the original refresh and starts another one, but will be cancelled as well.
+        final TestFuture future2 = new TestFuture();
+        testCache.get("bar", isCancelled::get, future2);
+        testCache.assertPendingRefreshes(2);
+
+        testCache.assertNextRefreshCancelled();
+        assertFalse(future1.isDone());
+        testCache.assertPendingRefreshes(1);
+        assertFalse(future2.isDone());
+
+        isCancelled.set(true);
+        // This next refresh should also fail with a cancellation exception.
+        testCache.completeNextRefresh("bar", 1);
+        expectThrows(ExecutionException.class, TaskCancelledException.class, future1::result);
+        expectThrows(ExecutionException.class, TaskCancelledException.class, future2::result);
     }
 
     public void testListenerCompletedWithCancellationExceptionIfRefreshCancelled() throws ExecutionException {