JAVA-1948: Close session properly when LBP fails to initialize

GregBestland · olim7t · commit bc84386793ed · 2018-10-05T15:17:57.000-07:00
diff --git a/changelog/README.md b/changelog/README.md
@@ -4,6 +4,7 @@
 
 ### 4.0.0-beta2 (in progress)
 
+- [bug] JAVA-1948: Close session properly when LBP fails to initialize
 - [improvement] JAVA-1949: Improve error message when contact points are wrong
 - [improvement] JAVA-1956: Add statementsCount accessor to BatchStatementBuilder
 - [bug] JAVA-1946: Ignore protocol version in equals comparison for UdtValue/TupleValue
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java b/core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java
@@ -366,7 +366,11 @@ private void afterInitialSchemaRefresh(CqlIdentifier keyspace) {
                   }
                 });
       } catch (Throwable throwable) {
-        initFuture.completeExceptionally(throwable);
+        forceCloseAsync()
+            .whenComplete(
+                (v, error) -> {
+                  initFuture.completeExceptionally(throwable);
+                });
       }
     }
 
diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/api/core/ConnectIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/api/core/ConnectIT.java
@@ -15,6 +15,9 @@
  */
 package com.datastax.oss.driver.api.core;
 
+import static com.datastax.oss.driver.api.testinfra.utils.ConditionChecker.checkThat;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static junit.framework.TestCase.fail;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
@@ -98,6 +101,30 @@ public void should_wait_for_contact_points_if_reconnection_enabled() throws Exce
     session.close();
   }
 
+  /**
+   * Test for JAVA-1948. This ensures that when the LBP initialization fails that any connections
+   * are cleaned up appropriately.
+   */
+  @Test
+  public void should_cleanup_on_lbp_init_failure() {
+    try {
+      DriverConfigLoader loader =
+          SessionUtils.configLoaderBuilder()
+              .without(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER)
+              .build();
+      CqlSession.builder()
+          .addContactPoints(simulacronRule.getContactPoints())
+          .withConfigLoader(loader)
+          .build();
+      fail("Should have thrown a DriverException for no DC with explicit contact point");
+    } catch (DriverException ignored) {
+    }
+    // One second should be plenty of time for connections to close server side
+    checkThat(() -> simulacronRule.cluster().getConnections().getConnections().isEmpty())
+        .before(1, SECONDS)
+        .becomesTrue();
+  }
+
   @SuppressWarnings("unchecked")
   private CompletionStage<? extends Session> newSessionAsync(
       SimulacronRule serverRule, DriverConfigLoader loader) {

Original file line number	Diff line number	Diff line change
`@@ -366,7 +366,11 @@ private void afterInitialSchemaRefresh(CqlIdentifier keyspace) {`
`366`	`366`	`}`
`367`	`367`	`});`
`368`	`368`	`} catch (Throwable throwable) {`
`369`		`- initFuture.completeExceptionally(throwable);`
	`369`	`+ forceCloseAsync()`
	`370`	`+ .whenComplete(`
	`371`	`+ (v, error) -> {`
	`372`	`+ initFuture.completeExceptionally(throwable);`
	`373`	`+ });`
`370`	`374`	`}`
`371`	`375`	`}`
`372`	`376`