Implemented block uploads.

core/src/main/scala/org/apache/spark/network/BlockTransferService.scala

@@ -54,9 +54,6 @@ abstract class BlockTransferService extends Closeable {
    * Fetch a sequence of blocks from a remote node asynchronously,
    * available only after [[init]] is invoked.
    *
-   * Note that [[BlockFetchingListener.onBlockFetchSuccess]] is called once per block,
-   * while [[BlockFetchingListener.onBlockFetchFailure]] is called once per failure (not per block).
-   *
    * Note that this API takes a sequence so the implementation can batch requests, and does not
    * return a future so the underlying implementation can invoke onBlockFetchSuccess as soon as
    * the data of a block is fetched, rather than waiting for all blocks to be fetched.

core/src/main/scala/org/apache/spark/network/exceptions.scala

@@ -0,0 +1,31 @@
+/*
+ * 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
+
+class BlockFetchFailureException(blockId: String, errorMsg: String, cause: Throwable)
+  extends Exception(errorMsg, cause) {
+
+  def this(blockId: String, errorMsg: String) = this(blockId, errorMsg, null)
+}
+
+
+class BlockUploadFailureException(blockId: String, cause: Throwable)
+  extends Exception(s"Failed to fetch block $blockId", cause) {
+
+  def this(blockId: String) = this(blockId, null)
+}

core/src/main/scala/org/apache/spark/network/netty/BlockClient.scala

@@ -20,10 +20,13 @@ package org.apache.spark.network.netty
 import java.io.Closeable
 import java.util.concurrent.TimeoutException
 
+import scala.concurrent.{Future, promise}
+
 import io.netty.channel.{ChannelFuture, ChannelFutureListener}
 
 import org.apache.spark.Logging
-import org.apache.spark.network.BlockFetchingListener
+import org.apache.spark.network.{ManagedBuffer, BlockFetchingListener}
+import org.apache.spark.storage.StorageLevel
 
 
 /**
@@ -58,19 +61,19 @@ class BlockClient(cf: ChannelFuture, handler: BlockClientHandler) extends Closea
   def fetchBlocks(blockIds: Seq[String], listener: BlockFetchingListener): Unit = {
     var startTime: Long = 0
     logTrace {
-      startTime = System.nanoTime()
+      startTime = System.currentTimeMillis()
       s"Sending request $blockIds to $serverAddr"
     }
 
     blockIds.foreach { blockId =>
-      handler.addRequest(blockId, listener)
+      handler.addFetchRequest(blockId, listener)
     }
 
     cf.channel().writeAndFlush(BlockFetchRequest(blockIds)).addListener(new ChannelFutureListener {
       override def operationComplete(future: ChannelFuture): Unit = {
         if (future.isSuccess) {
           logTrace {
-            val timeTaken = (System.nanoTime() - startTime).toDouble / 1000000
+            val timeTaken = System.currentTimeMillis() - startTime
             s"Sending request $blockIds to $serverAddr took $timeTaken ms"
           }
         } else {
@@ -79,14 +82,43 @@ class BlockClient(cf: ChannelFuture, handler: BlockClientHandler) extends Closea
             s"Failed to send request $blockIds to $serverAddr: ${future.cause.getMessage}"
           logError(errorMsg, future.cause)
           blockIds.foreach { blockId =>
-            handler.removeRequest(blockId)
+            handler.removeFetchRequest(blockId)
             listener.onBlockFetchFailure(blockId, new RuntimeException(errorMsg))
           }
         }
       }
     })
   }
 
+  def uploadBlock(blockId: String, data: ManagedBuffer, storageLevel: StorageLevel): Future[Unit] = {
+    var startTime: Long = 0
+    logTrace {
+      startTime = System.currentTimeMillis()
+      s"Uploading block ($blockId) to $serverAddr"
+    }
+    val f = cf.channel().writeAndFlush(new BlockUploadRequest(blockId, data, storageLevel))
+
+    val p = promise[Unit]()
+    handler.addUploadRequest(blockId, p)
+    f.addListener(new ChannelFutureListener {
+      override def operationComplete(future: ChannelFuture): Unit = {
+        if (future.isSuccess) {
+          logTrace {
+            val timeTaken = System.currentTimeMillis() - startTime
+            s"Uploading block ($blockId) to $serverAddr took $timeTaken ms"
+          }
+        } else {
+          // Fail all blocks.
+          val errorMsg =
+            s"Failed to upload block $blockId to $serverAddr: ${future.cause.getMessage}"
+          logError(errorMsg, future.cause)
+        }
+      }
+    })
+
+    p.future
+  }
+
   /** Close the connection. This does NOT block till the connection is closed. */
   def close(): Unit = cf.channel().close()
 }

core/src/main/scala/org/apache/spark/network/netty/BlockClientHandler.scala

@@ -19,10 +19,12 @@ package org.apache.spark.network.netty
 
 import java.util.concurrent.ConcurrentHashMap
 
+import scala.concurrent.Promise
+
 import io.netty.channel.{ChannelHandlerContext, SimpleChannelInboundHandler}
 
 import org.apache.spark.Logging
-import org.apache.spark.network.BlockFetchingListener
+import org.apache.spark.network.{BlockFetchFailureException, BlockUploadFailureException, BlockFetchingListener}
 
 
 /**
@@ -35,43 +37,57 @@ private[netty]
 class BlockClientHandler extends SimpleChannelInboundHandler[ServerResponse] with Logging {
 
   /** Tracks the list of outstanding requests and their listeners on success/failure. */
-  private[this] val outstandingRequests: java.util.Map[String, BlockFetchingListener] =
+  private[this] val outstandingFetches: java.util.Map[String, BlockFetchingListener] =
     new ConcurrentHashMap[String, BlockFetchingListener]
 
-  def addRequest(blockId: String, listener: BlockFetchingListener): Unit = {
-    outstandingRequests.put(blockId, listener)
+  private[this] val outstandingUploads: java.util.Map[String, Promise[Unit]] =
+    new ConcurrentHashMap[String, Promise[Unit]]
+
+  def addFetchRequest(blockId: String, listener: BlockFetchingListener): Unit = {
+    outstandingFetches.put(blockId, listener)
   }
 
-  def removeRequest(blockId: String): Unit = {
-    outstandingRequests.remove(blockId)
+  def removeFetchRequest(blockId: String): Unit = {
+    outstandingFetches.remove(blockId)
+  }
+
+  def addUploadRequest(blockId: String, promise: Promise[Unit]): Unit = {
+    outstandingUploads.put(blockId, promise)
   }
 
   /**
    * Fire the failure callback for all outstanding requests. This is called when we have an
    * uncaught exception or pre-mature connection termination.
    */
   private def failOutstandingRequests(cause: Throwable): Unit = {
-    val iter = outstandingRequests.entrySet().iterator()
-    while (iter.hasNext) {
-      val entry = iter.next()
+    val iter1 = outstandingFetches.entrySet().iterator()
+    while (iter1.hasNext) {
+      val entry = iter1.next()
       entry.getValue.onBlockFetchFailure(entry.getKey, cause)
     }
     // TODO(rxin): Maybe we need to synchronize the access? Otherwise we could clear new requests
     // as well. But I guess that is ok given the caller will fail as soon as any requests fail.
-    outstandingRequests.clear()
+    outstandingFetches.clear()
+
+    val iter2 = outstandingUploads.entrySet().iterator()
+    while (iter2.hasNext) {
+      val entry = iter2.next()
+      entry.getValue.failure(new RuntimeException(s"Failed to upload block ${entry.getKey}"))
+    }
+    outstandingUploads.clear()
   }
 
   override def channelUnregistered(ctx: ChannelHandlerContext): Unit = {
-    if (outstandingRequests.size() > 0) {
-      logError("Still have " + outstandingRequests.size() + " requests outstanding " +
+    if (outstandingFetches.size() > 0) {
+      logError("Still have " + outstandingFetches.size() + " requests outstanding " +
         s"when connection from ${ctx.channel.remoteAddress} is closed")
       failOutstandingRequests(new RuntimeException(
         s"Connection from ${ctx.channel.remoteAddress} closed"))
     }
   }
 
   override def exceptionCaught(ctx: ChannelHandlerContext, cause: Throwable): Unit = {
-    if (outstandingRequests.size() > 0) {
+    if (outstandingFetches.size() > 0) {
       logError(
         s"Exception in connection from ${ctx.channel.remoteAddress}: ${cause.getMessage}", cause)
       failOutstandingRequests(cause)
@@ -83,23 +99,39 @@ class BlockClientHandler extends SimpleChannelInboundHandler[ServerResponse] wit
     val server = ctx.channel.remoteAddress.toString
     response match {
       case BlockFetchSuccess(blockId, buf) =>
-        val listener = outstandingRequests.get(blockId)
+        val listener = outstandingFetches.get(blockId)
         if (listener == null) {
           logWarning(s"Got a response for block $blockId from $server but it is not outstanding")
           buf.release()
         } else {
-          outstandingRequests.remove(blockId)
+          outstandingFetches.remove(blockId)
           listener.onBlockFetchSuccess(blockId, buf)
           buf.release()
         }
       case BlockFetchFailure(blockId, errorMsg) =>
-        val listener = outstandingRequests.get(blockId)
+        val listener = outstandingFetches.get(blockId)
         if (listener == null) {
           logWarning(
             s"Got a response for block $blockId from $server ($errorMsg) but it is not outstanding")
         } else {
-          outstandingRequests.remove(blockId)
-          listener.onBlockFetchFailure(blockId, new RuntimeException(errorMsg))
+          outstandingFetches.remove(blockId)
+          listener.onBlockFetchFailure(blockId, new BlockFetchFailureException(blockId, errorMsg))
+        }
+      case BlockUploadSuccess(blockId) =>
+        val p = outstandingUploads.get(blockId)
+        if (p == null) {
+          logWarning(s"Got a response for upload $blockId from $server but it is not outstanding")
+        } else {
+          outstandingUploads.remove(blockId)
+          p.success(Unit)
+        }
+      case BlockUploadFailure(blockId, error) =>
+        val p = outstandingUploads.get(blockId)
+        if (p == null) {
+          logWarning(s"Got a response for upload $blockId from $server but it is not outstanding")
+        } else {
+          outstandingUploads.remove(blockId)
+          p.failure(new BlockUploadFailureException(blockId))
         }
     }
   }

core/src/main/scala/org/apache/spark/network/netty/BlockServerHandler.scala

@@ -21,6 +21,7 @@ import io.netty.channel._
 
 import org.apache.spark.Logging
 import org.apache.spark.network.{ManagedBuffer, BlockDataManager}
+import org.apache.spark.storage.StorageLevel
 
 
 /**
@@ -39,13 +40,13 @@ private[netty] class BlockServerHandler(dataProvider: BlockDataManager)
   override def channelRead0(ctx: ChannelHandlerContext, request: ClientRequest): Unit = {
     request match {
       case BlockFetchRequest(blockIds) =>
-        blockIds.foreach(processBlockRequest(ctx, _))
-      case BlockUploadRequest(blockId, data) =>
-        // TODO(rxin): handle upload.
+        blockIds.foreach(processFetchRequest(ctx, _))
+      case BlockUploadRequest(blockId, data, level) =>
+        processUploadRequest(ctx, blockId, data, level)
     }
   }  // end of channelRead0
 
-  private def processBlockRequest(ctx: ChannelHandlerContext, blockId: String): Unit = {
+  private def processFetchRequest(ctx: ChannelHandlerContext, blockId: String): Unit = {
     // A helper function to send error message back to the client.
     def client = ctx.channel.remoteAddress.toString
 
@@ -90,4 +91,35 @@ private[netty] class BlockServerHandler(dataProvider: BlockDataManager)
       }
     )
   }  // end of processBlockRequest
+
+  private def processUploadRequest(
+      ctx: ChannelHandlerContext,
+      blockId: String,
+      data: ManagedBuffer,
+      level: StorageLevel): Unit = {
+    // A helper function to send error message back to the client.
+    def client = ctx.channel.remoteAddress.toString
+
+    try {
+      dataProvider.putBlockData(blockId, data, level)
+      ctx.writeAndFlush(BlockUploadSuccess(blockId)).addListener(new ChannelFutureListener {
+        override def operationComplete(future: ChannelFuture): Unit = {
+          if (!future.isSuccess) {
+            logError(s"Error sending an ACK back to client $client")
+          }
+        }
+      })
+    } catch {
+      case e: Throwable =>
+        logError(s"Error processing uploaded block $blockId", e)
+        ctx.writeAndFlush(BlockUploadFailure(blockId, e.getMessage)).addListener(
+          new ChannelFutureListener {
+            override def operationComplete(future: ChannelFuture): Unit = {
+              if (!future.isSuccess) {
+                logError(s"Error sending an ACK back to client $client")
+              }
+            }
+          })
+    }
+  }  // end of processUploadRequest
 }

core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala

@@ -63,9 +63,9 @@ final class NettyBlockTransferService(conf: SparkConf) extends BlockTransferServ
       hostname: String,
       port: Int,
       blockId: String,
-      blockData: ManagedBuffer, level: StorageLevel): Future[Unit] = {
-    // TODO(rxin): Implement uploadBlock.
-    ???
+      blockData: ManagedBuffer,
+      level: StorageLevel): Future[Unit] = {
+    clientFactory.createClient(hostName, port).uploadBlock(blockId, blockData, level)
   }
 
   override def hostName: String = {