Add RpcEnvFactory to create RpcEnv

Author: zsxwing

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

@@ -111,27 +111,18 @@ private[spark] case class RpcEnvConfig(
     securityManager: SecurityManager)
 
 /**
- * A RpcEnv implementation must have a companion object with an
- * `apply(config: RpcEnvConfig): RpcEnv` method so that it can be created via Reflection.
- *
- * {{{
- * object MyCustomRpcEnv {
- *   def apply(config: RpcEnvConfig): RpcEnv = {
- *     ...
- *   }
- * }
- * }}}
+ * A RpcEnv implementation must have a [[RpcEnvFactory]] implementation with an empty constructor
+ * so that it can be created via Reflection.
  */
 private[spark] object RpcEnv {
 
-  private def getRpcEnvCompanion(conf: SparkConf): AnyRef = {
+  private def getRpcEnvFactory(conf: SparkConf): RpcEnvFactory = {
     // Add more RpcEnv implementations here
-    val rpcEnvNames = Map("akka" -> "org.apache.spark.rpc.akka.AkkaRpcEnv")
+    val rpcEnvNames = Map("akka" -> "org.apache.spark.rpc.akka.AkkaRpcEnvFactory")
     val rpcEnvName = conf.get("spark.rpc", "akka")
-    val rpcEnvClassName = rpcEnvNames.getOrElse(rpcEnvName.toLowerCase, rpcEnvName)
-    val companion = Class.forName(
-      rpcEnvClassName + "$", true, Utils.getContextOrSparkClassLoader).getField("MODULE$").get(null)
-    companion
+    val rpcEnvFactoryClassName = rpcEnvNames.getOrElse(rpcEnvName.toLowerCase, rpcEnvName)
+    Class.forName(rpcEnvFactoryClassName, true, Utils.getContextOrSparkClassLoader).
+      newInstance().asInstanceOf[RpcEnvFactory]
   }
 
   def create(
@@ -142,13 +133,20 @@ private[spark] object RpcEnv {
       securityManager: SecurityManager): RpcEnv = {
     // Using Reflection to create the RpcEnv to avoid to depend on Akka directly
     val config = RpcEnvConfig(conf, name, host, port, securityManager)
-    val companion = getRpcEnvCompanion(conf)
-    companion.getClass.getMethod("apply", classOf[RpcEnvConfig]).
-      invoke(companion, config).asInstanceOf[RpcEnv]
+    getRpcEnvFactory(conf).create(config)
   }
 
 }
 
+/**
+ * A factory class to create the [[RpcEnv]]. It must have an empty constructor so that it can be
+ * created using Reflection.
+ */
+private[spark] trait RpcEnvFactory {
+
+  def create(config: RpcEnvConfig): RpcEnv
+}
+
 /**
  * An end point for the RPC that defines what functions to trigger given a message.
  *

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

@@ -43,7 +43,7 @@ import org.apache.spark.util.{ActorLogReceive, AkkaUtils}
  * @param conf
  * @param boundPort
  */
-private[spark] class AkkaRpcEnv private (
+private[spark] class AkkaRpcEnv private[akka] (
     val actorSystem: ActorSystem, conf: SparkConf, boundPort: Int) extends RpcEnv with Logging {
 
   private val defaultAddress: RpcAddress = {
@@ -250,14 +250,13 @@ private[spark] class AkkaRpcEnv private (
   override def toString = s"${getClass.getSimpleName}($actorSystem)"
 }
 
-private[spark] object AkkaRpcEnv {
+private[spark] class AkkaRpcEnvFactory extends RpcEnvFactory {
 
-  def apply(config: RpcEnvConfig): RpcEnv = {
+  def create(config: RpcEnvConfig): RpcEnv = {
     val (actorSystem, boundPort) = AkkaUtils.createActorSystem(
-        config.name, config.host, config.port, config.conf, config.securityManager)
+      config.name, config.host, config.port, config.conf, config.securityManager)
     new AkkaRpcEnv(actorSystem, config.conf, boundPort)
   }
-
 }
 
 private[akka] class AkkaRpcEndpointRef(

core/src/test/scala/org/apache/spark/rpc/akka/AkkaRpcEnvSuite.scala

@@ -23,7 +23,8 @@ import org.apache.spark.{SecurityManager, SparkConf}
 class AkkaRpcEnvSuite extends RpcEnvSuite {
 
   override def createRpcEnv(conf: SparkConf, name: String, port: Int): RpcEnv = {
-    AkkaRpcEnv(RpcEnvConfig(conf, name, "localhost", port, new SecurityManager(conf)))
+    new AkkaRpcEnvFactory().create(
+      RpcEnvConfig(conf, name, "localhost", port, new SecurityManager(conf)))
   }
 
   test("setupEndpointRef: systemName, address, endpointName") {
@@ -35,8 +36,8 @@ class AkkaRpcEnvSuite extends RpcEnvSuite {
       }
     })
     val conf = new SparkConf()
-    val newRpcEnv =
-      AkkaRpcEnv(RpcEnvConfig(conf, "test", "localhost", 12346, new SecurityManager(conf)))
+    val newRpcEnv = new AkkaRpcEnvFactory().create(
+      RpcEnvConfig(conf, "test", "localhost", 12346, new SecurityManager(conf)))
     try {
       val newRef = newRpcEnv.setupEndpointRef("local", ref.address, "test_endpoint")
       assert("akka.tcp://local@localhost:12345/user/test_endpoint" ===