Handle runtime keyspace setup failures as failed futures

Author: dkropachev

driver-core/src/main/java/com/datastax/driver/core/Connection.java

@@ -913,7 +913,8 @@ ListenableFuture<Connection> setKeyspaceAsync(final String keyspace) {
           return ksFuture;
         } catch (RuntimeException e) {
           targetKeyspace.compareAndSet(attempt, defaultKeyspaceAttempt);
-          throw e;
+          ksFuture.setException(e);
+          return ksFuture;
         }
         Futures.addCallback(
             future,

driver-core/src/test/java/com/datastax/driver/core/HostConnectionPoolTest.java

@@ -462,6 +462,41 @@ public void should_restore_pool_accounting_when_keyspace_setup_fails_on_borrow()
     }
   }
 
+  /**
+   * Ensures that an unexpected synchronous failure while starting keyspace setup is reported
+   * through the borrow future and does not leak pool accounting.
+   *
+   * @test_category connection:connection_pool
+   */
+  @Test(groups = "short")
+  public void should_restore_pool_accounting_when_keyspace_setup_throws_on_borrow()
+      throws Exception {
+    Cluster cluster = createClusterBuilder().build();
+    try {
+      HostConnectionPool pool = createPool(cluster, 1, 1);
+      Connection connection = spy(pool.connections[0].get(0));
+      pool.connections[0].set(0, connection);
+
+      pool.manager.poolsState.setKeyspace("newkeyspace");
+      RuntimeException failure = new RuntimeException("Unexpected keyspace setup failure");
+      doThrow(failure).when(connection).write(any(Message.Request.class));
+
+      MockRequest request = MockRequest.send(pool);
+
+      try {
+        Uninterruptibles.getUninterruptibly(request.connectionFuture, 5, TimeUnit.SECONDS);
+        fail("Should have failed to borrow connection");
+      } catch (ExecutionException e) {
+        assertThat(e.getCause()).isSameAs(failure);
+      }
+      verify(connection).write(any(Message.Request.class));
+      assertThat(connection.inFlight.get()).isEqualTo(0);
+      assertThat(pool.totalInFlight.get()).isEqualTo(0);
+    } finally {
+      cluster.close();
+    }
+  }
+
   /**
    * Ensures that a synchronous write failure while setting the keyspace is reported through the
    * returned future and does not leave the failed keyspace attempt in-flight.