KAFKA-20860: Retry consumer reconciliation after synchronous failure

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

@@ -944,36 +944,43 @@ public void maybeReconcile(boolean canCommit) {
                 revokedPartitions
         );
 
-        // Mark partitions as pending revocation to stop fetching from the partitions (no new
-        // fetches sent out, and no in-flight fetches responses processed).
-        markPendingRevocationToPauseFetching(revokedPartitions);
-
-        // Commit offsets if auto-commit enabled before reconciling a new assignment. Request will
-        // be retried until it succeeds, fails with non-retriable error, or timer expires.
-        CompletableFuture<Void> commitResult = signalReconciliationStarted();
-
-        // Execute commit -> onPartitionsRevoked -> onPartitionsAssigned.
-        commitResult.whenComplete((__, commitReqError) -> {
-            if (commitReqError != null) {
-                // The call to commit, that includes retry logic for retriable errors, failed to
-                // complete within the time boundaries (fatal error or retriable that did not
-                // recover). Proceed with the revocation.
-                log.error("Auto-commit request before reconciling new assignment failed. " +
-                    "Will proceed with the reconciliation anyway.", commitReqError);
-            } else {
-                log.debug("Auto-commit before reconciling new assignment completed successfully.");
-            }
-
-            if (!maybeAbortReconciliation()) {
-                revokeAndAssign(resolvedAssignment, assignedTopicIdPartitions, revokedPartitions, addedPartitions);
-            }
+        try {
+            // Mark partitions as pending revocation to stop fetching from the partitions (no new
+            // fetches sent out, and no in-flight fetches responses processed).
+            markPendingRevocationToPauseFetching(revokedPartitions);
+
+            // Commit offsets if auto-commit enabled before reconciling a new assignment. Request will
+            // be retried until it succeeds, fails with non-retriable error, or timer expires.
+            CompletableFuture<Void> commitResult = signalReconciliationStarted();
+
+            // Execute commit -> onPartitionsRevoked -> onPartitionsAssigned.
+            commitResult.whenComplete((__, commitReqError) -> {
+                try {
+                    if (commitReqError != null) {
+                        // The call to commit, that includes retry logic for retriable errors, failed to
+                        // complete within the time boundaries (fatal error or retriable that did not
+                        // recover). Proceed with the revocation.
+                        log.error("Auto-commit request before reconciling new assignment failed. " +
+                            "Will proceed with the reconciliation anyway.", commitReqError);
+                    } else {
+                        log.debug("Auto-commit before reconciling new assignment completed successfully.");
+                    }
+
+                    if (!maybeAbortReconciliation()) {
+                        revokeAndAssign(resolvedAssignment, assignedTopicIdPartitions, revokedPartitions, addedPartitions);
+                    }
+                } catch (RuntimeException error) {
+                    handleReconciliationFailure(error);
+                }
+            });
+        } catch (RuntimeException error) {
+            handleReconciliationFailure(error);
+        }
+    }
 
-        }).exceptionally(error -> {
-            if (error != null) {
-                log.error("Reconciliation failed.", error);
-            }
-            return null;
-        });
+    private void handleReconciliationFailure(RuntimeException error) {
+        log.error("Reconciliation failed.", error);
+        markReconciliationCompleted();
     }
 
     long getDeadlineMsForTimeout(final long timeoutMs) {

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

@@ -92,6 +92,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -2002,6 +2003,48 @@ public void testListenerCallbacksThrowsErrorOnPartitionsRevoked() {
         testErrorsOnPartitionsRevoked(new IllegalArgumentException("Intentional onPartitionsRevoked() error"));
     }
 
+    @Test
+    public void testReconciliationRetriesAfterSynchronousFailure() {
+        ConsumerMembershipManager membershipManager = createMemberInStableState();
+        Uuid topicId = Uuid.randomUuid();
+        mockOwnedPartition(membershipManager, topicId, "topic1");
+        receiveEmptyAssignment(membershipManager);
+
+        doThrow(new IllegalStateException("Intentional reconciliation failure"))
+            .doNothing()
+            .when(membershipManager)
+            .markPendingRevocationToPauseFetching(anySet());
+
+        assertDoesNotThrow(() -> membershipManager.maybeReconcile(true));
+        assertFalse(membershipManager.reconciliationInProgress());
+
+        membershipManager.maybeReconcile(true);
+        processAssignmentEventNoCallback(membershipManager);
+
+        assertFalse(membershipManager.reconciliationInProgress());
+    }
+
+    @Test
+    public void testReconciliationRetriesAfterSynchronousCallbackSetupFailure() {
+        ConsumerMembershipManager membershipManager = createMemberInStableState();
+        Uuid topicId = Uuid.randomUuid();
+        mockOwnedPartition(membershipManager, topicId, "topic1");
+        receiveEmptyAssignment(membershipManager);
+
+        doThrow(new IllegalStateException("Intentional callback setup failure"))
+            .doNothing()
+            .when(membershipManager)
+            .signalPartitionsBeingRevoked(anySet());
+
+        assertDoesNotThrow(() -> membershipManager.maybeReconcile(true));
+        assertFalse(membershipManager.reconciliationInProgress());
+
+        membershipManager.maybeReconcile(true);
+        processAssignmentEventNoCallback(membershipManager);
+
+        assertFalse(membershipManager.reconciliationInProgress());
+    }
+
     private void testErrorsOnPartitionsRevoked(RuntimeException error) {
         // Step 1: set up mocks
         String topicName = "topic1";