3.x: fail fast on permanent keyspace setup failures

dkropachev · dkropachev · commit 087fe8ce70f7 · 2026-07-02T13:53:16.000-04:00
Keep keyspace setup failures on the borrow path from corrupting pool accounting, and surface validation failures from the internal USE query as request failures instead of treating them as host failures.

Validation errors such as missing keyspaces or authorization problems are now returned as typed query errors without defuncting the connection or exhausting the query plan. Transient connection and busy-pool failures continue through the existing retry/next-host path.

Tests cover failed keyspace future cleanup, pool counter restoration, non-defunct validation failures, and fail-fast request behavior.
diff --git a/driver-core/src/main/java/com/datastax/driver/core/Connection.java b/driver-core/src/main/java/com/datastax/driver/core/Connection.java
@@ -34,6 +34,7 @@
 import com.datastax.driver.core.exceptions.DriverInternalError;
 import com.datastax.driver.core.exceptions.FrameTooLongException;
 import com.datastax.driver.core.exceptions.OperationTimedOutException;
+import com.datastax.driver.core.exceptions.QueryValidationException;
 import com.datastax.driver.core.exceptions.TransportException;
 import com.datastax.driver.core.exceptions.UnsupportedProtocolVersionException;
 import com.datastax.driver.core.utils.MoreFutures;
@@ -900,7 +901,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>() {
@@ -915,7 +926,12 @@ public void onSuccess(Message.Response response) {
                   targetKeyspace.compareAndSet(attempt, defaultKeyspaceAttempt);
                   if (response.type == ERROR) {
                     Responses.Error error = (Responses.Error) response;
-                    ksFuture.setException(defunct(error.asException(endPoint)));
+                    DriverException exception = error.asException(endPoint);
+                    if (exception instanceof QueryValidationException) {
+                      ksFuture.setException(exception);
+                    } else {
+                      ksFuture.setException(defunct(exception));
+                    }
                   } else {
                     ksFuture.setException(
                         defunct(
diff --git a/driver-core/src/main/java/com/datastax/driver/core/HostConnectionPool.java b/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()) {
diff --git a/driver-core/src/main/java/com/datastax/driver/core/RequestHandler.java b/driver-core/src/main/java/com/datastax/driver/core/RequestHandler.java
@@ -31,6 +31,7 @@
 import com.datastax.driver.core.exceptions.NoHostAvailableException;
 import com.datastax.driver.core.exceptions.OperationTimedOutException;
 import com.datastax.driver.core.exceptions.OverloadedException;
+import com.datastax.driver.core.exceptions.QueryValidationException;
 import com.datastax.driver.core.exceptions.ReadFailureException;
 import com.datastax.driver.core.exceptions.ReadTimeoutException;
 import com.datastax.driver.core.exceptions.ServerError;
@@ -468,6 +469,10 @@ public void onSuccess(Connection connection) {
 
             @Override
             public void onFailure(Throwable t) {
+              if (t instanceof QueryValidationException) {
+                setFinalException(null, (QueryValidationException) t);
+                return;
+              }
               if (t instanceof BusyPoolException) {
                 logError(host.getEndPoint(), t);
               } else {
diff --git a/driver-core/src/test/java/com/datastax/driver/core/HostConnectionPoolTest.java b/driver-core/src/test/java/com/datastax/driver/core/HostConnectionPoolTest.java
@@ -50,6 +50,7 @@
 import com.datastax.driver.core.exceptions.BusyPoolException;
 import com.datastax.driver.core.exceptions.ConnectionException;
 import com.datastax.driver.core.exceptions.DriverException;
+import com.datastax.driver.core.exceptions.InvalidQueryException;
 import com.datastax.driver.core.exceptions.NoHostAvailableException;
 import com.datastax.driver.core.policies.ConstantReconnectionPolicy;
 import com.google.common.base.Function;
@@ -62,6 +63,9 @@
 import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.Uninterruptibles;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import io.netty.util.CharsetUtil;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.util.Collection;
@@ -134,6 +138,16 @@ private void assertBorrowedConnection(
     assertBorrowedConnections(requests, Collections.singletonList(expectedConnection));
   }
 
+  private static Responses.Error errorResponse(ExceptionCode code, String message) {
+    ByteBuf body = Unpooled.buffer();
+    body.writeInt(code.value);
+    byte[] messageBytes = message.getBytes(CharsetUtil.UTF_8);
+    body.writeShort(messageBytes.length);
+    body.writeBytes(messageBytes);
+    return Responses.Error.decoder.decode(
+        body, ProtocolVersion.V4, CodecRegistry.DEFAULT_INSTANCE, ProtocolFeatureStore.EMPTY);
+  }
+
   /**
    * Ensures that if a fixed-sized pool has filled its core connections and reached its maximum
    * number of enqueued requests, then borrowConnection will fail instead of creating a new
@@ -393,6 +407,158 @@ 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 keyspace validation errors are surfaced as query errors and do not defunct the
+   * connection.
+   *
+   * @test_category connection:connection_pool
+   */
+  @Test(groups = "short")
+  public void should_not_defunct_connection_when_keyspace_setup_fails_with_validation_error()
+      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);
+      SettableConnectionFuture useFuture =
+          new SettableConnectionFuture(new Requests.Query("USE \"missingkeyspace\""));
+
+      doReturn(useFuture).when(connection).write(any(Message.Request.class));
+      ListenableFuture<Connection> keyspaceFuture = connection.setKeyspaceAsync("missingkeyspace");
+      useFuture.setResponse(
+          errorResponse(ExceptionCode.INVALID, "Keyspace missingkeyspace does not exist"));
+
+      try {
+        Uninterruptibles.getUninterruptibly(keyspaceFuture, 5, TimeUnit.SECONDS);
+        fail("Should have failed keyspace attempt");
+      } catch (ExecutionException e) {
+        assertThat(e.getCause()).isInstanceOf(InvalidQueryException.class);
+      }
+      assertThat(connection.isDefunct()).isFalse();
+      assertThat(pool.opened()).isEqualTo(1);
+    } finally {
+      cluster.close();
+    }
+  }
+
+  /**
+   * Ensures that permanent keyspace setup failures fail the request directly instead of being
+   * retried across hosts as connection failures.
+   *
+   * @test_category connection:connection_pool
+   */
+  @Test(groups = "short")
+  public void should_fail_fast_when_keyspace_setup_fails_with_validation_error() throws Exception {
+    Cluster cluster = createClusterBuilder().build();
+    try {
+      SessionManager session = (SessionManager) cluster.connect();
+      Host host = TestUtils.findHost(cluster, 1);
+      HostConnectionPool pool = session.pools.get(host);
+      Connection connection = spy(pool.connections[0].get(0));
+      pool.connections[0].set(0, connection);
+      InvalidQueryException failure =
+          new InvalidQueryException(
+              pool.host.getEndPoint(), "Keyspace missingkeyspace does not exist");
+      session.poolsState.setKeyspace("missingkeyspace");
+      doReturn(Futures.<Connection>immediateFailedFuture(failure))
+          .when(connection)
+          .setKeyspaceAsync("missingkeyspace");
+      activityClient.clearAllRecordedActivity();
+
+      try {
+        session.execute("select * from tbl");
+        fail("Should have failed the request with the keyspace validation error");
+      } catch (InvalidQueryException e) {
+        // expected
+      }
+      assertThat(connection.isDefunct()).isFalse();
+      assertThat(pool.opened()).isEqualTo(1);
+      assertThat(activityClient.retrieveQueries()).extractingResultOf("getQuery").isEmpty();
+    } 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 +667,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();
@@ -1743,4 +1912,14 @@ public void onFailure(Throwable t) {
       return future;
     }
   }
+
+  private static class SettableConnectionFuture extends Connection.Future {
+    SettableConnectionFuture(Message.Request request) {
+      super(request);
+    }
+
+    void setResponse(Message.Response response) {
+      super.set(response);
+    }
+  }
 }
