Repository: apache/kafka
File: clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java (Java)

Code before:
        } else {
            log.trace("Received produce response from node {} with correlation id {}", response.destination(), correlationId);
            // if we have a response, parse it
            if (response.hasResponse()) {
                // Sender should exercise PartitionProduceResponse rather than ProduceResponse.PartitionResponse
                // https://issues.apache.org/jira/browse/KAFKA-10696
                ProduceResponse produceResponse = (ProduceResponse) response.responseBody();
                // This will be set by completeBatch.
                Map<TopicPartition, Metadata.LeaderIdAndEpoch> partitionsWithUpdatedLeaderInfo = new HashMap<>();
                produceResponse.data().responses().forEach(r -> r.partitionResponses().forEach(p -> {
                    ProduceResponse.PartitionResponse partResp = new ProduceResponse.PartitionResponse(
                            Errors.forCode(p.errorCode()),
                            p.baseOffset(),
                            p.logAppendTimeMs(),
                            p.logStartOffset(),
                            p.recordErrors()
                                .stream()
                                .map(e -> new ProduceResponse.RecordError(e.batchIndex(), e.batchIndexErrorMessage()))
                                .collect(Collectors.toList()),
                            p.errorMessage(),
                            p.currentLeader());
                    // Version 13 drop topic name and add support to topic id.
                    // We need to find batch based on topic id and partition index only as
                    // topic name in the response might be empty.
                    TopicIdPartition tpId = new TopicIdPartition(r.topicId(), p.index(), r.name());
                    ProducerBatch batch = batches.entrySet().stream()
                            .filter(entry -> entry.getKey().same(tpId))
                            .map(Map.Entry::getValue).findFirst().orElse(null);

                    if (batch == null) {
                        throw new IllegalStateException("batch created for " + tpId + " can't be found, " +
                                "topic might be recreated after the batch creation.");
                    }
                    completeBatch(batch, partResp, correlationId, now, partitionsWithUpdatedLeaderInfo);
                }));

                if (!partitionsWithUpdatedLeaderInfo.isEmpty()) {
                    List<Node> leaderNodes = produceResponse.data().nodeEndpoints().stream()
                        .map(e -> new Node(e.nodeId(), e.host(), e.port(), e.rack()))
                        .filter(e -> !e.equals(Node.noNode()))
                        .collect(
                            Collectors.toList());
                    Set<TopicPartition> updatedPartitions = metadata.updatePartitionLeadership(partitionsWithUpdatedLeaderInfo, leaderNodes);
                    if (log.isTraceEnabled()) {
                        updatedPartitions.forEach(
                            part -> log.debug("For {} leader was updated.", part)
                        );
                    }
                }

                this.sensors.recordLatency(response.destination(), response.requestLatencyMs());

Review comment (performance):
This changes a map lookup to an iteration. Could we do some produce perf test (with multiple topic/partitions) to verify there is no performance degradation?

Code after:
            log.trace("Received produce response from node {} with correlation id {}", response.destination(), correlationId);
            // if we have a response, parse it
            if (response.hasResponse()) {
                // Sender should exercise PartitionProduceResponse rather than ProduceResponse.PartitionResponse
                // https://issues.apache.org/jira/browse/KAFKA-10696
                ProduceResponse produceResponse = (ProduceResponse) response.responseBody();
                // This will be set by completeBatch.
                Map<TopicPartition, Metadata.LeaderIdAndEpoch> partitionsWithUpdatedLeaderInfo = new HashMap<>();
                produceResponse.data().responses().forEach(r -> r.partitionResponses().forEach(p -> {
                    ProduceResponse.PartitionResponse partResp = new ProduceResponse.PartitionResponse(
                            Errors.forCode(p.errorCode()),
                            p.baseOffset(),
                            p.logAppendTimeMs(),
                            p.logStartOffset(),
                            p.recordErrors()
                                .stream()
                                .map(e -> new ProduceResponse.RecordError(e.batchIndex(), e.batchIndexErrorMessage()))
                                .collect(Collectors.toList()),
                            p.errorMessage(),
                            p.currentLeader());

                    // Version 13 drop topic name and add support to topic id.
                    // We need to find batch based on topic id and partition index only as
                    // topic name in the response will be empty.
                    // For older versions, topic id is zero, and we will find the batch based on the topic name.
                    TopicPartition tp = (!r.topicId().equals(Uuid.ZERO_UUID) && topicNames.containsKey(r.topicId())) ?
                            new TopicPartition(topicNames.get(r.topicId()), p.index()) :
                            new TopicPartition(r.name(), p.index());

                    ProducerBatch batch = batches.get(tp);
                    if (batch == null) {
                        throw new IllegalStateException("Can't find batch created for topic id " + r.topicId() +
                                " topic name " + r.name() + " partition " + p.index() + " using " + topicNames);
                    }
                    completeBatch(batch, partResp, correlationId, now, partitionsWithUpdatedLeaderInfo);
                }));

                if (!partitionsWithUpdatedLeaderInfo.isEmpty()) {
                    List<Node> leaderNodes = produceResponse.data().nodeEndpoints().stream()
                        .map(e -> new Node(e.nodeId(), e.host(), e.port(), e.rack()))
                        .filter(e -> !e.equals(Node.noNode()))
                        .collect(
                            Collectors.toList());
                    Set<TopicPartition> updatedPartitions = metadata.updatePartitionLeadership(partitionsWithUpdatedLeaderInfo, leaderNodes);
                    if (log.isTraceEnabled()) {
                        updatedPartitions.forEach(
                            part -> log.debug("For {} leader was updated.", part)
                        );
                    }
                }