Source:
<code>void handleRevocation(final Collection<TopicPartition> revokedPartitions) { final Set<TopicPartition> remainingPartitions = new HashSet<>(revokedPartitions); final Map<TaskId, Map<TopicPartition, OffsetAndMetadata>> consumedOffsetsAndMetadataPerTask = new HashMap<>(); final Set<Task> suspended = new HashSet<>(); for (final Task task : tasks.values()) { if (remainingPartitions.containsAll(task.inputPartitions())) { task.prepareSuspend(); final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.committableOffsetsAndMetadata(); if (!committableOffsets.isEmpty()) { consumedOffsetsAndMetadataPerTask.put(task.id(), committableOffsets); } suspended.add(task); <START> } else if (eosBetaEnabled && !eosUpgradeModeEnabled) { <END> task.prepareCommit(); final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.committableOffsetsAndMetadata(); if (!committableOffsets.isEmpty()) { consumedOffsetsAndMetadataPerTask.put(task.id(), committableOffsets); } } remainingPartitions.removeAll(task.inputPartitions()); } if (!consumedOffsetsAndMetadataPerTask.isEmpty()) { commitOffsetsOrTransaction(consumedOffsetsAndMetadataPerTask); } for (final Task task : tasks.values()) { if (consumedOffsetsAndMetadataPerTask.containsKey(task.id())) { if (suspended.contains(task)) { task.suspend(); } else { task.postCommit(); } } } if (!remainingPartitions.isEmpty()) { log.warn("The following partitions are missing from the task partitions. It could potentially " + "due to race condition of consumer detecting the heartbeat failure, or the tasks " + "have been cleaned up by the handleAssignment callback.", remainingPartitions); } }</code><technical_language>for this PR: this is messy here... instead of interleaving suspending committing logic, do loop tasks gathering offsets, do another loop tasks revokedPartitions a suspend</technical_language>

Target:
void handleRevocation(final Collection<TopicPartition> revokedPartitions) { final Set<TopicPartition> remainingPartitions = new HashSet<>(revokedPartitions); final Map<TaskId, Map<TopicPartition, OffsetAndMetadata>> consumedOffsetsAndMetadataPerTask = new HashMap<>(); final Set<Task> suspended = new HashSet<>(); for (final Task task : tasks.values()) { if (remainingPartitions.containsAll(task.inputPartitions())) { task.prepareSuspend(); final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.committableOffsetsAndMetadata(); if (!committableOffsets.isEmpty()) { consumedOffsetsAndMetadataPerTask.put(task.id(), committableOffsets); } suspended.add(task); } else { if (task.isActive() && task.commitNeeded()) { task.prepareCommit(); final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.committableOffsetsAndMetadata(); if (!committableOffsets.isEmpty()) { consumedOffsetsAndMetadataPerTask.put(task.id(), committableOffsets); } } } remainingPartitions.removeAll(task.inputPartitions()); } if (!consumedOffsetsAndMetadataPerTask.isEmpty()) { commitOffsetsOrTransaction(consumedOffsetsAndMetadataPerTask); } for (final Task task : tasks.values()) { if (consumedOffsetsAndMetadataPerTask.containsKey(task.id())) { if (suspended.contains(task)) { task.suspend(); } else { task.postCommit(); } } } if (!remainingPartitions.isEmpty()) { log.warn("The following partitions are missing from the task partitions. It could potentially " + "due to race condition of consumer detecting the heartbeat failure, or the tasks " + "have been cleaned up by the handleAssignment callback.", remainingPartitions); } }