Beefed up server/client integration tests.

Author: rxin

core/src/main/scala/org/apache/spark/network/netty/client/BlockFetchingClientHandler.scala

@@ -36,15 +36,15 @@ class BlockFetchingClientHandler extends SimpleChannelInboundHandler[ByteBuf] wi
   override def channelRead0(ctx: ChannelHandlerContext, in: ByteBuf) {
     val totalLen = in.readInt()
     val blockIdLen = in.readInt()
-    val blockIdBytes = new Array[Byte](blockIdLen)
+    val blockIdBytes = new Array[Byte](math.abs(blockIdLen))
     in.readBytes(blockIdBytes)
     val blockId = new String(blockIdBytes)
-    val blockLen = math.abs(totalLen) - blockIdLen - 4
+    val blockLen = totalLen - math.abs(blockIdLen) - 4
 
     def server = ctx.channel.remoteAddress.toString
 
-    // totalLen is negative when it is an error message.
-    if (totalLen < 0) {
+    // blockIdLen is negative when it is an error message.
+    if (blockIdLen < 0) {
       val errorMessageBytes = new Array[Byte](blockLen)
       in.readBytes(errorMessageBytes)
       val errorMsg = new String(errorMessageBytes)

core/src/main/scala/org/apache/spark/network/netty/server/BlockHeaderEncoder.scala

@@ -33,8 +33,8 @@ class BlockHeaderEncoder extends MessageToByteEncoder[BlockHeader] {
     msg.error match {
       case Some(errorMsg) =>
         val errorBytes = errorMsg.getBytes
-        out.writeInt(-(4 + blockId.length + errorBytes.size))
-        out.writeInt(blockId.length)
+        out.writeInt(4 + blockId.length + errorBytes.size)
+        out.writeInt(-blockId.length)
         out.writeBytes(blockId)
         out.writeBytes(errorBytes)
       case None =>

core/src/main/scala/org/apache/spark/network/netty/server/BlockServer.scala

@@ -54,7 +54,7 @@ import org.apache.spark.util.Utils
  *   frame-length (4 bytes), block-id-length (4 bytes), block-id, block-data.
  *
  *   frame-length should not include the length of itself.
- *   If frame-length is negative, then this is an error message rather than block-data. The real
+ *   If block-id-length is negative, then this is an error message rather than block-data. The real
  *   length is the absolute value of the frame-length.
  *
  */

core/src/test/scala/org/apache/spark/network/netty/ServerClientIntegrationSuite.scala

@@ -75,8 +75,10 @@ class ServerClientIntegrationSuite extends FunSuite with BeforeAndAfterAll {
     clientFactory.stop()
   }
 
+  /** A ByteBuf for buffer_block */
   lazy val byteBufferBlockReference = Unpooled.wrappedBuffer(buf)
 
+  /** A ByteBuf for file_block */
   lazy val fileBlockReference = {
     val bytes = new Array[Byte](testFile.length.toInt)
     val fp = new RandomAccessFile(testFile, "r")
@@ -85,84 +87,68 @@ class ServerClientIntegrationSuite extends FunSuite with BeforeAndAfterAll {
     Unpooled.wrappedBuffer(bytes, 10, testFile.length.toInt - 25)
   }
 
-  test("fetch a ByteBuffer block") {
+  def fetchBlocks(blockIds: Seq[String]): (Set[String], Set[ReferenceCountedBuffer], Set[String]) =
+  {
     val client = clientFactory.createClient(server.hostName, server.port)
     val sem = new Semaphore(0)
-    var receivedBlockId: String = null
-    var receivedBuffer = null.asInstanceOf[ReferenceCountedBuffer]
+    val receivedBlockIds = Collections.synchronizedSet(new HashSet[String])
+    val errorBlockIds = Collections.synchronizedSet(new HashSet[String])
+    val receivedBuffers = Collections.synchronizedSet(new HashSet[ReferenceCountedBuffer])
 
     client.fetchBlocks(
-      Seq(bufferBlockId),
+      blockIds,
       (blockId, buf) => {
-        receivedBlockId = blockId
+        receivedBlockIds.add(blockId)
         buf.retain()
-        receivedBuffer = buf
+        receivedBuffers.add(buf)
         sem.release()
       },
-      (blockId, errorMsg) => sem.release()
+      (blockId, errorMsg) => {
+        errorBlockIds.add(blockId)
+        sem.release()
+      }
     )
-
-    // This should block until the blocks are fetched
-    sem.acquire()
-
-    assert(receivedBlockId === bufferBlockId)
-    assert(receivedBuffer.underlying == byteBufferBlockReference)
-    receivedBuffer.release()
+    sem.acquire(blockIds.size)
     client.close()
+    (receivedBlockIds.toSet, receivedBuffers.toSet, errorBlockIds.toSet)
   }
 
-  test("fetch a FileSegment block via zero-copy send") {
-    val client = clientFactory.createClient(server.hostName, server.port)
-    val sem = new Semaphore(0)
-    var receivedBlockId: String = null
-    var receivedBuffer = null.asInstanceOf[ReferenceCountedBuffer]
-
-    client.fetchBlocks(
-      Seq(fileBlockId),
-      (blockId, buf) => {
-        receivedBlockId = blockId
-        buf.retain()
-        receivedBuffer = buf
-        sem.release()
-      },
-      (blockId, errorMsg) => sem.release()
-    )
+  test("fetch a ByteBuffer block") {
+    val (blockIds, buffers, failBlockIds) = fetchBlocks(Seq(bufferBlockId))
+    assert(blockIds === Set(bufferBlockId))
+    assert(buffers.map(_.underlying) === Set(byteBufferBlockReference))
+    assert(failBlockIds.isEmpty)
+    buffers.foreach(_.release())
+  }
 
-    // This should block until the blocks are fetched
-    sem.acquire()
+  test("fetch a FileSegment block via zero-copy send") {
+    val (blockIds, buffers, failBlockIds) = fetchBlocks(Seq(fileBlockId))
+    assert(blockIds === Set(fileBlockId))
+    assert(buffers.map(_.underlying) === Set(fileBlockReference))
+    assert(failBlockIds.isEmpty)
+    buffers.foreach(_.release())
+  }
 
-    assert(receivedBlockId === fileBlockId)
-    assert(receivedBuffer.underlying == fileBlockReference)
-    receivedBuffer.release()
-    client.close()
+  test("fetch a non-existent block") {
+    val (blockIds, buffers, failBlockIds) = fetchBlocks(Seq("random-block"))
+    assert(blockIds.isEmpty)
+    assert(buffers.isEmpty)
+    assert(failBlockIds === Set("random-block"))
   }
 
   test("fetch both ByteBuffer block and FileSegment block") {
-    val client = clientFactory.createClient(server.hostName, server.port)
-    val sem = new Semaphore(0)
-    val receivedBlockIds = Collections.synchronizedSet(new HashSet[String])
-    val receivedBuffers = Collections.synchronizedSet(new HashSet[ReferenceCountedBuffer])
-
-    client.fetchBlocks(
-      Seq(bufferBlockId, fileBlockId),
-      (blockId, buf) => {
-        receivedBlockIds.add(blockId)
-        buf.retain()
-        receivedBuffers.add(buf)
-        sem.release()
-      },
-      (blockId, errorMsg) => sem.release()
-    )
-
-    sem.acquire(2)
-    assert(receivedBlockIds.contains(bufferBlockId))
-    assert(receivedBlockIds.contains(fileBlockId))
-
-    val byteBufferReference = byteBufferBlockReference
-    val fileReference = fileBlockReference
+    val (blockIds, buffers, failBlockIds) = fetchBlocks(Seq(bufferBlockId, fileBlockId))
+    assert(blockIds === Set(bufferBlockId, fileBlockId))
+    assert(buffers.map(_.underlying) === Set(byteBufferBlockReference, fileBlockReference))
+    assert(failBlockIds.isEmpty)
+    buffers.foreach(_.release())
+  }
 
-    assert(receivedBuffers.map(_.underlying) === Set(byteBufferReference, fileReference))
-    receivedBuffers.foreach(_.release())
-    client.close()
+  test("fetch both ByteBuffer block and a non-existent block") {
+    val (blockIds, buffers, failBlockIds) = fetchBlocks(Seq(bufferBlockId, "random-block"))
+    assert(blockIds === Set(bufferBlockId))
+    assert(buffers.map(_.underlying) === Set(byteBufferBlockReference))
+    assert(failBlockIds === Set("random-block"))
+    buffers.foreach(_.release())
   }
 }