KAFKA-13472: Correct last-committed offsets tracking for sink tasks after partial revocation

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java

@@ -274,8 +274,9 @@ private void onCommitCompleted(Throwable error, long seqno, Map<TopicPartition,
                 log.debug("{} Finished offset commit successfully in {} ms for sequence number {}: {}",
                         this, durationMillis, seqno, committedOffsets);
                 if (committedOffsets != null) {
-                    log.debug("{} Setting last committed offsets to {}", this, committedOffsets);
-                    lastCommittedOffsets = committedOffsets;
+                    log.trace("{} Adding to last committed offsets: {}", this, committedOffsets);
+                    lastCommittedOffsets.putAll(committedOffsets);
+                    log.debug("{} Last committed offsets are now {}", this, committedOffsets);
                     sinkTaskMetricsGroup.recordCommittedOffsets(committedOffsets);
                 }
                 commitFailures = 0;
@@ -651,9 +652,10 @@ private void closePartitions(Collection<TopicPartition> topicPartitions, boolean
                 workerErrantRecordReporter.cancelFutures(topicPartitions);
                 log.trace("Cancelled all reported errors for {}", topicPartitions);
             }
-            topicPartitions.forEach(currentOffsets::remove);
+            currentOffsets.keySet().removeAll(topicPartitions);
         }
         updatePartitionCount();
+        lastCommittedOffsets.keySet().removeAll(topicPartitions);
     }
 
     private void updatePartitionCount() {
@@ -892,10 +894,8 @@ void assignedOffsets(Map<TopicPartition, OffsetAndMetadata> offsets) {
         }
 
         void clearOffsets(Collection<TopicPartition> topicPartitions) {
-            topicPartitions.forEach(tp -> {
-                consumedOffsets.remove(tp);
-                committedOffsets.remove(tp);
-            });
+            consumedOffsets.keySet().removeAll(topicPartitions);
+            committedOffsets.keySet().removeAll(topicPartitions);
             computeSinkRecordLag();
         }
 

connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSinkTask.java

@@ -27,6 +27,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -44,7 +45,7 @@ public class VerifiableSinkTask extends SinkTask {
     private String name; // Connector name
     private int id; // Task ID
 
-    private ArrayList<Map<String, Object>> unflushed = new ArrayList<>();
+    private final Map<TopicPartition, List<Map<String, Object>>> unflushed = new HashMap<>();
 
     @Override
     public String version() {
@@ -80,25 +81,33 @@ public void put(Collection<SinkRecord> records) {
                 dataJson = "Bad data can't be written as json: " + e.getMessage();
             }
             System.out.println(dataJson);
-            unflushed.add(data);
+            unflushed.computeIfAbsent(
+                    new TopicPartition(record.topic(), record.kafkaPartition()),
+                    tp -> new ArrayList<>()
+            ).add(data);
         }
     }
 
     @Override
     public void flush(Map<TopicPartition, OffsetAndMetadata> offsets) {
         long nowMs = System.currentTimeMillis();
-        for (Map<String, Object> data : unflushed) {
-            data.put("time_ms", nowMs);
-            data.put("flushed", true);
-            String dataJson;
-            try {
-                dataJson = JSON_SERDE.writeValueAsString(data);
-            } catch (JsonProcessingException e) {
-                dataJson = "Bad data can't be written as json: " + e.getMessage();
+        for (TopicPartition topicPartition : offsets.keySet()) {
+            if (!unflushed.containsKey(topicPartition)) {
+                continue;
             }
-            System.out.println(dataJson);
+            for (Map<String, Object> data : unflushed.get(topicPartition)) {
+                data.put("time_ms", nowMs);
+                data.put("flushed", true);
+                String dataJson;
+                try {
+                    dataJson = JSON_SERDE.writeValueAsString(data);
+                } catch (JsonProcessingException e) {
+                    dataJson = "Bad data can't be written as json: " + e.getMessage();
+                }
+                System.out.println(dataJson);
+            }
+            unflushed.remove(topicPartition);
         }
-        unflushed.clear();
     }
 
     @Override

connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java

@@ -660,6 +660,90 @@ public void testPartialRevocationAndAssignment() throws Exception {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testPreCommitFailureAfterPartialRevocationAndAssignment() throws Exception {
+        createTask(initialState);
+
+        // First poll; assignment is [TP1, TP2]
+        expectInitializeTask();
+        expectTaskGetTopic(true);
+        expectPollInitialAssignment();
+
+        // Second poll; a single record is delivered from TP1
+        expectConsumerPoll(1);
+        expectConversionAndTransformation(1);
+        sinkTask.put(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        // Third poll; assignment changes to [TP2]
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
+                () -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(Collections.singleton(TOPIC_PARTITION));
+                    rebalanceListener.getValue().onPartitionsAssigned(Collections.emptySet());
+                    return ConsumerRecords.empty();
+                });
+        EasyMock.expect(consumer.assignment()).andReturn(Collections.singleton(TOPIC_PARTITION)).times(2);
+        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        sinkTask.preCommit(offsets);
+        EasyMock.expectLastCall().andReturn(offsets);
+        consumer.commitSync(offsets);
+        EasyMock.expectLastCall();
+        sinkTask.close(Collections.singleton(TOPIC_PARTITION));
+        EasyMock.expectLastCall();
+        sinkTask.put(Collections.emptyList());
+        EasyMock.expectLastCall();
+
+        // Fourth poll; assignment changes to [TP2, TP3]
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
+                () -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet());
+                    rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3));
+                    return ConsumerRecords.empty();
+                });
+        EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3))).times(2);
+        EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(FIRST_OFFSET);
+        sinkTask.open(Collections.singleton(TOPIC_PARTITION3));
+        EasyMock.expectLastCall();
+        sinkTask.put(Collections.emptyList());
+        EasyMock.expectLastCall();
+
+        // Fifth poll; an offset commit takes place
+        EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3))).times(2);
+        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
+        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        workerCurrentOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET));
+        sinkTask.preCommit(workerCurrentOffsets);
+        EasyMock.expectLastCall().andThrow(new ConnectException("Failed to flush"));
+
+        consumer.seek(TOPIC_PARTITION2, FIRST_OFFSET);
+        EasyMock.expectLastCall();
+        consumer.seek(TOPIC_PARTITION3, FIRST_OFFSET);
+        EasyMock.expectLastCall();
+
+        expectConsumerPoll(0);
+        sinkTask.put(EasyMock.eq(Collections.emptyList()));
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        // First iteration--first call to poll, first consumer assignment
+        workerTask.iteration();
+        // Second iteration--second call to poll, delivery of one record
+        workerTask.iteration();
+        // Third iteration--third call to poll, partial consumer revocation
+        workerTask.iteration();
+        // Fourth iteration--fourth call to poll, partial consumer assignment
+        workerTask.iteration();
+        // Fifth iteration--task-requested offset commit with failure in SinkTask::preCommit
+        sinkTaskContext.getValue().requestCommit();
+        workerTask.iteration();
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testWakeupInCommitSyncCausesRetry() throws Exception {
         createTask(initialState);
@@ -970,6 +1054,51 @@ public void testPreCommit() throws Exception {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testPreCommitFailure() throws Exception {
+        createTask(initialState);
+
+        expectInitializeTask();
+        expectTaskGetTopic(true);
+        EasyMock.expect(consumer.assignment()).andStubReturn(INITIAL_ASSIGNMENT);
+
+        // iter 1
+        expectPollInitialAssignment();
+
+        // iter 2
+        expectConsumerPoll(2);
+        expectConversionAndTransformation(2);
+        sinkTask.put(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        // iter 3
+        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
+        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
+        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        sinkTask.preCommit(workerCurrentOffsets);
+        EasyMock.expectLastCall().andThrow(new ConnectException("Failed to flush"));
+
+        consumer.seek(TOPIC_PARTITION, FIRST_OFFSET);
+        EasyMock.expectLastCall();
+        consumer.seek(TOPIC_PARTITION2, FIRST_OFFSET);
+        EasyMock.expectLastCall();
+
+        expectConsumerPoll(0);
+        sinkTask.put(EasyMock.eq(Collections.emptyList()));
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        workerTask.iteration(); // iter 1 -- initial assignment
+        workerTask.iteration(); // iter 2 -- deliver 2 records
+        sinkTaskContext.getValue().requestCommit();
+        workerTask.iteration(); // iter 3 -- commit
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testIgnoredCommit() throws Exception {
         createTask(initialState);