[SPARK-6028][Core]A new RPC implemetation based on the network module

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

@@ -44,7 +44,7 @@ private[spark] class MapOutputTrackerMasterEndpoint(
 
   override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
     case GetMapOutputStatuses(shuffleId: Int) =>
-      val hostPort = context.sender.address.hostPort
+      val hostPort = context.senderAddress.hostPort
       logInfo("Asked to send map output locations for shuffle " + shuffleId + " to " + hostPort)
       val mapOutputStatuses = tracker.getSerializedMapOutputStatuses(shuffleId)
       val serializedSize = mapOutputStatuses.size

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

@@ -20,11 +20,10 @@ package org.apache.spark
 import java.io.File
 import java.net.Socket
 
-import akka.actor.ActorSystem
-
 import scala.collection.mutable
 import scala.util.Properties
 
+import akka.actor.ActorSystem
 import com.google.common.collect.MapMaker
 
 import org.apache.spark.annotation.DeveloperApi
@@ -42,7 +41,7 @@ import org.apache.spark.serializer.Serializer
 import org.apache.spark.shuffle.{ShuffleMemoryManager, ShuffleManager}
 import org.apache.spark.storage._
 import org.apache.spark.unsafe.memory.{ExecutorMemoryManager, MemoryAllocator}
-import org.apache.spark.util.{RpcUtils, Utils}
+import org.apache.spark.util.{AkkaUtils, RpcUtils, Utils}
 
 /**
  * :: DeveloperApi ::
@@ -58,6 +57,7 @@ import org.apache.spark.util.{RpcUtils, Utils}
 class SparkEnv (
     val executorId: String,
     private[spark] val rpcEnv: RpcEnv,
+    _actorSystem: ActorSystem, // TODO Remove actorSystem
     val serializer: Serializer,
     val closureSerializer: Serializer,
     val cacheManager: CacheManager,
@@ -77,7 +77,7 @@ class SparkEnv (
 
   // TODO Remove actorSystem
   @deprecated("Actor system is no longer supported as of 1.4.0", "1.4.0")
-  val actorSystem: ActorSystem = rpcEnv.asInstanceOf[AkkaRpcEnv].actorSystem
+  val actorSystem: ActorSystem = _actorSystem
 
   private[spark] var isStopped = false
   private val pythonWorkers = mutable.HashMap[(String, Map[String, String]), PythonWorkerFactory]()
@@ -101,6 +101,9 @@ class SparkEnv (
       blockManager.master.stop()
       metricsSystem.stop()
       outputCommitCoordinator.stop()
+      if (!rpcEnv.isInstanceOf[AkkaRpcEnv]) {
+        actorSystem.shutdown()
+      }
       rpcEnv.shutdown()
 
       // Unfortunately Akka's awaitTermination doesn't actually wait for the Netty server to shut
@@ -250,7 +253,13 @@ object SparkEnv extends Logging {
     // Create the ActorSystem for Akka and get the port it binds to.
     val actorSystemName = if (isDriver) driverActorSystemName else executorActorSystemName
     val rpcEnv = RpcEnv.create(actorSystemName, hostname, port, conf, securityManager)
-    val actorSystem = rpcEnv.asInstanceOf[AkkaRpcEnv].actorSystem
+    val actorSystem: ActorSystem =
+      if (rpcEnv.isInstanceOf[AkkaRpcEnv]) {
+        rpcEnv.asInstanceOf[AkkaRpcEnv].actorSystem
+      } else {
+        // Create a ActorSystem for legacy codes
+        AkkaUtils.createActorSystem(actorSystemName, hostname, port, conf, securityManager)._1
+      }
 
     // Figure out which port Akka actually bound to in case the original port is 0 or occupied.
     if (isDriver) {
@@ -402,6 +411,7 @@ object SparkEnv extends Logging {
     val envInstance = new SparkEnv(
       executorId,
       rpcEnv,
+      actorSystem,
       serializer,
       closureSerializer,
       cacheManager,

core/src/main/scala/org/apache/spark/deploy/worker/WorkerWatcher.scala

@@ -27,11 +27,9 @@ import org.apache.spark.rpc._
 private[spark] class WorkerWatcher(override val rpcEnv: RpcEnv, workerUrl: String)
   extends RpcEndpoint with Logging {
 
-  override def onStart() {
-    logInfo(s"Connecting to worker $workerUrl")
-    if (!isTesting) {
-      rpcEnv.asyncSetupEndpointRefByURI(workerUrl)
-    }
+  logInfo(s"Connecting to worker $workerUrl")
+  if (!isTesting) {
+    rpcEnv.asyncSetupEndpointRefByURI(workerUrl)
   }
 
   // Used to avoid shutting down JVM during tests

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

@@ -37,5 +37,5 @@ private[spark] trait RpcCallContext {
   /**
    * The sender of this message.
    */
-  def sender: RpcEndpointRef
+  def senderAddress: RpcAddress
 }

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

@@ -39,7 +39,21 @@ private[spark] trait RpcEnvFactory {
  * However, there is no guarantee that the same thread will be executing the same
  * [[ThreadSafeRpcEndpoint]] for different messages.
  */
-private[spark] trait ThreadSafeRpcEndpoint extends RpcEndpoint
+private[spark] trait ThreadSafeRpcEndpoint extends RpcEndpoint {
+  /**
+   * Invoked before [[RpcEndpoint]] starts to handle any message.
+   */
+  def onStart(): Unit = {
+    // By default, do nothing.
+  }
+
+  /**
+   * Invoked when [[RpcEndpoint]] is stopping.
+   */
+  def onStop(): Unit = {
+    // By default, do nothing.
+  }
+}
 
 
 /**
@@ -100,20 +114,6 @@ private[spark] trait RpcEndpoint {
     throw cause
   }
 
-  /**
-   * Invoked before [[RpcEndpoint]] starts to handle any message.
-   */
-  def onStart(): Unit = {
-    // By default, do nothing.
-  }
-
-  /**
-   * Invoked when [[RpcEndpoint]] is stopping.
-   */
-  def onStop(): Unit = {
-    // By default, do nothing.
-  }
-
   /**
    * Invoked when `remoteAddress` is connected to the current node.
    */

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

@@ -0,0 +1,22 @@
+/*
+ * 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.rpc
+
+import org.apache.spark.SparkException
+
+private[rpc] class RpcEndpointNotFoundException(uri: String)
+  extends SparkException(s"Cannot find endpoint: $uri")

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

@@ -36,8 +36,11 @@ private[spark] object RpcEnv {
 
   private def getRpcEnvFactory(conf: SparkConf): RpcEnvFactory = {
     // Add more RpcEnv implementations here
-    val rpcEnvNames = Map("akka" -> "org.apache.spark.rpc.akka.AkkaRpcEnvFactory")
-    val rpcEnvName = conf.get("spark.rpc", "akka")
+    val rpcEnvNames = Map("akka" -> "org.apache.spark.rpc.akka.AkkaRpcEnvFactory",
+      "netty" -> "org.apache.spark.rpc.netty.NettyRpcEnvFactory")
+    // Use "netty" by default so that Jenkins can run all tests using NettyRpcEnv.
+    // Will change it back to "akka" before merging the new implementation.
+    val rpcEnvName = conf.get("spark.rpc", "netty")
     val rpcEnvFactoryClassName = rpcEnvNames.getOrElse(rpcEnvName.toLowerCase, rpcEnvName)
     Utils.classForName(rpcEnvFactoryClassName).newInstance().asInstanceOf[RpcEnvFactory]
   }

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

@@ -97,7 +97,9 @@ private[spark] class AkkaRpcEnv private[akka] (
         // Listen for remote client network events
         context.system.eventStream.subscribe(self, classOf[AssociationEvent])
         safelyCall(endpoint) {
-          endpoint.onStart()
+          if (endpoint.isInstanceOf[ThreadSafeRpcEndpoint]) {
+            endpoint.asInstanceOf[ThreadSafeRpcEndpoint].onStart()
+          }
         }
       }
 
@@ -140,7 +142,9 @@ private[spark] class AkkaRpcEnv private[akka] (
       override def postStop(): Unit = {
         unregisterEndpoint(endpoint.self)
         safelyCall(endpoint) {
-          endpoint.onStop()
+          if (endpoint.isInstanceOf[ThreadSafeRpcEndpoint]) {
+            endpoint.asInstanceOf[ThreadSafeRpcEndpoint].onStop()
+          }
         }
       }
 
@@ -166,9 +170,9 @@ private[spark] class AkkaRpcEnv private[akka] (
             _sender ! AkkaMessage(response, false)
           }
 
-          // Some RpcEndpoints need to know the sender's address
-          override val sender: RpcEndpointRef =
-            new AkkaRpcEndpointRef(defaultAddress, _sender, conf)
+          // Use "lazy" because most of RpcEndpoints don't need "senderAddress"
+          override lazy val senderAddress: RpcAddress =
+            new AkkaRpcEndpointRef(defaultAddress, _sender, conf).address
         })
       } else {
         endpoint.receive