KAFKA-6607: Commit correct offsets for transactional input data

clients/src/test/java/org/apache/kafka/test/TestUtils.java

@@ -282,7 +282,6 @@ public static Properties producerConfig(final String bootstrapServers,
         final Properties properties = new Properties();
         properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
         properties.put(ProducerConfig.ACKS_CONFIG, "all");
-        properties.put(ProducerConfig.RETRIES_CONFIG, 0);
         properties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySerializer);
         properties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSerializer);
         properties.putAll(additional);

streams/src/main/java/org/apache/kafka/streams/processor/internals/PartitionGroup.java

@@ -58,15 +58,14 @@ public class PartitionGroup {
     private int totalBuffered;
     private boolean allBuffered;
 
-
-    public static class RecordInfo {
+    static class RecordInfo {
         RecordQueue queue;
 
-        public ProcessorNode node() {
+        ProcessorNode node() {
             return queue.source();
         }
 
-        public TopicPartition partition() {
+        TopicPartition partition() {
             return queue.partition();
         }
 
@@ -84,20 +83,23 @@ RecordQueue queue() {
         streamTime = RecordQueue.UNKNOWN;
     }
 
-    // visible for testing
     long partitionTimestamp(final TopicPartition partition) {
         final RecordQueue queue = partitionQueues.get(partition);
+
         if (queue == null) {
-            throw new NullPointerException("Partition " + partition + " not found.");
+            throw new IllegalStateException("Partition " + partition + " not found.");
         }
+
         return queue.partitionTime();
     }
 
     void setPartitionTime(final TopicPartition partition, final long partitionTime) {
         final RecordQueue queue = partitionQueues.get(partition);
+
         if (queue == null) {
-            throw new NullPointerException("Partition " + partition + " not found.");
+            throw new IllegalStateException("Partition " + partition + " not found.");
         }
+
         if (streamTime < partitionTime) {
             streamTime = partitionTime;
         }
@@ -152,6 +154,10 @@ record = queue.poll();
     int addRawRecords(final TopicPartition partition, final Iterable<ConsumerRecord<byte[], byte[]>> rawRecords) {
         final RecordQueue recordQueue = partitionQueues.get(partition);
 
+        if (recordQueue == null) {
+            throw new IllegalStateException("Partition " + partition + " not found.");
+        }
+
         final int oldSize = recordQueue.size();
         final int newSize = recordQueue.addRawRecords(rawRecords);
 
@@ -172,25 +178,35 @@ int addRawRecords(final TopicPartition partition, final Iterable<ConsumerRecord<
         return newSize;
     }
 
-    public Set<TopicPartition> partitions() {
+    Set<TopicPartition> partitions() {
         return Collections.unmodifiableSet(partitionQueues.keySet());
     }
 
     /**
      * Return the stream-time of this partition group defined as the largest timestamp seen across all partitions
      */
-    public long streamTime() {
+    long streamTime() {
         return streamTime;
     }
 
+    Long headRecordOffset(final TopicPartition partition) {
+        final RecordQueue recordQueue = partitionQueues.get(partition);
+
+        if (recordQueue == null) {
+            throw new IllegalStateException("Partition " + partition + " not found.");
+        }
+
+        return recordQueue.headRecordOffset();
+    }
+
     /**
      * @throws IllegalStateException if the record's partition does not belong to this partition group
      */
     int numBuffered(final TopicPartition partition) {
         final RecordQueue recordQueue = partitionQueues.get(partition);
 
         if (recordQueue == null) {
-            throw new IllegalStateException(String.format("Record's partition %s does not belong to this partition-group.", partition));
+            throw new IllegalStateException("Partition " + partition + " not found.");
         }
 
         return recordQueue.size();
@@ -204,14 +220,15 @@ boolean allPartitionsBuffered() {
         return allBuffered;
     }
 
-    public void close() {
+    void close() {
         clear();
         partitionQueues.clear();
     }
 
-    public void clear() {
+    void clear() {
         nonEmptyQueuesByTime.clear();
         streamTime = RecordQueue.UNKNOWN;
+        totalBuffered = 0;
         for (final RecordQueue queue : partitionQueues.values()) {
             queue.clear();
         }

streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordQueue.java

@@ -75,7 +75,7 @@ public class RecordQueue {
         );
         this.log = logContext.logger(RecordQueue.class);
     }
- 
+
     void setPartitionTime(final long partitionTime) {
         this.partitionTime = partitionTime;
     }
@@ -156,6 +156,10 @@ public long headRecordTimestamp() {
         return headRecord == null ? UNKNOWN : headRecord.timestamp;
     }
 
+    public Long headRecordOffset() {
+        return headRecord == null ? null : headRecord.offset();
+    }
+
     /**
      * Clear the fifo queue of its elements, also clear the time tracker's kept stamped elements
      */

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamTask.java

@@ -496,9 +496,24 @@ void commit(final boolean startNewTransaction, final Map<TopicPartition, Long> p
         }
 
         final Map<TopicPartition, OffsetAndMetadata> consumedOffsetsAndMetadata = new HashMap<>(consumedOffsets.size());
+
         for (final Map.Entry<TopicPartition, Long> entry : consumedOffsets.entrySet()) {
             final TopicPartition partition = entry.getKey();
-            final long offset = entry.getValue() + 1;
+            Long offset = partitionGroup.headRecordOffset(partition);
+            if (offset == null) {
+                try {
+                    offset = consumer.position(partition);
+                } catch (final TimeoutException error) {
+                    // the `consumer.position()` call should never block, because we know that we did process data
+                    // for the requested partition and thus the consumer should have a valid local position
+                    // that it can return immediately
+
+                    // hence, a `TimeoutException` indicates a bug and thus we rethrow it as fatal `IllegalStateException`
+                    throw new IllegalStateException(error);
+                } catch (final KafkaException fatal) {
+                    throw new StreamsException(fatal);
+                }
+            }
             final long partitionTime = partitionTimes.get(partition);
             consumedOffsetsAndMetadata.put(partition, new OffsetAndMetadata(offset, encodeTimestamp(partitionTime)));
         }
@@ -621,6 +636,8 @@ void suspend(final boolean clean,
             try {
                 commit(false, partitionTimes);
             } finally {
+                partitionGroup.clear();
+
                 if (eosEnabled) {
                     stateMgr.checkpoint(activeTaskCheckpointableOffsets());
 
@@ -677,8 +694,6 @@ private void maybeAbortTransactionAndCloseRecordCollector(final boolean isZombie
     private void closeTopology() {
         log.trace("Closing processor topology");
 
-        partitionGroup.clear();
-
         // close the processors
         // make sure close() is called for each node even when there is a RuntimeException
         RuntimeException exception = null;

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java

@@ -754,6 +754,13 @@ void runOnce() {
             throw new StreamsException(logPrefix + "Unexpected state " + state + " during normal iteration");
         }
 
+        final long pollLatency = advanceNowAndComputeLatency();
+
+        if (records != null && !records.isEmpty()) {
+            pollSensor.record(pollLatency, now);
+            addRecordsToTasks(records);
+        }
+
         // Shutdown hook could potentially be triggered and transit the thread state to PENDING_SHUTDOWN during #pollRequests().
         // The task manager internal states could be uninitialized if the state transition happens during #onPartitionsAssigned().
         // Should only proceed when the thread is still running after #pollRequests(), because no external state mutation
@@ -763,13 +770,6 @@ void runOnce() {
             return;
         }
 
-        final long pollLatency = advanceNowAndComputeLatency();
-
-        if (records != null && !records.isEmpty()) {
-            pollSensor.record(pollLatency, now);
-            addRecordsToTasks(records);
-        }
-
         // only try to initialize the assigned tasks
         // if the state is still in PARTITION_ASSIGNED after the poll call
         if (state == State.PARTITIONS_ASSIGNED) {
@@ -917,6 +917,13 @@ private void addRecordsToTasks(final ConsumerRecords<byte[], byte[]> records) {
             final StreamTask task = taskManager.activeTask(partition);
 
             if (task == null) {
+                if (!isRunning()) {
+                    // if we are in PENDING_SHUTDOWN and don't find the task it implies that it was a newly assigned
+                    // task that we just skipped to create;
+                    // hence, we just skip adding the corresponding records
+                    continue;
+                }
+
                 log.error(
                     "Unable to locate active task for received-record partition {}. Current tasks: {}",
                     partition,