Fix checkpointing doesn't retain driver port issue

Author: jerryshao

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -94,6 +94,11 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
   // contains a map from hostname to a list of input format splits on the host.
   private[spark] var preferredNodeLocationData: Map[String, Set[SplitInfo]] = Map()
 
+  // This is used for Spark Streaming to check whether driver host and port are set by user,
+  // if these two configurations are set by user, so the recovery mechanism should not remove this.
+  private[spark] val isDriverHostSetByUser = config.contains("spark.driver.host")
+  private[spark] val isDriverPortSetByUser = config.contains("spark.driver.port")
+
   val startTime = System.currentTimeMillis()
 
   private val stopped: AtomicBoolean = new AtomicBoolean(false)

streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala

@@ -41,7 +41,9 @@ class Checkpoint(@transient ssc: StreamingContext, val checkpointTime: Time)
   val checkpointDuration = ssc.checkpointDuration
   val pendingTimes = ssc.scheduler.getPendingTimes().toArray
   val delaySeconds = MetadataCleaner.getDelaySeconds(ssc.conf)
-  val sparkConfPairs = ssc.conf.getAll
+  val sparkConfPairs = ssc.conf.getAll.filterNot { kv =>
+    (!ssc.sc.isDriverHostSetByUser && kv._1 == "spark.driver.host") ||
+      (!ssc.sc.isDriverPortSetByUser && kv._1 == "spark.driver.port")
 
   def createSparkConf(): SparkConf = {
 
@@ -50,7 +52,9 @@ class Checkpoint(@transient ssc: StreamingContext, val checkpointTime: Time)
     val propertiesToReload = List(
       "spark.master",
       "spark.yarn.keytab",
-      "spark.yarn.principal")
+      "spark.yarn.principal",
+      "spark.driver.host",
+      "spark.driver.port")
 
     val newSparkConf = new SparkConf(loadDefaults = false).setAll(sparkConfPairs)
       .remove("spark.driver.host")

streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala

@@ -191,6 +191,38 @@ class CheckpointSuite extends TestSuiteBase {
     }
   }
 
+  // This tests if "spark.driver.host" and "spark.driver.port" is set by user, can be recovered
+  // with correct value.
+  test("correctly recover spark.driver.[host|port] from checkpoint") {
+    val conf = Map("spark.driver.host" -> "localhost", "spark.driver.port" -> "9999")
+    conf.foreach(kv => System.setProperty(kv._1, kv._2))
+    ssc = new StreamingContext(master, framework, batchDuration)
+    val originalConf = ssc.conf
+    assert(ssc.sc.isDriverHostSetByUser === true)
+    assert(ssc.sc.isDriverPortSetByUser === true)
+    assert(originalConf.get("spark.driver.host") === "localhost")
+    assert(originalConf.get("spark.driver.port") === "9999")
+
+    val cp = new Checkpoint(ssc, Time(1000))
+    ssc.stop()
+
+    // Serialize/deserialize to simulate write to storage and reading it back
+    val newCp = Utils.deserialize[Checkpoint](Utils.serialize(cp))
+
+    val newCpConf = newCp.sparkConf
+    assert(newCpConf.contains("spark.driver.host"))
+    assert(newCpConf.contains("spark.driver.port"))
+    assert(originalConf.get("spark.driver.host") === "localhost")
+    assert(originalConf.get("spark.driver.port") === "9999")
+
+    // Check if all the parameters have been restored
+    ssc = new StreamingContext(null, newCp, null)
+    val restoredConf = ssc.conf
+    assert(ssc.sc.isDriverHostSetByUser === true)
+    assert(ssc.sc.isDriverPortSetByUser === true)
+    assert(restoredConf.get("spark.driver.host") === "localhost")
+    assert(restoredConf.get("spark.driver.port") === "9999")
+  }
 
   // This tests whether the systm can recover from a master failure with simple
   // non-stateful operations. This assumes as reliable, replayable input