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

Code before:
            final Set<TopicPartition> partitionsToRevoke = activeTasksToRevokeWithPartitions.values().stream()
                .flatMap(Collection::stream)
                .collect(Collectors.toSet());

            final long start = time.milliseconds();
            try {
                log.info("Revoking active tasks {}.", tasks);
                taskManager.handleRevocation(partitionsToRevoke);
            } finally {
                final long latency = time.milliseconds() - start;
                tasksRevokedSensor.record(latency);
                log.info("partition revocation took {} ms.", latency);
            }
            if (streamThread.state() != StreamThread.State.PENDING_SHUTDOWN) {
                streamThread.setState(StreamThread.State.PARTITIONS_REVOKED);
            }
        } catch (final Exception exception) {
            return Optional.of(exception);
        }
        return Optional.empty();
    }

    @Override
    public Optional<Exception> onTasksAssigned(final StreamsRebalanceData.Assignment assignment) {
        final long start = time.milliseconds();
        try {
            final Map<TaskId, Set<TopicPartition>> activeTasksWithPartitions =
                pairWithTopicPartitions(assignment.activeTasks().stream());
            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);

Review comment (bug):
The start time is captured outside the try block, but if an exception occurs during the method execution, the sensor will still record the latency in the finally block. This could lead to recording metrics for failed operations that didn't complete normally, potentially skewing the latency measurements.

Code after:
            .collect(Collectors.toSet());

        final long start = time.milliseconds();
        try {
            log.info("Revoking active tasks {}.", tasks);
            taskManager.handleRevocation(partitionsToRevoke);
        } finally {
            final long latency = time.milliseconds() - start;
            tasksRevokedSensor.record(latency);
            log.info("partition revocation took {} ms.", latency);
        }
        if (streamThread.state() != StreamThread.State.PENDING_SHUTDOWN) {
            streamThread.setState(StreamThread.State.PARTITIONS_REVOKED);
        }
    }

    @Override
    public void onTasksAssigned(final StreamsRebalanceData.Assignment assignment) {
        final long start = time.milliseconds();
        final Map<TaskId, Set<TopicPartition>> activeTasksWithPartitions =
            pairWithTopicPartitions(assignment.activeTasks().stream());
        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);

        try {
            taskManager.handleAssignment(activeTasksWithPartitions, standbyTasksWithPartitions);
            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,