MapR [SPARK-482] Spark streaming app fails to start by UnknownTopicOrPartitionException with checkpoint (apache#453)

Author: ekrivokonmapr

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

@@ -105,7 +105,8 @@ private case class Subscribe[K, V](
       val shouldSuppress =
         aor != null && aor.asInstanceOf[String].toUpperCase(Locale.ROOT) == "NONE"
       try {
-        consumer.poll(0)
+        // consumer.poll(0)
+        KafkaUtils.waitForConsumerAssignment(consumer, toSeek.keySet())
       } catch {
         case x: NoOffsetForPartitionException if shouldSuppress =>
           logWarning("Catching NoOffsetForPartitionException since " +
@@ -158,7 +159,8 @@ private case class SubscribePattern[K, V](
       val shouldSuppress =
         aor != null && aor.asInstanceOf[String].toUpperCase(Locale.ROOT) == "NONE"
       try {
-        consumer.poll(0)
+        // consumer.poll(0)
+        KafkaUtils.waitForConsumerAssignment(consumer, toSeek.keySet())
       } catch {
         case x: NoOffsetForPartitionException if shouldSuppress =>
           logWarning("Catching NoOffsetForPartitionException since " +

external/kafka-0-10/src/main/scala/org/apache/spark/streaming/kafka010/KafkaUtils.scala

@@ -25,7 +25,7 @@ import net.razorvine.pickle.{IObjectPickler, Opcodes, Pickler}
 import org.apache.kafka.clients.consumer._
 import org.apache.kafka.common.TopicPartition
 
-import org.apache.spark.SparkContext
+import org.apache.spark.{SparkContext, SparkEnv}
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.{JavaRDD, JavaSparkContext}
 import org.apache.spark.api.python.SerDeUtil
@@ -230,6 +230,19 @@ object KafkaUtils extends Logging {
     }
   }
 
+  def waitForConsumerAssignment[K, V](consumer: KafkaConsumer[K, V],
+                                      partitions: ju.Set[TopicPartition]): Unit = {
+    val waitingForAssigmentTimeout = SparkEnv.get.conf.
+      getLong("spark.mapr.WaitingForAssignmentTimeout", 600000)
+
+    var timeout = 0
+    while ((consumer.assignment().isEmpty || consumer.assignment().size() < partitions.size)
+      && timeout < waitingForAssigmentTimeout) {
+
+      Thread.sleep(500)
+      timeout += 500
+    }
+  }
 }
 
 object KafkaUtilsPythonHelper {

external/kafka-0-9/src/main/scala/org/apache/spark/streaming/kafka09/ConsumerStrategy.scala

@@ -106,7 +106,8 @@ private case class Subscribe[K, V](
       val shouldSuppress =
         aor != null && aor.asInstanceOf[String].toUpperCase(Locale.ROOT) == "NONE"
       try {
-        consumer.poll(0)
+        // consumer.poll(0)
+        KafkaUtils.waitForConsumerAssignment(consumer, toSeek.keySet())
       } catch {
         case x: NoOffsetForPartitionException if shouldSuppress =>
           logWarning("Catching NoOffsetForPartitionException since " +
@@ -159,7 +160,8 @@ private case class SubscribePattern[K, V](
       val shouldSuppress =
         aor != null && aor.asInstanceOf[String].toUpperCase(Locale.ROOT) == "NONE"
       try {
-        consumer.poll(0)
+        // consumer.poll(0)
+        KafkaUtils.waitForConsumerAssignment(consumer, toSeek.keySet())
       } catch {
         case x: NoOffsetForPartitionException if shouldSuppress =>
           logWarning("Catching NoOffsetForPartitionException since " +

external/kafka-0-9/src/main/scala/org/apache/spark/streaming/kafka09/KafkaUtils.scala

@@ -25,7 +25,7 @@ import net.razorvine.pickle.{IObjectPickler, Opcodes, Pickler}
 import org.apache.kafka.clients.consumer._
 import org.apache.kafka.common.TopicPartition
 
-import org.apache.spark.SparkContext
+import org.apache.spark.{SparkContext, SparkEnv}
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.{JavaRDD, JavaSparkContext}
 import org.apache.spark.api.python.SerDeUtil
@@ -231,6 +231,19 @@ object KafkaUtils extends Logging {
     }
   }
 
+  def waitForConsumerAssignment[K, V](consumer: KafkaConsumer[K, V],
+                                      partitions: ju.Set[TopicPartition]): Unit = {
+    val waitingForAssigmentTimeout = SparkEnv.get.conf.
+      getLong("spark.mapr.WaitingForAssignmentTimeout", 600000)
+
+    var timeout = 0
+    while ((consumer.assignment().isEmpty || consumer.assignment().size() < partitions.size)
+      && timeout < waitingForAssigmentTimeout) {
+
+      Thread.sleep(500)
+      timeout += 500
+    }
+  }
 }
 
 @deprecated("Use kafka10 package instead of kafka09", "MapR Spark-2.3.2")