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

Code before:
            final Map<TaskId, Set<TopicPartition>> standbyTasksWithPartitions =
                pairWithTopicPartitions(Stream.concat(assignment.standbyTasks().stream(), assignment.warmupTasks().stream()));

            log.info("Processing new assignment {} from Streams Rebalance Protocol", assignment);

            taskManager.handleAssignment(activeTasksWithPartitions, standbyTasksWithPartitions);
            streamThread.setState(StreamThread.State.PARTITIONS_ASSIGNED);
            taskManager.handleRebalanceComplete();
        } catch (final Exception exception) {
            return Optional.of(exception);
        } finally {
            tasksAssignedSensor.record(time.milliseconds() - start);
        }
        return Optional.empty();
    }

    @Override
    public Optional<Exception> onAllTasksLost() {
        final long start = time.milliseconds();
        try {
            taskManager.handleLostAll();
        } catch (final Exception exception) {
            return Optional.of(exception);
        } finally {
            tasksLostSensor.record(time.milliseconds() - start);
        }
        return Optional.empty();
    }

    private Map<TaskId, Set<TopicPartition>> pairWithTopicPartitions(final Stream<StreamsRebalanceData.TaskId> taskIdStream) {
        return taskIdStream
            .collect(Collectors.toMap(
                this::toTaskId,
                task -> toTopicPartitions(task, streamsRebalanceData.subtopologies().get(task.subtopologyId()))
            ));
    }

    private TaskId toTaskId(final StreamsRebalanceData.TaskId task) {
        return new TaskId(Integer.parseInt(task.subtopologyId()), task.partitionId());
    }

    private Set<TopicPartition> toTopicPartitions(final StreamsRebalanceData.TaskId task,
                                                  final StreamsRebalanceData.Subtopology subTopology) {
        return
            Stream.concat(
                    subTopology.sourceTopics().stream(),
                    subTopology.repartitionSourceTopics().keySet().stream()
                )
                .map(t -> new TopicPartition(t, task.partitionId()))
                .collect(Collectors.toSet());
    }

Review comment (bug):
Similar to onTasksAssigned, this method records latency metrics even when exceptions occur. The finally block will execute regardless of whether the operation completed successfully or failed, which may not provide accurate latency measurements for successful operations.

Code after:
            streamThread.setState(StreamThread.State.PARTITIONS_ASSIGNED);
            taskManager.handleRebalanceComplete();
            streamsRebalanceData.setReconciledAssignment(assignment);
        } finally {
            tasksAssignedSensor.record(time.milliseconds() - start);
        }
    }

    @Override
    public void onAllTasksLost() {
        final long start = time.milliseconds();
        try {
            taskManager.handleLostAll();
            streamsRebalanceData.setReconciledAssignment(StreamsRebalanceData.Assignment.EMPTY);
        } finally {
            tasksLostSensor.record(time.milliseconds() - start);
        }
    }

    private Map<TaskId, Set<TopicPartition>> pairWithTopicPartitions(final Stream<StreamsRebalanceData.TaskId> taskIdStream) {
        return taskIdStream
            .collect(Collectors.toMap(
                this::toTaskId,
                task -> toTopicPartitions(task, streamsRebalanceData.subtopologies().get(task.subtopologyId()))
            ));
    }

    private TaskId toTaskId(final StreamsRebalanceData.TaskId task) {
        return new TaskId(Integer.parseInt(task.subtopologyId()), task.partitionId());
    }

    private Set<TopicPartition> toTopicPartitions(final StreamsRebalanceData.TaskId task,
                                                  final StreamsRebalanceData.Subtopology subTopology) {
        return
            Stream.concat(
                    subTopology.sourceTopics().stream(),
                    subTopology.repartitionSourceTopics().keySet().stream()
                )
                .map(t -> new TopicPartition(t, task.partitionId()))
                .collect(Collectors.toSet());
    }
}