JAVA-2943: Prevent session leak with wrong keyspace name

Author: adutra

changelog/README.md

@@ -4,6 +4,7 @@
 
 ### 4.11.2 (in progress)
 
+- [bug] JAVA-2943: Prevent session leak with wrong keyspace name
 - [bug] JAVA-2938: OverloadedException message is misleading
 
 ### 4.11.1

core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java

@@ -384,16 +384,30 @@ private void init(CqlIdentifier keyspace) {
           .getTopologyMonitor()
           .init()
           .thenCompose(v -> metadataManager.refreshNodes())
-          .thenAccept(v -> afterInitialNodeListRefresh(keyspace))
-          .exceptionally(
-              error -> {
-                initFuture.completeExceptionally(error);
-                RunOrSchedule.on(adminExecutor, this::close);
-                return null;
+          .thenCompose(v -> checkProtocolVersion())
+          .thenCompose(v -> initialSchemaRefresh())
+          .thenCompose(v -> initializePools(keyspace))
+          .whenComplete(
+              (v, error) -> {
+                if (error == null) {
+                  LOG.debug("[{}] Initialization complete, ready", logPrefix);
+                  notifyListeners();
+                  initFuture.complete(DefaultSession.this);
+                } else {
+                  LOG.debug("[{}] Initialization failed, force closing", logPrefix, error);
+                  forceCloseAsync()
+                      .whenComplete(
+                          (v1, error1) -> {
+                            if (error1 != null) {
+                              error.addSuppressed(error1);
+                            }
+                            initFuture.completeExceptionally(error);
+                          });
+                }
               });
     }
 
-    private void afterInitialNodeListRefresh(CqlIdentifier keyspace) {
+    private CompletionStage<Void> checkProtocolVersion() {
       try {
         boolean protocolWasForced =
             context.getConfig().getDefaultProfile().isDefined(DefaultDriverOption.PROTOCOL_VERSION);
@@ -426,48 +440,39 @@ private void afterInitialNodeListRefresh(CqlIdentifier keyspace) {
                 bestVersion);
           }
         }
-        metadataManager
+        return CompletableFuture.completedFuture(null);
+      } catch (Throwable throwable) {
+        return CompletableFutures.failedFuture(throwable);
+      }
+    }
+
+    private CompletionStage<RefreshSchemaResult> initialSchemaRefresh() {
+      try {
+        return metadataManager
             .refreshSchema(null, false, true)
-            .whenComplete(
-                (metadata, error) -> {
-                  if (error != null) {
-                    Loggers.warnWithException(
-                        LOG,
-                        "[{}] Unexpected error while refreshing schema during initialization, "
-                            + "keeping previous version",
-                        logPrefix,
-                        error);
-                  }
-                  afterInitialSchemaRefresh(keyspace);
+            .exceptionally(
+                error -> {
+                  Loggers.warnWithException(
+                      LOG,
+                      "[{}] Unexpected error while refreshing schema during initialization, "
+                          + "proceeding without schema metadata",
+                      logPrefix,
+                      error);
+                  return null;
                 });
       } catch (Throwable throwable) {
-        initFuture.completeExceptionally(throwable);
+        return CompletableFutures.failedFuture(throwable);
       }
     }
 
-    private void afterInitialSchemaRefresh(CqlIdentifier keyspace) {
+    private CompletionStage<Void> initializePools(CqlIdentifier keyspace) {
       try {
         nodeStateManager.markInitialized();
         context.getLoadBalancingPolicyWrapper().init();
         context.getConfigLoader().onDriverInit(context);
-        LOG.debug("[{}] Initialization complete, ready", logPrefix);
-        poolManager
-            .init(keyspace)
-            .whenComplete(
-                (v, error) -> {
-                  if (error != null) {
-                    initFuture.completeExceptionally(error);
-                  } else {
-                    notifyListeners();
-                    initFuture.complete(DefaultSession.this);
-                  }
-                });
+        return poolManager.init(keyspace);
       } catch (Throwable throwable) {
-        forceCloseAsync()
-            .whenComplete(
-                (v, error) -> {
-                  initFuture.completeExceptionally(throwable);
-                });
+        return CompletableFutures.failedFuture(throwable);
       }
     }
 

integration-tests/src/test/java/com/datastax/oss/driver/core/SessionLeakIT.java

@@ -16,6 +16,7 @@
 package com.datastax.oss.driver.core;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Fail.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
@@ -25,7 +26,9 @@
 import ch.qos.logback.classic.Logger;
 import ch.qos.logback.classic.spi.ILoggingEvent;
 import ch.qos.logback.core.Appender;
+import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.InvalidKeyspaceException;
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
 import com.datastax.oss.driver.api.core.session.Session;
@@ -34,6 +37,7 @@
 import com.datastax.oss.driver.categories.IsolatedTests;
 import com.datastax.oss.driver.internal.core.session.DefaultSession;
 import com.datastax.oss.simulacron.common.cluster.ClusterSpec;
+import com.datastax.oss.simulacron.common.stubbing.PrimeDsl;
 import java.util.HashSet;
 import java.util.Set;
 import org.junit.Before;
@@ -103,4 +107,28 @@ public void should_warn_when_session_count_exceeds_threshold() {
     verify(appender, never()).doAppend(any());
     session.close();
   }
+
+  @Test
+  public void should_never_warn_when_session_init_fails() {
+    SIMULACRON_RULE
+        .cluster()
+        .prime(PrimeDsl.when("USE \"non_existent_keyspace\"").then(PrimeDsl.invalid("irrelevant")));
+    int threshold = 4;
+    // Set the config option explicitly, in case it gets overridden in the test application.conf:
+    DriverConfigLoader configLoader =
+        DriverConfigLoader.programmaticBuilder()
+            .withInt(DefaultDriverOption.SESSION_LEAK_THRESHOLD, threshold)
+            .build();
+    // Go over the threshold, no warnings expected
+    for (int i = 0; i < threshold + 1; i++) {
+      try (Session session =
+          SessionUtils.newSession(
+              SIMULACRON_RULE, CqlIdentifier.fromCql("non_existent_keyspace"), configLoader)) {
+        fail("Session %s should have failed to initialize", session.getName());
+      } catch (InvalidKeyspaceException e) {
+        assertThat(e.getMessage()).isEqualTo("Invalid keyspace non_existent_keyspace");
+      }
+    }
+    verify(appender, never()).doAppend(any());
+  }
 }