Add 'sender' to RpcCallContext and rename the parameter of receiveAndReply to 'context'

Author: zsxwing

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

@@ -191,8 +191,8 @@ private[spark] trait RpcEndpoint {
   /**
    * Process messages from [[RpcEndpointRef.sendWithReply]] or [[RpcCallContext.replyWithSender)]]
    */
-  def receiveAndReply(response: RpcCallContext): PartialFunction[Any, Unit] = {
-    case _ => response.sendFailure(new SparkException(self + " won't reply anything"))
+  def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
+    case _ => context.sendFailure(new SparkException(self + " won't reply anything"))
   }
 
   /**
@@ -360,19 +360,12 @@ private[spark] trait RpcCallContext {
   def reply(response: Any): Unit
 
   /**
-   * Reply a message to the corresponding [[RpcEndpoint.receiveAndReply]]. If you use this one to
-   * reply, it means you expect the target [[RpcEndpoint]] should reply you something.
-   *
-   * TODO better method name?
-   *
-   * @param response the response message
-   * @param sender who replies this message. The target [[RpcEndpoint]] will use `sender` to send
-   *               back something.
+   * Report a failure to the sender.
    */
-  def replyWithSender(response: Any, sender: RpcEndpointRef): Unit
+  def sendFailure(e: Throwable): Unit
 
   /**
-   * Report a failure to the sender.
+   * The sender of this message.
    */
-  def sendFailure(e: Throwable): Unit
+  def sender: RpcEndpointRef
 }

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

@@ -146,10 +146,9 @@ private[spark] class AkkaRpcEnv private (
                       s ! AkkaMessage(response, false)
                     }
 
-                    override def replyWithSender(response: Any, sender: RpcEndpointRef): Unit = {
-                      s.!(AkkaMessage(response, true))(
-                        sender.asInstanceOf[AkkaRpcEndpointRef].actorRef)
-                    }
+                    // Some RpcEndpoints need to know the sender's address
+                    override val sender: RpcEndpointRef =
+                      new AkkaRpcEndpointRef(defaultAddress, s, conf)
                   })
                 } else {
                   endpoint.receive

core/src/main/scala/org/apache/spark/scheduler/OutputCommitCoordinator.scala

@@ -156,13 +156,13 @@ private[spark] object OutputCommitCoordinator {
       override val rpcEnv: RpcEnv, outputCommitCoordinator: OutputCommitCoordinator)
     extends RpcEndpoint with Logging {
 
-    override def receiveAndReply(response: RpcCallContext) = {
+    override def receiveAndReply(context: RpcCallContext) = {
       case AskPermissionToCommitOutput(stage, partition, taskAttempt) =>
-        response.reply(
+        context.reply(
           outputCommitCoordinator.handleAskPermissionToCommit(stage, partition, taskAttempt))
       case StopCoordinator =>
         logInfo("OutputCommitCoordinator stopped!")
-        response.reply(true)
+        context.reply(true)
         stop()
     }
   }

core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala

@@ -93,9 +93,9 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
     val endpoint = new RpcEndpoint {
       override val rpcEnv = env
 
-      override def receiveAndReply(response: RpcCallContext) = {
-        case "Hello" => response.reply(self)
-        case "Echo" => response.reply("Echo")
+      override def receiveAndReply(context: RpcCallContext) = {
+        case "Hello" => context.reply(self)
+        case "Echo" => context.reply("Echo")
       }
     }
     val rpcEndpointRef = env.setupEndpoint("send-ref", endpoint)
@@ -109,9 +109,9 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
     val rpcEndpointRef = env.setupEndpoint("ask-locally", new RpcEndpoint {
       override val rpcEnv = env
 
-      override def receiveAndReply(response: RpcCallContext) = {
+      override def receiveAndReply(context: RpcCallContext) = {
         case msg: String => {
-          response.reply(msg)
+          context.reply(msg)
         }
       }
     })
@@ -123,9 +123,9 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
     env.setupEndpoint("ask-remotely", new RpcEndpoint {
       override val rpcEnv = env
 
-      override def receiveAndReply(response: RpcCallContext) = {
+      override def receiveAndReply(context: RpcCallContext) = {
         case msg: String => {
-          response.reply(msg)
+          context.reply(msg)
         }
       }
     })
@@ -146,10 +146,10 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
     env.setupEndpoint("ask-timeout", new RpcEndpoint {
       override val rpcEnv = env
 
-      override def receiveAndReply(response: RpcCallContext) = {
+      override def receiveAndReply(context: RpcCallContext) = {
         case msg: String => {
           Thread.sleep(100)
-          response.reply(msg)
+          context.reply(msg)
         }
       }
     })
@@ -175,8 +175,8 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
     val pongRef = env.setupEndpoint("pong", new RpcEndpoint {
       override val rpcEnv = env
 
-      override def receiveAndReply(response: RpcCallContext) = {
-        case Ping(id) => response.replyWithSender(Pong(id), self)
+      override def receiveAndReply(context: RpcCallContext) = {
+        case Ping(id) => context.sender.sendWithReply(Pong(id), self)
       }
     })
 
@@ -185,14 +185,14 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
 
       var requester: RpcCallContext = _
 
-      override def receiveAndReply(response: RpcCallContext) = {
+      override def receiveAndReply(context: RpcCallContext) = {
         case Start => {
-          requester = response
+          requester = context
           pongRef.sendWithReply(Ping(1), self)
         }
         case p @ Pong(id) => {
           if (id < 10) {
-            response.replyWithSender(Ping(id + 1), self)
+            context.sender.sendWithReply(Ping(id + 1), self)
           } else {
             requester.reply(p)
           }
@@ -430,8 +430,8 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
     val endpointRef = env.setupEndpoint("sendWithReply", new RpcEndpoint {
       override val rpcEnv = env
 
-      override def receiveAndReply(response: RpcCallContext) = {
-        case m => response.reply("ack")
+      override def receiveAndReply(context: RpcCallContext) = {
+        case m => context.reply("ack")
       }
     })
 
@@ -446,8 +446,8 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
     env.setupEndpoint("sendWithReply-remotely", new RpcEndpoint {
       override val rpcEnv = env
 
-      override def receiveAndReply(response: RpcCallContext) = {
-        case m => response.reply("ack")
+      override def receiveAndReply(context: RpcCallContext) = {
+        case m => context.reply("ack")
       }
     })
 
@@ -468,8 +468,8 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
     val endpointRef = env.setupEndpoint("sendWithReply-error", new RpcEndpoint {
       override val rpcEnv = env
 
-      override def receiveAndReply(response: RpcCallContext) = {
-        case m => response.sendFailure(new SparkException("Oops"))
+      override def receiveAndReply(context: RpcCallContext) = {
+        case m => context.sendFailure(new SparkException("Oops"))
       }
     })
 
@@ -486,8 +486,8 @@ abstract class RpcEnvSuite extends FunSuite with BeforeAndAfterAll {
     env.setupEndpoint("sendWithReply-remotely-error", new RpcEndpoint {
       override val rpcEnv = env
 
-      override def receiveAndReply(response: RpcCallContext) = {
-        case msg: String => response.sendFailure(new SparkException("Oops"))
+      override def receiveAndReply(context: RpcCallContext) = {
+        case msg: String => context.sendFailure(new SparkException("Oops"))
       }
     })