impl

Author: lucasbru

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

@@ -984,6 +984,8 @@ public void close(final Timer timer) {
             }
         } finally {
             super.close(timer);
+            // Super-class close may wait for more commit callbacks to complete.
+            invokeCompletedOffsetCommitCallbacks();
         }
     }
 
@@ -1033,16 +1035,22 @@ public RequestFuture<Void> commitOffsetsAsync(final Map<TopicPartition, OffsetAn
             lookupCoordinator().addListener(new RequestFutureListener<Void>() {
                 @Override
                 public void onSuccess(Void value) {
-                    pendingAsyncCommits.decrementAndGet();
-                    doCommitOffsetsAsync(offsets, callback);
-                    client.pollNoWakeup();
+                    try {
+                        doCommitOffsetsAsync(offsets, callback);
+                        client.pollNoWakeup();
+                    } finally {
+                        pendingAsyncCommits.decrementAndGet();
+                    }
                 }
 
                 @Override
                 public void onFailure(RuntimeException e) {
-                    pendingAsyncCommits.decrementAndGet();
-                    completedOffsetCommits.add(new OffsetCommitCompletion(callback, offsets,
+                    try {
+                        completedOffsetCommits.add(new OffsetCommitCompletion(callback, offsets,
                             new RetriableCommitFailedException(e)));
+                    } finally {
+                        pendingAsyncCommits.decrementAndGet();
+                    }
                 }
             });
         }
@@ -1061,25 +1069,29 @@ private RequestFuture<Void> doCommitOffsetsAsync(final Map<TopicPartition, Offse
         future.addListener(new RequestFutureListener<Void>() {
             @Override
             public void onSuccess(Void value) {
-                inFlightAsyncCommits.decrementAndGet();
-
-                if (interceptors != null)
-                    interceptors.onCommit(offsets);
-                completedOffsetCommits.add(new OffsetCommitCompletion(cb, offsets, null));
+                try {
+                    if (interceptors != null)
+                        interceptors.onCommit(offsets);
+                    completedOffsetCommits.add(new OffsetCommitCompletion(cb, offsets, null));
+                } finally {
+                    inFlightAsyncCommits.decrementAndGet();
+                }
             }
 
             @Override
             public void onFailure(RuntimeException e) {
-                inFlightAsyncCommits.decrementAndGet();
+                try {
+                    Exception commitException = e;
 
-                Exception commitException = e;
-
-                if (e instanceof RetriableException) {
-                    commitException = new RetriableCommitFailedException(e);
-                }
-                completedOffsetCommits.add(new OffsetCommitCompletion(cb, offsets, commitException));
-                if (commitException instanceof FencedInstanceIdException) {
-                    asyncCommitFenced.set(true);
+                    if (e instanceof RetriableException) {
+                        commitException = new RetriableCommitFailedException(e);
+                    }
+                    completedOffsetCommits.add(new OffsetCommitCompletion(cb, offsets, commitException));
+                    if (commitException instanceof FencedInstanceIdException) {
+                        asyncCommitFenced.set(true);
+                    }
+                } finally {
+                    inFlightAsyncCommits.decrementAndGet();
                 }
             }
         });
@@ -1164,7 +1176,8 @@ public void maybeAutoCommitOffsetsAsync(long now) {
     }
 
     private boolean invokePendingAsyncCommits(Timer timer) {
-        if (inFlightAsyncCommits.get() == 0) {
+        if (pendingAsyncCommits.get() == 0 && inFlightAsyncCommits.get() == 0) {
+            invokeCompletedOffsetCommitCallbacks();
             return true;
         }
 
@@ -1174,7 +1187,8 @@ private boolean invokePendingAsyncCommits(Timer timer) {
             client.poll(timer);
             invokeCompletedOffsetCommitCallbacks();
 
-            if (inFlightAsyncCommits.get() == 0) {
+            if (pendingAsyncCommits.get() == 0 && inFlightAsyncCommits.get() == 0) {
+                invokeCompletedOffsetCommitCallbacks();
                 return true;
             }
 

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

@@ -229,7 +229,11 @@ private GroupRebalanceConfig buildRebalanceConfig(Optional<String> groupInstance
     @AfterEach
     public void teardown() {
         this.metrics.close();
-        this.coordinator.close(time.timer(0));
+        try {
+            this.coordinator.close(time.timer(0));
+        } catch (Exception e) {
+            // ignore
+        }
     }
 
     @Test

core/src/test/scala/integration/kafka/api/PlaintextConsumerCommitTest.scala

@@ -304,6 +304,62 @@ class PlaintextConsumerCommitTest extends AbstractConsumerTest {
     consumeAndVerifyRecords(consumer = otherConsumer, numRecords = 1, startingOffset = 5, startingTimestamp = startingTimestamp)
   }
 
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testCommitAsyncCompletedBeforeConsumerCloses(quorum: String, groupProtocol: String): Unit = {
+    // This is testing the contract that asynchronous offset commit are completed before the consumer
+    // is closed, even when no commit sync is performed as part of the close (due to auto-commit
+    // disabled, or simply because there no consumed offsets).
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 3, tp)
+    sendRecords(producer, numRecords = 3, tp2)
+
+    val consumer = createConsumer()
+    consumer.assign(List(tp, tp2).asJava)
+
+    // Try without looking up the coordinator first
+    val cb = new CountConsumerCommitCallback
+    consumer.commitAsync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(1L))).asJava, cb)
+    consumer.commitAsync(Map[TopicPartition, OffsetAndMetadata]((tp2, new OffsetAndMetadata(1L))).asJava, cb)
+    consumer.close()
+    assertEquals(2, cb.successCount)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testCommitAsyncCompletedBeforeCommitSyncReturns(quorum: String, groupProtocol: String): Unit = {
+    // This is testing the contract that asynchronous offset commits sent previously with the
+    // `commitAsync` are guaranteed to have their callbacks invoked prior to completion of
+    // `commitSync` (given that it does not time out).
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 3, tp)
+    sendRecords(producer, numRecords = 3, tp2)
+
+    val consumer = createConsumer()
+    consumer.assign(List(tp, tp2).asJava)
+
+    // Try without looking up the coordinator first
+    val cb = new CountConsumerCommitCallback
+    consumer.commitAsync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(1L))).asJava, cb)
+    consumer.commitSync(Map.empty[TopicPartition, OffsetAndMetadata].asJava)
+    assertEquals(1, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(1, cb.successCount)
+
+    // Try with coordinator known
+    consumer.commitAsync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(2L))).asJava, cb)
+    consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp2, new OffsetAndMetadata(2L))).asJava)
+    assertEquals(2, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(2, consumer.committed(Set(tp2).asJava).get(tp2).offset)
+    assertEquals(2, cb.successCount)
+
+    // Try with empty sync commit
+    consumer.commitAsync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(3L))).asJava, cb)
+    consumer.commitSync(Map.empty[TopicPartition, OffsetAndMetadata].asJava)
+    assertEquals(3, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(2, consumer.committed(Set(tp2).asJava).get(tp2).offset)
+    assertEquals(3, cb.successCount)
+  }
+
   def changeConsumerSubscriptionAndValidateAssignment[K, V](consumer: Consumer[K, V],
                                                             topicsToSubscribe: List[String],
                                                             expectedAssignment: Set[TopicPartition],
@@ -314,6 +370,8 @@ class PlaintextConsumerCommitTest extends AbstractConsumerTest {
 }
 
 object PlaintextConsumerCommitTest {
+  def getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly: Stream[Arguments] =
+    BaseConsumerTest.getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly()
 
   def getTestQuorumAndGroupProtocolParametersAll: Stream[Arguments] =
     BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll()