[SPARK-19529][BRANCH-1.6] Backport PR #16866 to branch-1.6

liancheng · JoshRosen · commit a50ef3d9a06f · 2017-02-14T11:27:37.000-08:00
## What changes were proposed in this pull request? This PR backports PR #16866 to branch-1.6 ## How was this patch tested? Existing tests. Author: Cheng Lian <lian@databricks.com> Closes #16917 from liancheng/spark-19529-1.6-backport.
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java b/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
@@ -120,7 +120,8 @@ public TransportClientFactory(
    *
    * Concurrency: This method is safe to call from multiple threads.
    */
-  public TransportClient createClient(String remoteHost, int remotePort) throws IOException {
+  public TransportClient createClient(String remoteHost, int remotePort)
+      throws IOException, InterruptedException {
     // Get connection from the connection pool first.
     // If it is not found or not active, create a new one.
     final InetSocketAddress address = new InetSocketAddress(remoteHost, remotePort);
@@ -176,13 +177,14 @@ public TransportClient createClient(String remoteHost, int remotePort) throws IO
    * As with {@link #createClient(String, int)}, this method is blocking.
    */
   public TransportClient createUnmanagedClient(String remoteHost, int remotePort)
-      throws IOException {
+      throws IOException, InterruptedException {
     final InetSocketAddress address = new InetSocketAddress(remoteHost, remotePort);
     return createClient(address);
   }
 
   /** Create a completely new {@link TransportClient} to the remote address. */
-  private TransportClient createClient(InetSocketAddress address) throws IOException {
+  private TransportClient createClient(InetSocketAddress address)
+      throws IOException, InterruptedException {
     logger.debug("Creating new connection to " + address);
 
     Bootstrap bootstrap = new Bootstrap();
@@ -209,7 +211,7 @@ public void initChannel(SocketChannel ch) {
     // Connect to the remote server
     long preConnect = System.nanoTime();
     ChannelFuture cf = bootstrap.connect(address);
-    if (!cf.awaitUninterruptibly(conf.connectionTimeoutMs())) {
+    if (!cf.await(conf.connectionTimeoutMs())) {
       throw new IOException(
         String.format("Connecting to %s timed out (%s ms)", address, conf.connectionTimeoutMs()));
     } else if (cf.cause() != null) {
diff --git a/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java b/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java
@@ -99,6 +99,8 @@ public void run() {
             clients.add(client);
           } catch (IOException e) {
             failed.incrementAndGet();
+          } catch (InterruptedException e) {
+            throw new RuntimeException(e);
           }
         }
       };
@@ -140,7 +142,7 @@ public void reuseClientsUpToConfigVariableConcurrent() throws Exception {
   }
 
   @Test
-  public void returnDifferentClientsForDifferentServers() throws IOException {
+  public void returnDifferentClientsForDifferentServers() throws IOException, InterruptedException {
     TransportClientFactory factory = context.createClientFactory();
     TransportClient c1 = factory.createClient(TestUtils.getLocalHost(), server1.getPort());
     TransportClient c2 = factory.createClient(TestUtils.getLocalHost(), server2.getPort());
@@ -169,7 +171,7 @@ public void neverReturnInactiveClients() throws IOException, InterruptedExceptio
   }
 
   @Test
-  public void closeBlockClientsWithFactory() throws IOException {
+  public void closeBlockClientsWithFactory() throws IOException, InterruptedException {
     TransportClientFactory factory = context.createClientFactory();
     TransportClient c1 = factory.createClient(TestUtils.getLocalHost(), server1.getPort());
     TransportClient c2 = factory.createClient(TestUtils.getLocalHost(), server2.getPort());
diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java
@@ -101,7 +101,7 @@ public void fetchBlocks(
         new RetryingBlockFetcher.BlockFetchStarter() {
           @Override
           public void createAndStart(String[] blockIds, BlockFetchingListener listener)
-              throws IOException {
+              throws IOException, InterruptedException {
             TransportClient client = clientFactory.createClient(host, port);
             new OneForOneBlockFetcher(client, appId, execId, blockIds, listener).start();
           }
@@ -136,7 +136,7 @@ public void registerWithShuffleServer(
       String host,
       int port,
       String execId,
-      ExecutorShuffleInfo executorInfo) throws IOException {
+      ExecutorShuffleInfo executorInfo) throws IOException, InterruptedException {
     checkInit();
     TransportClient client = clientFactory.createUnmanagedClient(host, port);
     try {
diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/RetryingBlockFetcher.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/RetryingBlockFetcher.java
@@ -57,7 +57,8 @@ public static interface BlockFetchStarter {
      * {@link org.apache.spark.network.client.TransportClientFactory} in order to fix connection
      * issues.
      */
-    void createAndStart(String[] blockIds, BlockFetchingListener listener) throws IOException;
+    void createAndStart(String[] blockIds, BlockFetchingListener listener)
+         throws IOException, InterruptedException;
   }
 
   /** Shared executor service used for waiting and retrying. */
diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/mesos/MesosExternalShuffleClient.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/mesos/MesosExternalShuffleClient.java
@@ -53,7 +53,8 @@ public MesosExternalShuffleClient(
     super(conf, secretKeyHolder, saslEnabled, saslEncryptionEnabled);
   }
 
-  public void registerDriverWithShuffleService(String host, int port) throws IOException {
+  public void registerDriverWithShuffleService(String host, int port)
+      throws IOException, InterruptedException {
     checkInit();
     ByteBuffer registerDriver = new RegisterDriver(appId).toByteBuffer();
     TransportClient client = clientFactory.createClient(host, port);
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java
@@ -102,7 +102,7 @@ public void afterEach() {
   }
 
   @Test
-  public void testGoodClient() throws IOException {
+  public void testGoodClient() throws IOException, InterruptedException {
     clientFactory = context.createClientFactory(
       Lists.<TransportClientBootstrap>newArrayList(
         new SaslClientBootstrap(conf, "app-1", secretKeyHolder)));
@@ -132,7 +132,7 @@ public void testBadClient() {
   }
 
   @Test
-  public void testNoSaslClient() throws IOException {
+  public void testNoSaslClient() throws IOException, InterruptedException {
     clientFactory = context.createClientFactory(
       Lists.<TransportClientBootstrap>newArrayList());
 
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleIntegrationSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleIntegrationSuite.java
@@ -273,7 +273,7 @@ public void testFetchNoServer() throws Exception {
   }
 
   private void registerExecutor(String executorId, ExecutorShuffleInfo executorInfo)
-      throws IOException {
+      throws IOException, InterruptedException {
     ExternalShuffleClient client = new ExternalShuffleClient(conf, null, false, false);
     client.init(APP_ID);
     client.registerWithShuffleServer(TestUtils.getLocalHost(), server.getPort(),
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleSecuritySuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleSecuritySuite.java
@@ -60,7 +60,7 @@ public void afterEach() {
   }
 
   @Test
-  public void testValid() throws IOException {
+  public void testValid() throws IOException, InterruptedException {
     validate("my-app-id", "secret", false);
   }
 
@@ -83,12 +83,13 @@ public void testBadSecret() {
   }
 
   @Test
-  public void testEncryption() throws IOException {
+  public void testEncryption() throws IOException, InterruptedException {
     validate("my-app-id", "secret", true);
   }
 
   /** Creates an ExternalShuffleClient and attempts to register with the server. */
-  private void validate(String appId, String secretKey, boolean encrypt) throws IOException {
+  private void validate(String appId, String secretKey, boolean encrypt)
+      throws IOException, InterruptedException {
     ExternalShuffleClient client =
       new ExternalShuffleClient(conf, new TestSecretKeyHolder(appId, secretKey), true, encrypt);
     client.init(appId);
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/RetryingBlockFetcherSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/RetryingBlockFetcherSuite.java
@@ -66,7 +66,7 @@ public void afterEach() {
   }
 
   @Test
-  public void testNoFailures() throws IOException {
+  public void testNoFailures() throws IOException, InterruptedException {
     BlockFetchingListener listener = mock(BlockFetchingListener.class);
 
     List<? extends Map<String, Object>> interactions = Arrays.asList(
@@ -85,7 +85,7 @@ public void testNoFailures() throws IOException {
   }
 
   @Test
-  public void testUnrecoverableFailure() throws IOException {
+  public void testUnrecoverableFailure() throws IOException, InterruptedException {
     BlockFetchingListener listener = mock(BlockFetchingListener.class);
 
     List<? extends Map<String, Object>> interactions = Arrays.asList(
@@ -104,7 +104,7 @@ public void testUnrecoverableFailure() throws IOException {
   }
 
   @Test
-  public void testSingleIOExceptionOnFirst() throws IOException {
+  public void testSingleIOExceptionOnFirst() throws IOException, InterruptedException {
     BlockFetchingListener listener = mock(BlockFetchingListener.class);
 
     List<? extends Map<String, Object>> interactions = Arrays.asList(
@@ -127,7 +127,7 @@ public void testSingleIOExceptionOnFirst() throws IOException {
   }
 
   @Test
-  public void testSingleIOExceptionOnSecond() throws IOException {
+  public void testSingleIOExceptionOnSecond() throws IOException, InterruptedException {
     BlockFetchingListener listener = mock(BlockFetchingListener.class);
 
     List<? extends Map<String, Object>> interactions = Arrays.asList(
@@ -149,7 +149,7 @@ public void testSingleIOExceptionOnSecond() throws IOException {
   }
 
   @Test
-  public void testTwoIOExceptions() throws IOException {
+  public void testTwoIOExceptions() throws IOException, InterruptedException {
     BlockFetchingListener listener = mock(BlockFetchingListener.class);
 
     List<? extends Map<String, Object>> interactions = Arrays.asList(
@@ -177,7 +177,7 @@ public void testTwoIOExceptions() throws IOException {
   }
 
   @Test
-  public void testThreeIOExceptions() throws IOException {
+  public void testThreeIOExceptions() throws IOException, InterruptedException {
     BlockFetchingListener listener = mock(BlockFetchingListener.class);
 
     List<? extends Map<String, Object>> interactions = Arrays.asList(
@@ -209,7 +209,7 @@ public void testThreeIOExceptions() throws IOException {
   }
 
   @Test
-  public void testRetryAndUnrecoverable() throws IOException {
+  public void testRetryAndUnrecoverable() throws IOException, InterruptedException {
     BlockFetchingListener listener = mock(BlockFetchingListener.class);
 
     List<? extends Map<String, Object>> interactions = Arrays.asList(
@@ -252,7 +252,7 @@ public void testRetryAndUnrecoverable() throws IOException {
   @SuppressWarnings("unchecked")
   private static void performInteractions(List<? extends Map<String, Object>> interactions,
                                           BlockFetchingListener listener)
-    throws IOException {
+      throws IOException, InterruptedException {
 
     TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     BlockFetchStarter fetchStarter = mock(BlockFetchStarter.class);

Original file line number	Diff line number	Diff line change
`@@ -99,6 +99,8 @@ public void run() {`
`99`	`99`	`clients.add(client);`
`100`	`100`	`} catch (IOException e) {`
`101`	`101`	`failed.incrementAndGet();`
	`102`	`+ } catch (InterruptedException e) {`
	`103`	`+ throw new RuntimeException(e);`
`102`	`104`	`}`
`103`	`105`	`}`
`104`	`106`	`};`
`@@ -140,7 +142,7 @@ public void reuseClientsUpToConfigVariableConcurrent() throws Exception {`
`140`	`142`	`}`
`141`	`143`
`142`	`144`	`@Test`
`143`		`- public void returnDifferentClientsForDifferentServers() throws IOException {`
	`145`	`+ public void returnDifferentClientsForDifferentServers() throws IOException, InterruptedException {`
`144`	`146`	`TransportClientFactory factory = context.createClientFactory();`
`145`	`147`	`TransportClient c1 = factory.createClient(TestUtils.getLocalHost(), server1.getPort());`
`146`	`148`	`TransportClient c2 = factory.createClient(TestUtils.getLocalHost(), server2.getPort());`
`@@ -169,7 +171,7 @@ public void neverReturnInactiveClients() throws IOException, InterruptedExceptio`
`169`	`171`	`}`
`170`	`172`
`171`	`173`	`@Test`
`172`		`- public void closeBlockClientsWithFactory() throws IOException {`
	`174`	`+ public void closeBlockClientsWithFactory() throws IOException, InterruptedException {`
`173`	`175`	`TransportClientFactory factory = context.createClientFactory();`
`174`	`176`	`TransportClient c1 = factory.createClient(TestUtils.getLocalHost(), server1.getPort());`
`175`	`177`	`TransportClient c2 = factory.createClient(TestUtils.getLocalHost(), server2.getPort());`
Original file line number	Diff line number	Diff line change
`@@ -57,7 +57,8 @@ public static interface BlockFetchStarter {`
`57`	`57`	`* {@link org.apache.spark.network.client.TransportClientFactory} in order to fix connection`
`58`	`58`	`* issues.`
`59`	`59`	`*/`
`60`		`- void createAndStart(String[] blockIds, BlockFetchingListener listener) throws IOException;`
	`60`	`+ void createAndStart(String[] blockIds, BlockFetchingListener listener)`
	`61`	`+ throws IOException, InterruptedException;`
`61`	`62`	`}`
`62`	`63`
`63`	`64`	`/** Shared executor service used for waiting and retrying. */`
Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,8 @@ public MesosExternalShuffleClient(`
`53`	`53`	`super(conf, secretKeyHolder, saslEnabled, saslEncryptionEnabled);`
`54`	`54`	`}`
`55`	`55`
`56`		`- public void registerDriverWithShuffleService(String host, int port) throws IOException {`
	`56`	`+ public void registerDriverWithShuffleService(String host, int port)`
	`57`	`+ throws IOException, InterruptedException {`
`57`	`58`	`checkInit();`
`58`	`59`	`ByteBuffer registerDriver = new RegisterDriver(appId).toByteBuffer();`
`59`	`60`	`TransportClient client = clientFactory.createClient(host, port);`
Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,7 @@ public void afterEach() {`
`102`	`102`	`}`
`103`	`103`
`104`	`104`	`@Test`
`105`		`- public void testGoodClient() throws IOException {`
	`105`	`+ public void testGoodClient() throws IOException, InterruptedException {`
`106`	`106`	`clientFactory = context.createClientFactory(`
`107`	`107`	`Lists.<TransportClientBootstrap>newArrayList(`
`108`	`108`	`new SaslClientBootstrap(conf, "app-1", secretKeyHolder)));`
`@@ -132,7 +132,7 @@ public void testBadClient() {`
`132`	`132`	`}`
`133`	`133`
`134`	`134`	`@Test`
`135`		`- public void testNoSaslClient() throws IOException {`
	`135`	`+ public void testNoSaslClient() throws IOException, InterruptedException {`
`136`	`136`	`clientFactory = context.createClientFactory(`
`137`	`137`	`Lists.<TransportClientBootstrap>newArrayList());`
`138`	`138`
Original file line number	Diff line number	Diff line change
`@@ -273,7 +273,7 @@ public void testFetchNoServer() throws Exception {`
`273`	`273`	`}`
`274`	`274`
`275`	`275`	`private void registerExecutor(String executorId, ExecutorShuffleInfo executorInfo)`
`276`		`- throws IOException {`
	`276`	`+ throws IOException, InterruptedException {`
`277`	`277`	`ExternalShuffleClient client = new ExternalShuffleClient(conf, null, false, false);`
`278`	`278`	`client.init(APP_ID);`
`279`	`279`	`client.registerWithShuffleServer(TestUtils.getLocalHost(), server.getPort(),`
Original file line number	Diff line number	Diff line change
`@@ -60,7 +60,7 @@ public void afterEach() {`
`60`	`60`	`}`
`61`	`61`
`62`	`62`	`@Test`
`63`		`- public void testValid() throws IOException {`
	`63`	`+ public void testValid() throws IOException, InterruptedException {`
`64`	`64`	`validate("my-app-id", "secret", false);`
`65`	`65`	`}`
`66`	`66`
`@@ -83,12 +83,13 @@ public void testBadSecret() {`
`83`	`83`	`}`
`84`	`84`
`85`	`85`	`@Test`
`86`		`- public void testEncryption() throws IOException {`
	`86`	`+ public void testEncryption() throws IOException, InterruptedException {`
`87`	`87`	`validate("my-app-id", "secret", true);`
`88`	`88`	`}`
`89`	`89`
`90`	`90`	`/** Creates an ExternalShuffleClient and attempts to register with the server. */`
`91`		`- private void validate(String appId, String secretKey, boolean encrypt) throws IOException {`
	`91`	`+ private void validate(String appId, String secretKey, boolean encrypt)`
	`92`	`+ throws IOException, InterruptedException {`
`92`	`93`	`ExternalShuffleClient client =`
`93`	`94`	`new ExternalShuffleClient(conf, new TestSecretKeyHolder(appId, secretKey), true, encrypt);`
`94`	`95`	`client.init(appId);`