Add docs, timeout config, better failure handling

Author: aarondav

core/src/main/scala/org/apache/spark/SparkConf.scala

@@ -218,8 +218,8 @@ class SparkConf(loadDefaults: Boolean) extends Cloneable with Logging {
     getAll.filter { case (k, _) => isAkkaConf(k) }
 
   /**
-   * Returns the Spark application id, valid in the Driver after TaskScheduler registration in the
-   * driver and from the start in the Executor.
+   * Returns the Spark application id, valid in the Driver after TaskScheduler registration and
+   * from the start in the Executor.
    */
   def getAppId: String = get("spark.app.id")
 

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark.network.netty
 
+import org.apache.spark.network.util.TransportConf
+
 import scala.collection.JavaConversions._
 import scala.concurrent.{Future, Promise}
 
@@ -41,6 +43,7 @@ class NettyBlockTransferService(conf: SparkConf, securityManager: SecurityManage
   // TODO: Don't use Java serialization, use a more cross-version compatible serialization format.
   private val serializer = new JavaSerializer(conf)
   private val authEnabled = securityManager.isAuthenticationEnabled()
+  private val transportConf = SparkTransportConf.fromSparkConf(conf)
 
   private[this] var transportContext: TransportContext = _
   private[this] var server: TransportServer = _
@@ -53,10 +56,10 @@ class NettyBlockTransferService(conf: SparkConf, securityManager: SecurityManage
         (nettyRpcHandler, None)
       } else {
         (new SaslRpcHandler(nettyRpcHandler, securityManager),
-          Some(new SaslBootstrap(conf.getAppId, securityManager)))
+          Some(new SaslBootstrap(transportConf, conf.getAppId, securityManager)))
       }
     }
-    transportContext = new TransportContext(SparkTransportConf.fromSparkConf(conf), rpcHandler)
+    transportContext = new TransportContext(transportConf, rpcHandler)
     clientFactory = transportContext.createClientFactory(bootstrap.toList)
     server = transportContext.createServer()
     logInfo("Server created on " + server.getPort)

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -176,8 +176,8 @@ private[spark] class BlockManager(
    * the appId may not be known at BlockManager instantiation time (in particular for the driver,
    * where it is only learned after registration with the TaskScheduler).
    *
-   * This method initializes the BlockTransferService and ShuffleClient registers with the
-   * BlockManagerMaster, starts theBlockManagerWorker actor, and registers with a local shuffle
+   * This method initializes the BlockTransferService and ShuffleClient, registers with the
+   * BlockManagerMaster, starts the BlockManagerWorker actor, and registers with a local shuffle
    * service if configured.
    */
   def initialize(appId: String): Unit = {

network/common/src/main/java/org/apache/spark/network/client/TransportClientBootstrap.java

@@ -17,6 +17,16 @@
 
 package org.apache.spark.network.client;
 
+/**
+ * A bootstrap which is executed on a TransportClient before it is returned to the user.
+ * This enables an initial exchange of information (e.g., SASL authentication tokens) on a once-per-
+ * connection basis.
+ *
+ * Since connections (and TransportClients) are reused as much as possible, it is generally
+ * reasonable to perform an expensive bootstrapping operation, as they often share a lifespan with
+ * the JVM itself.
+ */
 public interface TransportClientBootstrap {
-  public void doBootstrap(TransportClient client);
+  /** Performs the bootstrapping operation, throwing an exception on failure. */
+  public void doBootstrap(TransportClient client) throws RuntimeException;
 }

network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java

@@ -26,6 +26,8 @@
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
 
+import com.google.common.base.Preconditions;
+import com.google.common.base.Throwables;
 import com.google.common.collect.Lists;
 import io.netty.bootstrap.Bootstrap;
 import io.netty.buffer.PooledByteBufAllocator;
@@ -42,6 +44,7 @@
 import org.apache.spark.network.TransportContext;
 import org.apache.spark.network.server.TransportChannelHandler;
 import org.apache.spark.network.util.IOMode;
+import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.NettyUtils;
 import org.apache.spark.network.util.TransportConf;
 
@@ -60,19 +63,19 @@ public class TransportClientFactory implements Closeable {
 
   private final TransportContext context;
   private final TransportConf conf;
-  private final ConcurrentHashMap<SocketAddress, TransportClient> connectionPool;
   private final List<TransportClientBootstrap> clientBootstraps;
+  private final ConcurrentHashMap<SocketAddress, TransportClient> connectionPool;
 
   private final Class<? extends Channel> socketChannelClass;
   private EventLoopGroup workerGroup;
 
   public TransportClientFactory(
       TransportContext context,
       List<TransportClientBootstrap> clientBootstraps) {
-    this.context = context;
+    this.context = Preconditions.checkNotNull(context);
     this.conf = context.getConf();
+    this.clientBootstraps = Lists.newArrayList(Preconditions.checkNotNull(clientBootstraps));
     this.connectionPool = new ConcurrentHashMap<SocketAddress, TransportClient>();
-    this.clientBootstraps = Lists.newArrayList(clientBootstraps);
 
     IOMode ioMode = IOMode.valueOf(conf.ioMode());
     this.socketChannelClass = NettyUtils.getClientChannelClass(ioMode);
@@ -138,10 +141,16 @@ public void initChannel(SocketChannel ch) {
     TransportClient client = clientRef.get();
     assert client != null : "Channel future completed successfully with null client";
 
-    logger.debug("Connection to {} successful, running bootstraps...", address);
     // Execute any client bootstraps synchronously before marking the Client as successful.
-    for (TransportClientBootstrap clientBootstrap : clientBootstraps) {
-      clientBootstrap.doBootstrap(client);
+    logger.debug("Connection to {} successful, running bootstraps...", address);
+    try {
+      for (TransportClientBootstrap clientBootstrap : clientBootstraps) {
+        clientBootstrap.doBootstrap(client);
+      }
+    } catch (Exception e) { // catch Exception as the bootstrap may be written in Scala
+      logger.error("Exception while bootstrapping client", e);
+      client.close();
+      throw Throwables.propagate(e);
     }
 
     logger.debug("Successfully executed {} bootstraps for {}", clientBootstraps.size(), address);

network/common/src/main/java/org/apache/spark/network/util/TransportConf.java

@@ -55,4 +55,7 @@ public int connectionTimeoutMs() {
 
   /** Send buffer size (SO_SNDBUF). */
   public int sendBuf() { return conf.getInt("spark.shuffle.io.sendBuffer", -1); }
+
+  /** Timeout for a single round trip of SASL token exchange, in milliseconds. */
+  public int saslRTTimeout() { return conf.getInt("spark.shuffle.sasl.timeout", 30000); }
 }

network/shuffle/src/main/java/org/apache/spark/network/sasl/SaslBootstrap.java

@@ -24,31 +24,43 @@
 
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.client.TransportClientBootstrap;
+import org.apache.spark.network.util.TransportConf;
 
+/**
+ * Bootstraps a {@link TransportClient} by performing SASL authentication on the connection. The
+ * server should be setup with a {@link SaslRpcHandler} with matching keys for the given appId.
+ */
 public class SaslBootstrap implements TransportClientBootstrap {
   private final Logger logger = LoggerFactory.getLogger(SaslBootstrap.class);
 
-  private final String secretKeyId;
+  private final TransportConf conf;
+  private final String appId;
   private final SecretKeyHolder secretKeyHolder;
 
-  public SaslBootstrap(String secretKeyId, SecretKeyHolder secretKeyHolder) {
-    this.secretKeyId = secretKeyId;
+  public SaslBootstrap(TransportConf conf, String appId, SecretKeyHolder secretKeyHolder) {
+    this.conf = conf;
+    this.appId = appId;
     this.secretKeyHolder = secretKeyHolder;
   }
 
+  /**
+   * Performs SASL authentication by sending a token, and then proceeding with the SASL
+   * challenge-response tokens until we either successfully authenticate or throw an exception
+   * due to mismatch.
+   */
   public void doBootstrap(TransportClient client) {
-    SparkSaslClient saslClient = new SparkSaslClient(secretKeyId, secretKeyHolder);
+    SparkSaslClient saslClient = new SparkSaslClient(appId, secretKeyHolder);
     try {
       byte[] payload = saslClient.firstToken();
 
       while (!saslClient.isComplete()) {
-        SaslMessage msg = new SaslMessage(secretKeyId, payload);
-        logger.info("Sending msg {} {}", secretKeyId, payload.length);
+        SaslMessage msg = new SaslMessage(appId, payload);
+        logger.info("Sending msg {} {}", appId, payload.length);
         ByteBuf buf = Unpooled.buffer(msg.encodedLength());
         msg.encode(buf);
 
-        byte[] response = client.sendRpcSync(buf.array(), 300000);
-        logger.info("Got response {} {}", secretKeyId, response.length);
+        byte[] response = client.sendRpcSync(buf.array(), conf.saslRTTimeout());
+        logger.info("Got response {} {}", appId, response.length);
         payload = saslClient.response(response);
       }
     } finally {

network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java

@@ -96,7 +96,7 @@ public void afterEach() {
   public void testGoodClient() {
     clientFactory = context.createClientFactory(
       Lists.<TransportClientBootstrap>newArrayList(
-        new SaslBootstrap("app-id", new TestSecretKeyHolder("good-key"))));
+        new SaslBootstrap(conf, "app-id", new TestSecretKeyHolder("good-key"))));
 
     TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
     String msg = "Hello, World!";
@@ -108,7 +108,7 @@ public void testGoodClient() {
   public void testBadClient() {
     clientFactory = context.createClientFactory(
       Lists.<TransportClientBootstrap>newArrayList(
-        new SaslBootstrap("app-id", new TestSecretKeyHolder("bad-key"))));
+        new SaslBootstrap(conf, "app-id", new TestSecretKeyHolder("bad-key"))));
 
     try {
       // Bootstrap should fail on startup.
@@ -146,7 +146,7 @@ public void testNoSaslServer() {
     TransportContext context = new TransportContext(conf, handler);
     clientFactory = context.createClientFactory(
       Lists.<TransportClientBootstrap>newArrayList(
-        new SaslBootstrap("app-id", new TestSecretKeyHolder("key"))));
+        new SaslBootstrap(conf, "app-id", new TestSecretKeyHolder("key"))));
     TransportServer server = context.createServer();
     try {
       clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());