MINOR: Keep pendingTask as WakeupFuture if currentTask is completed already. (#21586)

System tests that use VerifiableConsumer are flaky because
VerifiableConsumer isn't shutting down on request in certain situations.
There can be a race condition in the commitSync method, as the future
that we set as the active task to the wakeupTrigger can be already
completed by the time we are setting it. Which leads to the wakeup
request never being fulfilled.  Added a check if the task we are
receiving in setActiveTask was triggered when we complete it
exceptionally.

Also added additional logging when a shutdown is requested to make
debugging easier.

Reviewers: Kirk True <ktrue@confluent.io>, Bill Bejeck
 <bbejeck@apache.org>
(cherry picked from commit fdece9c)

Author: Nikita-Shupletsov

clients/src/main/java/org/apache/kafka/clients/consumer/internals/WakeupTrigger.java

@@ -78,8 +78,12 @@ public <T> CompletableFuture<T> setActiveTask(final CompletableFuture<T> current
             if (task == null) {
                 return new ActiveFuture(currentTask);
             } else if (task instanceof WakeupFuture) {
-                currentTask.completeExceptionally(new WakeupException());
-                return null;
+                boolean wasTriggered = currentTask.completeExceptionally(new WakeupException());
+
+                // If the Future was *already* completed when we invoke completeExceptionally, the WakeupException
+                // will be ignored. If it was already completed, we then need to return a new WakeupFuture so that the
+                // next call to setActiveTask will throw the WakeupException.
+                return wasTriggered ? null : new WakeupFuture();
             } else if (task instanceof DisabledWakeups) {
                 return task;
             }

clients/src/test/java/org/apache/kafka/clients/consumer/internals/WakeupTriggerTest.java

@@ -219,6 +219,39 @@ public void testExceptionTriggeredWhenTaskAsynchronouslyCancelled() {
         assertThrows(WakeupException.class, () -> wakeupTrigger.maybeTriggerWakeup());
     }
 
+    @Test
+    public void testExceptionTriggeredWhenTaskAsynchronouslyCompletedBeforeSet() {
+        final CompletableFuture<Void> task = new CompletableFuture<>();
+        task.complete(null);
+        wakeupTrigger.wakeup();
+        wakeupTrigger.setActiveTask(task);
+        assertNotNull(wakeupTrigger.getPendingTask());
+        assertInstanceOf(WakeupTrigger.WakeupFuture.class, wakeupTrigger.getPendingTask());
+        assertThrows(WakeupException.class, () -> wakeupTrigger.maybeTriggerWakeup());
+    }
+
+    @Test
+    public void testExceptionTriggeredWhenTaskAsynchronouslyFailedBeforeSet() {
+        final CompletableFuture<Void> task = new CompletableFuture<>();
+        task.completeExceptionally(new RuntimeException("Simulated error"));
+        wakeupTrigger.wakeup();
+        wakeupTrigger.setActiveTask(task);
+        assertNotNull(wakeupTrigger.getPendingTask());
+        assertInstanceOf(WakeupTrigger.WakeupFuture.class, wakeupTrigger.getPendingTask());
+        assertThrows(WakeupException.class, () -> wakeupTrigger.maybeTriggerWakeup());
+    }
+
+    @Test
+    public void testExceptionTriggeredWhenTaskAsynchronouslyCancelledBeforeSet() {
+        final CompletableFuture<Void> task = new CompletableFuture<>();
+        task.cancel(true);
+        wakeupTrigger.wakeup();
+        wakeupTrigger.setActiveTask(task);
+        assertNotNull(wakeupTrigger.getPendingTask());
+        assertInstanceOf(WakeupTrigger.WakeupFuture.class, wakeupTrigger.getPendingTask());
+        assertThrows(WakeupException.class, () -> wakeupTrigger.maybeTriggerWakeup());
+    }
+
     private void assertWakeupExceptionIsThrown(final CompletableFuture<?> future) {
         assertTrue(future.isCompletedExceptionally());
         assertInstanceOf(WakeupException.class,

tests/kafkatest/services/verifiable_consumer.py

@@ -344,6 +344,8 @@ def _worker(self, idx, node):
                         handler.handle_partitions_revoked(event, node, self.logger)
                     elif name == "partitions_assigned":
                         handler.handle_partitions_assigned(event, node, self.logger)
+                    elif name == "shutdown_requested":
+                        self.logger.debug("Shutdown has been requested")
                     else:
                         self.logger.debug("%s: ignoring unknown event: %s" % (str(node.account), event))
 

tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java

@@ -259,6 +259,7 @@ public void run() {
     public void close() {
         boolean interrupted = false;
         try {
+            printJson(new ShutdownRequested());
             consumer.wakeup();
             while (true) {
                 try {
@@ -295,6 +296,14 @@ public String name() {
         }
     }
 
+    private static class ShutdownRequested extends ConsumerEvent {
+
+        @Override
+        public String name() {
+            return "shutdown_requested";
+        }
+    }
+
     private static class ShutdownComplete extends ConsumerEvent {
 
         @Override