[SPARK-53904][SS] Use Consumer.poll(Duration) instead of Consumer.poll(long)

connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.kafka010
 
 import java.{util => ju}
+import java.time.Duration
 
 import scala.collection.mutable.ArrayBuffer
 import scala.jdk.CollectionConverters._
@@ -135,7 +136,7 @@ private[kafka010] class KafkaOffsetReaderConsumer(
     uninterruptibleThreadRunner.runUninterruptibly {
       assert(Thread.currentThread().isInstanceOf[UninterruptibleThread])
       // Poll to get the latest assigned partitions
-      consumer.poll(0)
+      consumer.poll(Duration.ofMillis(500))
       val partitions = consumer.assignment()
       consumer.pause(partitions)
       partitions.asScala.toSet
@@ -579,7 +580,7 @@ private[kafka010] class KafkaOffsetReaderConsumer(
 
     withRetriesWithoutInterrupt {
       // Poll to get the latest assigned partitions
-      consumer.poll(0)
+      consumer.poll(Duration.ofMillis(500))
       val partitions = consumer.assignment()
 
       if (!fetchingEarliestOffset) {

connector/kafka-0-10/src/main/scala/org/apache/spark/streaming/kafka010/ConsumerStrategy.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.streaming.kafka010
 
 import java.{lang => jl, util => ju}
+import java.time.Duration
 import java.util.Locale
 
 import scala.jdk.CollectionConverters._
@@ -105,7 +106,7 @@ private case class Subscribe[K, V](
       val shouldSuppress =
         aor != null && aor.asInstanceOf[String].toUpperCase(Locale.ROOT) == "NONE"
       try {
-        consumer.poll(0)
+        consumer.poll(Duration.ZERO)
       } catch {
         case x: NoOffsetForPartitionException if shouldSuppress =>
           logWarning(log"Catching NoOffsetForPartitionException since " +
@@ -159,7 +160,7 @@ private case class SubscribePattern[K, V](
       val shouldSuppress =
         aor != null && aor.asInstanceOf[String].toUpperCase(Locale.ROOT) == "NONE"
       try {
-        consumer.poll(0)
+        consumer.poll(Duration.ZERO)
       } catch {
         case x: NoOffsetForPartitionException if shouldSuppress =>
           logWarning(log"Catching NoOffsetForPartitionException since " +

connector/kafka-0-10/src/main/scala/org/apache/spark/streaming/kafka010/DirectKafkaInputDStream.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.streaming.kafka010
 
 import java.{ util => ju }
+import java.time.Duration
 import java.util.concurrent.ConcurrentLinkedQueue
 import java.util.concurrent.atomic.AtomicReference
 
@@ -170,15 +171,15 @@ private[spark] class DirectKafkaInputDStream[K, V](
   private def paranoidPoll(c: Consumer[K, V]): Unit = {
     // don't actually want to consume any messages, so pause all partitions
     c.pause(c.assignment())
-    val msgs = c.poll(0)
+    val msgs = c.poll(Duration.ZERO)
     if (!msgs.isEmpty) {
       // position should be minimum offset per topicpartition
       msgs.asScala.foldLeft(Map[TopicPartition, Long]()) { (acc, m) =>
         val tp = new TopicPartition(m.topic, m.partition)
         val off = acc.get(tp).map(o => Math.min(o, m.offset)).getOrElse(m.offset)
         acc + (tp -> off)
       }.foreach { case (tp, off) =>
-          logInfo(log"poll(0) returned messages, seeking ${MDC(TOPIC_PARTITION, tp)} to " +
+          logInfo(log"poll(ZERO) returned messages, seeking ${MDC(TOPIC_PARTITION, tp)} to " +
             log"${MDC(OFFSET, off)} to compensate")
           c.seek(tp, off)
       }

connector/kafka-0-10/src/test/scala/org/apache/spark/streaming/kafka010/DirectKafkaStreamSuite.scala

@@ -19,6 +19,7 @@ package org.apache.spark.streaming.kafka010
 
 import java.io.File
 import java.lang.{Long => JLong}
+import java.time.Duration
 import java.util.{Arrays, HashMap => JHashMap, Map => JMap, UUID}
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.ConcurrentLinkedQueue
@@ -255,7 +256,7 @@ class DirectKafkaStreamSuite
         preferredHosts,
         ConsumerStrategies.Subscribe[String, String](List(topic), kafkaParams.asScala),
         new DefaultPerPartitionConfig(sparkConf))
-      s.consumer().poll(0)
+      s.consumer().poll(Duration.ofMillis(500))
       assert(
         s.consumer().position(topicPartition) >= offsetBeforeStart,
         "Start offset not from latest"
@@ -311,7 +312,7 @@ class DirectKafkaStreamSuite
           kafkaParams.asScala,
           Map(topicPartition -> 11L)),
         new DefaultPerPartitionConfig(sparkConf))
-      s.consumer().poll(0)
+      s.consumer().poll(Duration.ZERO)
       assert(
         s.consumer().position(topicPartition) >= offsetBeforeStart,
         "Start offset not from latest"
@@ -473,7 +474,7 @@ class DirectKafkaStreamSuite
     ssc.stop()
     val consumer = new KafkaConsumer[String, String](kafkaParams)
     consumer.subscribe(Arrays.asList(topic))
-    consumer.poll(0)
+    consumer.poll(Duration.ofMillis(500))
     committed.asScala.foreach {
       case (k, v) =>
         // commits are async, not exactly once