Fix borrow accounting on keyspace setup failure

Author: dkropachev

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

@@ -900,7 +900,14 @@ 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 (RuntimeException e) {
+          targetKeyspace.compareAndSet(attempt, existingAttempt);
+          ksFuture.setException(e);
+          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,30 @@ ListenableFuture<Connection> borrowConnection(
       if (totalInFlightCount > currentCapacity) maybeSpawnNewConnection(shardId);
     }
 
-    return leastBusy.setKeyspaceAsync(manager.poolsState.keyspace);
+    final Connection borrowedConnection = leastBusy;
+    ListenableFuture<Connection> setKeyspaceFuture;
+    try {
+      setKeyspaceFuture = borrowedConnection.setKeyspaceAsync(manager.poolsState.keyspace);
+    } catch (BusyConnectionException e) {
+      borrowedConnection.release(true);
+      throw e;
+    } catch (RuntimeException e) {
+      borrowedConnection.release(false);
+      throw e;
+    }
+    Futures.addCallback(
+        setKeyspaceFuture,
+        new FutureCallback<Connection>() {
+          @Override
+          public void onSuccess(Connection connection) {}
+
+          @Override
+          public void onFailure(Throwable t) {
+            borrowedConnection.release(false);
+          }
+        },
+        MoreExecutors.directExecutor());
+    return setKeyspaceFuture;
   }
 
   private ListenableFuture<Connection> enqueue(

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

@@ -62,6 +62,7 @@
 import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.Uninterruptibles;
+import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.util.Collection;
@@ -434,6 +435,33 @@ public void should_adjust_connection_keyspace_on_dequeue_if_pool_state_is_differ
             .contains(
                 "Aborting attempt to set keyspace to 'newks' 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);
+    } finally {
+      cluster.close();
+    }
+  }
+
+  @Test(groups = "short")
+  public void should_restore_borrow_counts_when_set_keyspace_throws_synchronously()
+      throws Exception {
+    Cluster cluster = createClusterBuilder().build();
+    try {
+      HostConnectionPool pool = createPool(cluster, 1, 1);
+      Connection connection = pool.connections[0].get(0);
+
+      pool.manager.poolsState.setKeyspace("newkeyspace");
+      markDefunctWithoutClosing(connection);
+
+      try {
+        pool.borrowConnection(5000, MILLISECONDS, 0, null, null, null, null);
+        fail("Expected a ConnectionException");
+      } catch (ConnectionException e) {
+        assertThat(e.getMessage()).contains("Write attempt on defunct connection");
+      }
+
+      assertThat(connection.inFlight.get()).isEqualTo(0);
+      assertThat(pool.totalInFlight.get()).isEqualTo(0);
     } finally {
       cluster.close();
     }
@@ -1412,6 +1440,12 @@ private HostConnectionPool createPool(Cluster cluster, int coreConnections, int
     return sm.pools.get(host);
   }
 
+  private void markDefunctWithoutClosing(Connection connection) throws Exception {
+    Field isDefunct = Connection.class.getDeclaredField("isDefunct");
+    isDefunct.setAccessible(true);
+    ((AtomicBoolean) isDefunct.get(connection)).set(true);
+  }
+
   /**
    * This test uses a table named "Java349" with 1000 column and performs asynchronously 100k
    * insertions. While the insertions are being executed, the number of opened connection is