Handle keyspace setup failures during borrow

Author: dkropachev

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

@@ -900,7 +900,17 @@ ListenableFuture<Connection> setKeyspaceAsync(final String keyspace)
         logger.debug("{} Setting keyspace {}", this, keyspace);
         // Note: we quote the keyspace below, because the name is the one coming from Cassandra, so
         // it's in the right case already
-        Future future = write(new Requests.Query("USE \"" + keyspace + '"'));
+        Future future;
+        try {
+          future = write(new Requests.Query("USE \"" + keyspace + '"'));
+        } catch (ConnectionException | BusyConnectionException e) {
+          targetKeyspace.compareAndSet(attempt, defaultKeyspaceAttempt);
+          ksFuture.setException(e);
+          return ksFuture;
+        } catch (RuntimeException e) {
+          targetKeyspace.compareAndSet(attempt, defaultKeyspaceAttempt);
+          throw e;
+        }
         Futures.addCallback(
             future,
             new FutureCallback<Message.Response>() {

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

@@ -27,6 +27,7 @@
 import static com.datastax.driver.core.Connection.State.TRASHED;
 
 import com.datastax.driver.core.exceptions.AuthenticationException;
+import com.datastax.driver.core.exceptions.BusyConnectionException;
 import com.datastax.driver.core.exceptions.BusyPoolException;
 import com.datastax.driver.core.exceptions.ConnectionException;
 import com.datastax.driver.core.exceptions.UnsupportedProtocolVersionException;
@@ -611,7 +612,22 @@ ListenableFuture<Connection> borrowConnection(
       if (totalInFlightCount > currentCapacity) maybeSpawnNewConnection(shardId);
     }
 
-    return leastBusy.setKeyspaceAsync(manager.poolsState.keyspace);
+    final Connection borrowedConnection = leastBusy;
+    ListenableFuture<Connection> setKeyspaceFuture =
+        borrowedConnection.setKeyspaceAsync(manager.poolsState.keyspace);
+    Futures.addCallback(
+        setKeyspaceFuture,
+        new FutureCallback<Connection>() {
+          @Override
+          public void onSuccess(Connection connection) {}
+
+          @Override
+          public void onFailure(Throwable t) {
+            borrowedConnection.release(t instanceof BusyConnectionException);
+          }
+        },
+        MoreExecutors.directExecutor());
+    return setKeyspaceFuture;
   }
 
   private ListenableFuture<Connection> enqueue(
@@ -698,8 +714,14 @@ private void dequeue(final Connection connection) {
         pendingBorrowCount.decrementAndGet();
         // Ensure that the keyspace set on the connection is the one set on the pool state, in the
         // general case it will be.
-        ListenableFuture<Connection> setKeyspaceFuture =
-            connection.setKeyspaceAsync(manager.poolsState.keyspace);
+        ListenableFuture<Connection> setKeyspaceFuture;
+        try {
+          setKeyspaceFuture = connection.setKeyspaceAsync(manager.poolsState.keyspace);
+        } catch (ConnectionException | BusyConnectionException e) {
+          pendingBorrow.setException(e);
+          connection.inFlight.decrementAndGet();
+          continue;
+        }
         // Slight optimization, if the keyspace was already correct the future will be complete, so
         // simply complete it here.
         if (setKeyspaceFuture.isDone()) {

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

@@ -393,6 +393,85 @@ public void should_adjust_connection_keyspace_on_dequeue_if_pool_state_is_differ
     }
   }
 
+  /**
+   * Ensures that if keyspace setup fails while borrowing a connection, the borrow future fails and
+   * pool accounting is restored.
+   *
+   * @test_category connection:connection_pool
+   */
+  @Test(groups = "short")
+  public void should_restore_pool_accounting_when_keyspace_setup_fails_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");
+      ConnectionException failure =
+          new ConnectionException(pool.host.getEndPoint(), "Write attempt on defunct connection");
+      doReturn(Futures.<Connection>immediateFailedFuture(failure))
+          .when(connection)
+          .setKeyspaceAsync("newkeyspace");
+
+      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);
+      }
+      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.
+   *
+   * @test_category connection:connection_pool
+   */
+  @Test(groups = "short")
+  public void should_fail_keyspace_future_and_clear_attempt_when_connection_is_defunct()
+      throws Exception {
+    Cluster cluster = createClusterBuilder().build();
+    try {
+      HostConnectionPool pool = createPool(cluster, 1, 1);
+      Connection connection = pool.connections[0].get(0);
+
+      connection.defunct(
+          new ConnectionException(pool.host.getEndPoint(), "Test defunct connection"));
+
+      ListenableFuture<Connection> firstAttempt = connection.setKeyspaceAsync("newkeyspace");
+      ListenableFuture<Connection> secondAttempt = connection.setKeyspaceAsync("newkeyspace");
+
+      assertThat(firstAttempt).isNotSameAs(secondAttempt);
+      try {
+        Uninterruptibles.getUninterruptibly(firstAttempt, 5, TimeUnit.SECONDS);
+        fail("Should have failed keyspace attempt");
+      } catch (ExecutionException e) {
+        assertThat(e.getCause())
+            .isInstanceOf(ConnectionException.class)
+            .hasMessageContaining("Write attempt on defunct connection");
+      }
+      try {
+        Uninterruptibles.getUninterruptibly(secondAttempt, 5, TimeUnit.SECONDS);
+        fail("Should have failed keyspace attempt");
+      } catch (ExecutionException e) {
+        assertThat(e.getCause())
+            .isInstanceOf(ConnectionException.class)
+            .hasMessageContaining("Write attempt on defunct connection");
+      }
+    } finally {
+      cluster.close();
+    }
+  }
+
   /**
    * Ensures that on borrowConnection if a set keyspace attempt is in progress on that connection
    * for a different keyspace than the pool state that the borrowConnection future returned is
@@ -501,6 +580,9 @@ public void should_adjust_connection_keyspace_on_dequeue_if_pool_state_is_differ
             .contains(
                 "Aborting attempt to set keyspace to 'newkeyspace' since there is already an in flight attempt to set keyspace to 'slowks'.");
       }
+      assertThat(connection.inFlight.get()).isEqualTo(0);
+      assertThat(pool.totalInFlight.get()).isEqualTo(0);
+      assertThat(pool.pendingBorrowCount.get()).isEqualTo(0);
     } finally {
       MockRequest.completeAll(requests);
       cluster.close();