Added test cases for BlockHeaderEncoder and BlockFetchingClientHandlerSuite.

Author: rxin

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

@@ -33,6 +33,8 @@ import org.apache.spark.Logging
  * Client for fetching data blocks from [[org.apache.spark.network.netty.server.BlockServer]].
  * Use [[BlockFetchingClientFactory]] to instantiate this client.
  *
+ * The constructor blocks until a connection is successfully established.
+ *
  * See [[org.apache.spark.network.netty.server.BlockServer]] for client/server protocol.
  *
  * Concurrency: [[BlockFetchingClient]] is not thread safe and should not be shared.

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

@@ -83,6 +83,8 @@ class BlockFetchingClientFactory(val conf: NettyConfig) {
   /**
    * Create a new BlockFetchingClient connecting to the given remote host / port.
    *
+   * This blocks until a connection is successfully established.
+   *
    * Concurrency: This method is safe to call from multiple threads.
    */
   def createClient(remoteHost: String, remotePort: Int): BlockFetchingClient = {

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

@@ -44,19 +44,19 @@ class BlockFetchingClientHandler extends SimpleChannelInboundHandler[ByteBuf] wi
     val blockIdBytes = new Array[Byte](math.abs(blockIdLen))
     in.readBytes(blockIdBytes)
     val blockId = new String(blockIdBytes)
-    val blockLen = totalLen - math.abs(blockIdLen) - 4
+    val blockSize = totalLen - math.abs(blockIdLen) - 4
 
     def server = ctx.channel.remoteAddress.toString
 
     // blockIdLen is negative when it is an error message.
     if (blockIdLen < 0) {
-      val errorMessageBytes = new Array[Byte](blockLen)
+      val errorMessageBytes = new Array[Byte](blockSize)
       in.readBytes(errorMessageBytes)
       val errorMsg = new String(errorMessageBytes)
-      logTrace(s"Received block $blockId ($blockLen B) with error $errorMsg from $server")
+      logTrace(s"Received block $blockId ($blockSize B) with error $errorMsg from $server")
       blockFetchFailureCallback(blockId, errorMsg)
     } else {
-      logTrace(s"Received block $blockId ($blockLen B) from $server")
+      logTrace(s"Received block $blockId ($blockSize B) from $server")
       blockFetchSuccessCallback(blockId, new ReferenceCountedBuffer(in))
     }
   }

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

@@ -20,7 +20,7 @@ package org.apache.spark.network.netty
 import java.io.{RandomAccessFile, File}
 import java.nio.ByteBuffer
 import java.util.{Collections, HashSet}
-import java.util.concurrent.Semaphore
+import java.util.concurrent.{TimeUnit, Semaphore}
 
 import scala.collection.JavaConversions._
 
@@ -34,6 +34,9 @@ import org.apache.spark.network.netty.server.BlockServer
 import org.apache.spark.storage.{FileSegment, BlockDataProvider}
 
 
+/**
+ * Test suite that makes sure the server and the client implementations share the same protocol.
+ */
 class ServerClientIntegrationSuite extends FunSuite with BeforeAndAfterAll {
 
   val bufSize = 100000
@@ -108,7 +111,9 @@ class ServerClientIntegrationSuite extends FunSuite with BeforeAndAfterAll {
         sem.release()
       }
     )
-    sem.acquire(blockIds.size)
+    if (!sem.tryAcquire(blockIds.size, 30, TimeUnit.SECONDS)) {
+      fail("Timeout getting response from the server")
+    }
     client.close()
     (receivedBlockIds.toSet, receivedBuffers.toSet, errorBlockIds.toSet)
   }

core/src/test/scala/org/apache/spark/network/netty/client/BlockFetchingClientHandlerSuite.scala

@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.netty.client
+
+import java.nio.ByteBuffer
+
+import io.netty.buffer.Unpooled
+import io.netty.channel.embedded.EmbeddedChannel
+
+import org.scalatest.FunSuite
+
+
+class BlockFetchingClientHandlerSuite extends FunSuite {
+
+  test("handling block data (successful fetch)") {
+    val blockId = "test_block"
+    val blockData = "blahblahblahblahblah"
+    val totalLength = 4 + blockId.length + blockData.length
+
+    var parsedBlockId: String = ""
+    var parsedBlockData: String = ""
+    val handler = new BlockFetchingClientHandler
+    handler.blockFetchSuccessCallback = (bid, refCntBuf) => {
+      parsedBlockId = bid
+      val bytes = new Array[Byte](refCntBuf.byteBuffer().remaining)
+      refCntBuf.byteBuffer().get(bytes)
+      parsedBlockData = new String(bytes)
+    }
+
+    val channel = new EmbeddedChannel(handler)
+    val buf = ByteBuffer.allocate(totalLength + 4)  // 4 bytes for the length field itself
+    buf.putInt(totalLength)
+    buf.putInt(blockId.length)
+    buf.put(blockId.getBytes)
+    buf.put(blockData.getBytes)
+    buf.flip()
+
+    channel.writeInbound(Unpooled.wrappedBuffer(buf))
+    assert(parsedBlockId === blockId)
+    assert(parsedBlockData === blockData)
+
+    channel.close()
+  }
+
+  test("handling error message (failed fetch)") {
+    val blockId = "test_block"
+    val errorMsg = "error erro5r error err4or error3 error6 error erro1r"
+    val totalLength = 4 + blockId.length + errorMsg.length
+
+    var parsedBlockId: String = ""
+    var parsedErrorMsg: String = ""
+    val handler = new BlockFetchingClientHandler
+    handler.blockFetchFailureCallback = (bid, msg) => {
+      parsedBlockId = bid
+      parsedErrorMsg = msg
+    }
+
+    val channel = new EmbeddedChannel(handler)
+    val buf = ByteBuffer.allocate(totalLength + 4)  // 4 bytes for the length field itself
+    buf.putInt(totalLength)
+    buf.putInt(-blockId.length)
+    buf.put(blockId.getBytes)
+    buf.put(errorMsg.getBytes)
+    buf.flip()
+
+    channel.writeInbound(Unpooled.wrappedBuffer(buf))
+    assert(parsedBlockId === blockId)
+    assert(parsedErrorMsg === errorMsg)
+
+    channel.close()
+  }
+}

core/src/test/scala/org/apache/spark/network/netty/server/BlockHeaderEncoderSuite.scala

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.netty.server
+
+import io.netty.buffer.ByteBuf
+import io.netty.channel.embedded.EmbeddedChannel
+
+import org.scalatest.FunSuite
+
+
+class BlockHeaderEncoderSuite extends FunSuite {
+
+  test("encode normal block data") {
+    val blockId = "test_block"
+    val channel = new EmbeddedChannel(new BlockHeaderEncoder)
+    channel.writeOutbound(new BlockHeader(17, blockId, None))
+    val out = channel.readOutbound().asInstanceOf[ByteBuf]
+    assert(out.readInt() === 4 + blockId.length + 17)
+    assert(out.readInt() === blockId.length)
+
+    val blockIdBytes = new Array[Byte](blockId.length)
+    out.readBytes(blockIdBytes)
+    assert(new String(blockIdBytes) === blockId)
+    assert(out.readableBytes() === 0)
+
+    channel.close()
+  }
+
+  test("encode error message") {
+    val blockId = "error_block"
+    val errorMsg = "error encountered"
+    val channel = new EmbeddedChannel(new BlockHeaderEncoder)
+    channel.writeOutbound(new BlockHeader(17, blockId, Some(errorMsg)))
+    val out = channel.readOutbound().asInstanceOf[ByteBuf]
+    assert(out.readInt() === 4 + blockId.length + errorMsg.length)
+    assert(out.readInt() === -blockId.length)
+
+    val blockIdBytes = new Array[Byte](blockId.length)
+    out.readBytes(blockIdBytes)
+    assert(new String(blockIdBytes) === blockId)
+
+    val errorMsgBytes = new Array[Byte](errorMsg.length)
+    out.readBytes(errorMsgBytes)
+    assert(new String(errorMsgBytes) === errorMsg)
+    assert(out.readableBytes() === 0)
+
+    channel.close()
+  }
+}