Address

Author: zsxwing

core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala

@@ -37,7 +37,7 @@ import org.apache.spark.network.crypto.{AuthClientBootstrap, AuthServerBootstrap
 import org.apache.spark.network.netty.SparkTransportConf
 import org.apache.spark.network.server._
 import org.apache.spark.rpc._
-import org.apache.spark.serializer.{JavaSerializer, JavaSerializerInstance}
+import org.apache.spark.serializer.{JavaSerializer, JavaSerializerInstance, SerializationStream}
 import org.apache.spark.util.{ByteBufferInputStream, ByteBufferOutputStream, ThreadUtils, Utils}
 
 private[netty] class NettyRpcEnv(
@@ -253,6 +253,13 @@ private[netty] class NettyRpcEnv(
     javaSerializerInstance.serialize(content)
   }
 
+  /**
+   * Returns [[SerializationStream]] that forwards the serialized bytes to `out`.
+   */
+  private[netty] def serializeStream(out: OutputStream): SerializationStream = {
+    javaSerializerInstance.serializeStream(out)
+  }
+
   private[netty] def deserialize[T: ClassTag](client: TransportClient, bytes: ByteBuffer): T = {
     NettyRpcEnv.currentClient.withValue(client) {
       deserialize { () =>
@@ -530,19 +537,23 @@ private[netty] class NettyRpcEndpointRef(
  */
 private[netty] class RequestMessage(
     val senderAddress: RpcAddress,
-    val receiver: NettyRpcEndpointRef, val content: Any) {
+    val receiver: NettyRpcEndpointRef,
+    val content: Any) {
 
-  /** Manually serialize [[RequestMessage]] to minimize the size of bytes. */
+  /** Manually serialize [[RequestMessage]] to minimize the size. */
   def serialize(nettyEnv: NettyRpcEnv): ByteBuffer = {
     val bos = new ByteBufferOutputStream()
     val out = new DataOutputStream(bos)
     try {
       writeRpcAddress(out, senderAddress)
       writeRpcAddress(out, receiver.address)
       out.writeUTF(receiver.name)
-      val contentBytes = nettyEnv.serialize(content)
-      assert(contentBytes.hasArray)
-      out.write(contentBytes.array, contentBytes.arrayOffset, contentBytes.remaining)
+      val s = nettyEnv.serializeStream(out)
+      try {
+        s.writeObject(content)
+      } finally {
+        s.close()
+      }
     } finally {
       out.close()
     }

core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcEnvSuite.scala

@@ -17,10 +17,13 @@
 
 package org.apache.spark.rpc.netty
 
+import org.scalatest.mock.MockitoSugar
+
 import org.apache.spark._
+import org.apache.spark.network.client.TransportClient
 import org.apache.spark.rpc._
 
-class NettyRpcEnvSuite extends RpcEnvSuite {
+class NettyRpcEnvSuite extends RpcEnvSuite with MockitoSugar {
 
   override def createRpcEnv(
       conf: SparkConf,
@@ -53,4 +56,32 @@ class NettyRpcEnvSuite extends RpcEnvSuite {
     }
   }
 
+  test("RequestMessage serialization") {
+    def assertRequestMessageEquals(expected: RequestMessage, actual: RequestMessage): Unit = {
+      assert(expected.senderAddress === actual.senderAddress)
+      assert(expected.receiver === actual.receiver)
+      assert(expected.content === actual.content)
+    }
+
+    val nettyEnv = env.asInstanceOf[NettyRpcEnv]
+    val client = mock[TransportClient]
+    val senderAddress = RpcAddress("locahost", 12345)
+    val receiverAddress = RpcEndpointAddress("localhost", 54321, "test")
+    val receiver = new NettyRpcEndpointRef(nettyEnv.conf, receiverAddress, nettyEnv)
+
+    val msg = new RequestMessage(senderAddress, receiver, "foo")
+    assertRequestMessageEquals(
+      msg,
+      RequestMessage(nettyEnv, client, msg.serialize(nettyEnv)))
+
+    val msg2 = new RequestMessage(null, receiver, "foo")
+    assertRequestMessageEquals(
+      msg2,
+      RequestMessage(nettyEnv, client, msg2.serialize(nettyEnv)))
+
+    val msg3 = new RequestMessage(senderAddress, receiver, null)
+    assertRequestMessageEquals(
+      msg3,
+      RequestMessage(nettyEnv, client, msg3.serialize(nettyEnv)))
+  }
 }