Code review comment: Note:  This changes the blocking behavior of this call.  https://cwiki.apache.org/confluence/display/KAFKA/KIP-266%3A+Fix+consumer+indefinite+blocking+behavior  I'd like to see a proper ticket for this particular change.

Code under review:
public void run() {
        
        if (!running) {
            return;
        }

        
        final Handover handover = this.handover;

        
        
        
        try {
            this.consumer = getConsumer(kafkaProperties);
        } catch (Throwable t) {
            handover.reportError(t);
            return;
        }

        
        try {
            
            if (useMetrics) {
                
                Map<MetricName, ? extends Metric> metrics = consumer.metrics();
                if (metrics == null) {
                    
                    log.info("Consumer implementation does not support metrics");
                } else {
                    
                    for (Map.Entry<MetricName, ? extends Metric> metric : metrics.entrySet()) {
                        consumerMetricGroup.gauge(
                                metric.getKey().name(), new KafkaMetricWrapper(metric.getValue()));

                        
                        
                        subtaskMetricGroup.gauge(
                                metric.getKey().name(), new KafkaMetricWrapper(metric.getValue()));
                    }
                }
            }

            
            if (!running) {
                return;
            }

            
            
            ConsumerRecords<byte[], byte[]> records = null;

            
            
            
            List<KafkaTopicPartitionState<T, TopicPartition>> newPartitions;

            
            while (running) {

                
                if (!commitInProgress) {
                    
                    
                    final Tuple2<Map<TopicPartition, OffsetAndMetadata>, KafkaCommitCallback>
                            commitOffsetsAndCallback = nextOffsetsToCommit.getAndSet(null);

                    if (commitOffsetsAndCallback != null) {
                        log.debug("Sending async offset commit request to Kafka broker");

                        
                        
                        commitInProgress = true;
                        retryOnceOnWakeup(
                                () ->
                                        consumer.commitAsync(
                                                commitOffsetsAndCallback.f0,
                                                new CommitCallback(commitOffsetsAndCallback.f1)),
                                "commitAsync");
                    }
                }

                try {
                    if (hasAssignedPartitions) {
                        newPartitions = unassignedPartitionsQueue.pollBatch();
                    } else {
                        
                        
                        
                        
                        newPartitions = unassignedPartitionsQueue.getBatchBlocking();
                    }
                    if (newPartitions != null) {
                        reassignPartitions(newPartitions);
                    }
                } catch (AbortedReassignmentException e) {
                    continue;
                }

                if (!hasAssignedPartitions) {
                    
                    continue;
                }

                
                
                if (records == null) {
                    try {
                        records = consumer.poll(Duration.ofMillis(pollTimeout));
                    } catch (WakeupException we) {
                        continue;
                    }
                }

                try {
                    handover.produce(records);
                    records = null;
                } catch (Handover.WakeupException e) {
                    
                }
            }
            
        } catch (Throwable t) {
            
            
            
            handover.reportError(t);
        } finally {
            
            handover.close();

            
            try {
                consumer.close();
            } catch (Throwable t) {
                log.warn("Error while closing Kafka consumer", t);
            }
        }
    }

Line commented on: records = consumer.poll(Duration.ofMillis(pollTimeout));

Revised code:
public void run() {
        
        if (!running) {
            return;
        }

        
        final Handover handover = this.handover;

        
        
        
        try {
            this.consumer = getConsumer(kafkaProperties);
        } catch (Throwable t) {
            handover.reportError(t);
            return;
        }

        
        try {
            
            if (useMetrics) {
                
                Map<MetricName, ? extends Metric> metrics = consumer.metrics();
                if (metrics == null) {
                    
                    log.info("Consumer implementation does not support metrics");
                } else {
                    
                    for (Map.Entry<MetricName, ? extends Metric> metric : metrics.entrySet()) {
                        consumerMetricGroup.gauge(
                                metric.getKey().name(), new KafkaMetricWrapper(metric.getValue()));

                        
                        
                        subtaskMetricGroup.gauge(
                                metric.getKey().name(), new KafkaMetricWrapper(metric.getValue()));
                    }
                }
            }

            
            if (!running) {
                return;
            }

            
            
            ConsumerRecords<byte[], byte[]> records = null;

            
            
            
            List<KafkaTopicPartitionState<T, TopicPartition>> newPartitions;

            
            while (running) {

                
                if (!commitInProgress) {
                    
                    
                    final Tuple2<Map<TopicPartition, OffsetAndMetadata>, KafkaCommitCallback>
                            commitOffsetsAndCallback = nextOffsetsToCommit.getAndSet(null);

                    if (commitOffsetsAndCallback != null) {
                        log.debug("Sending async offset commit request to Kafka broker");

                        
                        
                        commitInProgress = true;
                        retryOnceOnWakeup(
                                () ->
                                        consumer.commitAsync(
                                                commitOffsetsAndCallback.f0,
                                                new CommitCallback(commitOffsetsAndCallback.f1)),
                                "commitAsync");
                    }
                }

                try {
                    if (hasAssignedPartitions) {
                        newPartitions = unassignedPartitionsQueue.pollBatch();
                    } else {
                        
                        
                        
                        
                        newPartitions = unassignedPartitionsQueue.getBatchBlocking();
                    }
                    if (newPartitions != null) {
                        reassignPartitions(newPartitions);
                    }
                } catch (AbortedReassignmentException e) {
                    continue;
                }

                if (!hasAssignedPartitions) {
                    
                    continue;
                }

                
                
                if (records == null) {
                    try {
                        records = consumer.poll(Duration.ofMillis(pollTimeout));
                    } catch (WakeupException we) {
                        continue;
                    }
                }

                try {
                    handover.produce(records);
                    records = null;
                } catch (Handover.WakeupException e) {
                    
                }
            }
            
        } catch (Throwable t) {
            
            
            
            handover.reportError(t);
        } finally {
            
            handover.close();

            
            try {
                consumer.close();
            } catch (Throwable t) {
                log.warn("Error while closing Kafka consumer", t);
            }
        }
    }