Repository: apache/kafka
File: streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamTask.java (Java)

Code before:
                        log.debug("Skipped preparing {} task with id {} for commit since the task is getting closed dirty.", state(), id);
                        return null;
                    }
                    hasPendingTxCommit = eosEnabled;

                    log.debug("Prepared {} task for committing", state());
                    return committableOffsetsAndMetadata();
                } else {
                    log.debug("Skipped preparing {} task for commit since there is nothing to commit", state());
                    return Collections.emptyMap();
                }

            case CLOSED:
                throw new IllegalStateException("Illegal state " + state() + " while preparing active task " + id + " for committing");

            default:
                throw new IllegalStateException("Unknown state " + state() + " while preparing active task " + id + " for committing");
        }
    }

    private Optional<OffsetAndMetadata> findOffsetAndMetadata(final TopicPartition partition) {
        Long offset = partitionGroup.headRecordOffset(partition);
        Optional<Integer> leaderEpoch = partitionGroup.headRecordLeaderEpoch(partition);
        final long partitionTime = partitionGroup.partitionTimestamp(partition);
        if (offset == null) {
            if (nextOffsetsAndMetadataToBeConsumed.containsKey(partition)) {
                final OffsetAndMetadata offsetAndMetadata = nextOffsetsAndMetadataToBeConsumed.get(partition);
                offset = offsetAndMetadata.offset();
                leaderEpoch = offsetAndMetadata.leaderEpoch();
            } else {
                // it may be that we have not yet consumed any record from this partition, hence nothing to commit
                return Optional.empty();
            }
        }
        return Optional.of(new OffsetAndMetadata(offset,
                leaderEpoch,
                new TopicPartitionMetadata(partitionTime, processorContext.processorMetadata()).encode()));
    }

    private Map<TopicPartition, OffsetAndMetadata> committableOffsetsAndMetadata() {
        switch (state()) {
            case CREATED:
            case RESTORING:
                return Collections.emptyMap();

            case RUNNING:
            case SUSPENDED:

                // If there's processor metadata to be committed. We need to commit them to all
                // input partitions
                final Set<TopicPartition> partitionsNeedCommit = processorContext.processorMetadata().needsCommit() ?

Review comment (performance):
the change that reworked that code to use cached offsets: https://github.com/apache/kafka/pull/17091/files#diff-a76674468cda8772230fb8411717cf9068b1a363a792f32c602fb2ec5ba9efd7R472

@aliehsaeedii @lucasbru, as you folks worked on it, could you please take a look? thanks

Code after:
                        log.debug("Skipped preparing {} task with id {} for commit since the task is getting closed dirty.", state(), id);
                        return null;
                    }
                    hasPendingTxCommit = eosEnabled;

                    log.debug("Prepared {} task for committing", state());
                    return committableOffsetsAndMetadata();
                } else {
                    log.debug("Skipped preparing {} task for commit since there is nothing to commit", state());
                    return Collections.emptyMap();
                }

            case CLOSED:
                throw new IllegalStateException("Illegal state " + state() + " while preparing active task " + id + " for committing");

            default:
                throw new IllegalStateException("Unknown state " + state() + " while preparing active task " + id + " for committing");
        }
    }

    private Optional<OffsetAndMetadata> findOffsetAndMetadata(final TopicPartition partition) {
        Long offset = partitionGroup.headRecordOffset(partition);
        Optional<Integer> leaderEpoch = partitionGroup.headRecordLeaderEpoch(partition);
        final long partitionTime = partitionGroup.partitionTimestamp(partition);
        if (offset == null) {
            final OffsetAndMetadata offsetAndMetadata = nextOffsetsAndMetadataToBeConsumed.get(partition);
            if (offsetAndMetadata == null) {
                // it may be that we have not yet consumed any record from this partition, hence nothing to commit
                return Optional.empty();
            }
            offset = offsetAndMetadata.offset();
            leaderEpoch = offsetAndMetadata.leaderEpoch();
        }
        return Optional.of(new OffsetAndMetadata(offset,
                leaderEpoch,
                new TopicPartitionMetadata(partitionTime, processorContext.processorMetadata()).encode()));
    }

    private Map<TopicPartition, OffsetAndMetadata> committableOffsetsAndMetadata() {
        switch (state()) {
            case CREATED:
            case RESTORING:
                return Collections.emptyMap();

            case RUNNING:
            case SUSPENDED:

                // If there's processor metadata to be committed. We need to commit them to all
                // input partitions
                final Set<TopicPartition> partitionsNeedCommit = processorContext.processorMetadata().needsCommit() ?
                    inputPartitions() : consumedOffsets.keySet();