Spark RPC 学习总结

  val securityMgr = new SecurityManager(conf)
  val rpcEnv = RpcEnv.create(SYSTEM_NAME, host, port, conf, securityMgr)


  def create(
      name: String,
      host: String,
      port: Int,
      conf: SparkConf,
      securityManager: SecurityManager,
      clientMode: Boolean = false): RpcEnv = {
    create(name, host, host, port, conf, securityManager, 0, clientMode)
  }


  def create(
      name: String,
      bindAddress: String,
      advertiseAddress: String,
      port: Int,
      conf: SparkConf,
      securityManager: SecurityManager,
      numUsableCores: Int,
      clientMode: Boolean): RpcEnv = {
    val config = RpcEnvConfig(conf, name, bindAddress, advertiseAddress, port, securityManager,
      numUsableCores, clientMode)
    new NettyRpcEnvFactory().create(config)
  }

  def createRpcEnv(
                    conf: SparkConf,
                    name: String,
                    port: Int,
                    clientMode: Boolean = false): RpcEnv = {
    val config = RpcEnvConfig(conf, name, "localhost", "localhost", port,
      new SecurityManager(conf), 0, clientMode)
    new NettyRpcEnvFactory().create(config)
  }

def setupEndpoint(name: String, endpoint: RpcEndpoint): RpcEndpointRef

  /**
   * Retrieve the [[RpcEndpointRef]] represented by `uri` asynchronously.
   */
  def asyncSetupEndpointRefByURI(uri: String): Future[RpcEndpointRef]


  /**
   * Retrieve the [[RpcEndpointRef]] represented by `uri`. This is a blocking action.
   */
  def setupEndpointRefByURI(uri: String): RpcEndpointRef = {
    defaultLookupTimeout.awaitResult(asyncSetupEndpointRefByURI(uri))
  }


  /**
   * Retrieve the [[RpcEndpointRef]] represented by `address` and `endpointName`.
   * This is a blocking action.
   */
  def setupEndpointRef(address: RpcAddress, endpointName: String): RpcEndpointRef = {
    setupEndpointRefByURI(RpcEndpointAddress(address, endpointName).toString)
  }

// Worer 中返回 Master 的引用
val masterEndpoint = rpcEnv.setupEndpointRef(masterAddress, Master.ENDPOINT_NAME)
// CoarseGrainedExecutorBackend 中获取 Driver 的引用
driver = fetcher.setupEndpointRefByURI(arguments.driverUrl)
rpcEnv.asyncSetupEndpointRefByURI(driverUrl)

package org.apache.spark.rpc


import org.apache.spark.{SparkConf, SparkEnv}
import org.scalatest.concurrent.Eventually.{eventually, interval, timeout}


import scala.concurrent.duration._


object RpcLocalTest extends RpcParent {
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf()
    val env = createRpcEnv(conf, "local", 8000)
    @volatile var message: String = null
    val rpcEndpointRef = env.setupEndpoint("send-locally", new RpcEndpoint {
      override val rpcEnv = env


      override def onStart(): Unit = {
        println("start hello endpoint")
      }


      override def receive = {
        case msg: String =>
          println(msg)
          message = msg
      }
    })
    rpcEndpointRef.send("hello")


    eventually(timeout(5.seconds), interval(10.milliseconds)) {
      assert("hello" == message)
    }


    if (env != null) {
      env.shutdown()
    }
    SparkEnv.set(null)
  }
}

package org.apache.spark.rpc


import org.apache.spark.SparkConf


import java.util.concurrent.CountDownLatch


object RpcRemoteServer extends RpcParent {
  def main(args: Array[String]): Unit = {
    val shutdownLatch = new CountDownLatch(1)


    val env = createRpcEnv(new SparkConf(), "local", 8000)
    println(s"地址：${env.address}")
    env.setupEndpoint("ask-remotely", new RpcEndpoint {
      override val rpcEnv = env
      override def onStart(): Unit = {
        println("onStart 被调用")
      }
      override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
        case msg: String =>
          context.reply(msg)
      }
    })
    shutdownLatch.await()
  }
}

package org.apache.spark.rpc


import org.apache.spark.SparkConf


import java.util.concurrent.CountDownLatch


object RpcRemoteTest extends RpcParent {


  def main(args: Array[String]): Unit = {
    val shutdownLatch = new CountDownLatch(1)
    val anotherEnv = createRpcEnv(new SparkConf(), "remote-client", 0, clientMode = true)
    println(s"地址：${anotherEnv.address}")


    val rpcEndpointRef = anotherEnv.setupEndpointRef(new RpcAddress("localhost", 8000), "ask-remotely")
    val reply = rpcEndpointRef.askSync[String]("hello Remote")
    println(reply)
    shutdownLatch.await()
  }
}

  test("onStart and onStop") {
    val stopLatch = new CountDownLatch(1)
    val calledMethods = mutable.ArrayBuffer[String]()


    val endpoint = new RpcEndpoint {
      override val rpcEnv = env


      override def onStart(): Unit = {
        println("onStart 被调用")
        calledMethods += "start"
      }


      override def receive: PartialFunction[Any, Unit] = {
        case msg: String =>
      }


      override def onStop(): Unit = {
        println("onStop 被调用")
        calledMethods += "stop"
//        stopLatch.countDown()
      }
    }
    println("调用setupEndpoint前")
    val rpcEndpointRef = env.setupEndpoint("start-stop-test", endpoint)
    println("调用setupEndpoint后")
    println("调用stop前")
    env.stop(rpcEndpointRef)
    println("调用stop后")
    stopLatch.await(10, TimeUnit.SECONDS)
    assert(List("start", "stop") === calledMethods)
  }