Move setupDriverEndpointRef to RpcUtils and rename to makeDriverRef

Author: zsxwing

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

@@ -41,7 +41,7 @@ import org.apache.spark.scheduler.OutputCommitCoordinator.OutputCommitCoordinato
 import org.apache.spark.serializer.Serializer
 import org.apache.spark.shuffle.{ShuffleMemoryManager, ShuffleManager}
 import org.apache.spark.storage._
-import org.apache.spark.util.{AkkaUtils, Utils}
+import org.apache.spark.util.{AkkaUtils, RpcUtils, Utils}
 
 /**
  * :: DeveloperApi ::
@@ -300,7 +300,7 @@ object SparkEnv extends Logging {
         logInfo("Registering " + name)
         rpcEnv.setupEndpoint(name, endpointCreator)
       } else {
-        rpcEnv.setupDriverEndpointRef(name)
+        RpcUtils.makeDriverRef(name, conf, rpcEnv)
       }
     }
 

core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala

@@ -58,11 +58,6 @@ private[spark] trait RpcEnv {
    */
   def setupThreadSafeEndpoint(name: String, endpoint: RpcEndpoint): RpcEndpointRef
 
-  /**
-   * Retrieve a [[RpcEndpointRef]] which is located in the driver via its name.
-   */
-  def setupDriverEndpointRef(name: String): RpcEndpointRef
-
   /**
    * Retrieve the [[RpcEndpointRef]] represented by `url`.
    */

core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala

@@ -210,12 +210,8 @@ private[spark] class AkkaRpcEnv private (
       address.port.getOrElse(defaultAddress.port))
   }
 
-  override def setupDriverEndpointRef(name: String): RpcEndpointRef = {
-    new AkkaRpcEndpointRef(defaultAddress, AkkaUtils.makeDriverRef(name, conf, actorSystem), conf)
-  }
-
   override def setupEndpointRefByUrl(url: String): RpcEndpointRef = {
-    val timeout = Duration.create(conf.getLong("spark.akka.lookupTimeout", 30), "seconds")
+    val timeout = AkkaUtils.lookupTimeout(conf)
     val ref = Await.result(actorSystem.actorSelection(url).resolveOne(timeout), timeout)
     // TODO defaultAddress is wrong
     new AkkaRpcEndpointRef(defaultAddress, ref, conf)

core/src/main/scala/org/apache/spark/util/RpcUtils.scala

@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.util
+
+import org.apache.spark.{SparkEnv, SparkConf}
+import org.apache.spark.rpc.{RpcAddress, RpcEndpointRef, RpcEnv}
+
+object RpcUtils {
+
+  /**
+   * Retrieve a [[RpcEndpointRef]] which is located in the driver via its name.
+   */
+  def makeDriverRef(name: String, conf: SparkConf, rpcEnv: RpcEnv): RpcEndpointRef = {
+    val driverActorSystemName = SparkEnv.driverActorSystemName
+    val driverHost: String = conf.get("spark.driver.host", "localhost")
+    val driverPort: Int = conf.getInt("spark.driver.port", 7077)
+    Utils.checkHost(driverHost, "Expected hostname")
+    rpcEnv.setupEndpointRef(driverActorSystemName, RpcAddress(driverHost, driverPort), name)
+  }
+}