Close sessions properly in DefaultMetadataTabletMapIT

Moves additional sessions used in this test into try-with-resources
statements, ensuring that they are closed afterwards.

Author: Bouncheck

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

@@ -165,173 +165,181 @@ public static void setup() {
 
   @Test
   public void every_statement_should_deliver_tablet_info() {
-    Map<String, Supplier<CqlSession>> sessions = new HashMap<>();
-    sessions.put(
-        "REGULAR",
-        () -> CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build());
-    sessions.put(
-        "WITH_KEYSPACE",
-        () ->
+    try (CqlSession regularSession =
+            CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build();
+        CqlSession withKsSession =
             CqlSession.builder()
                 .addContactEndPoints(CCM_RULE.getContactPoints())
                 .withKeyspace(KEYSPACE_NAME)
-                .build());
-    sessions.put(
-        "USE_KEYSPACE",
-        () -> {
-          CqlSession s =
-              CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build();
-          s.execute("USE " + KEYSPACE_NAME);
-          return s;
-        });
-
-    Map<String, Function<CqlSession, Statement>> statements = new HashMap<>();
-    statements.put("SELECT_CONCRETE", s -> STMT_SELECT_CONCRETE);
-    statements.put("SELECT_PREPARED", s -> s.prepare(STMT_SELECT).bind(2, 2));
-    statements.put("SELECT_NO_KS_PREPARED", s -> s.prepare(STMT_SELECT_NO_KS).bind(2, 2));
-    statements.put("SELECT_CONCRETE_PREPARED", s -> s.prepare(STMT_SELECT_CONCRETE).bind());
-    statements.put("SELECT_PK_CONCRETE_PREPARED", s -> s.prepare(STMT_SELECT_PK_CONCRETE).bind(2));
-    statements.put("SELECT_CK_CONCRETE_PREPARED", s -> s.prepare(STMT_SELECT_CK_CONCRETE).bind(2));
-    statements.put("INSERT_CONCRETE", s -> STMT_INSERT_CONCRETE);
-    statements.put("INSERT_PREPARED", s -> s.prepare(STMT_INSERT).bind(2, 2));
-    statements.put("INSERT_NO_KS_PREPARED", s -> s.prepare(STMT_INSERT_NO_KS).bind(2, 2));
-    statements.put("INSERT_CONCRETE_PREPARED", s -> s.prepare(STMT_INSERT_CONCRETE).bind());
-    statements.put("INSERT_PK_CONCRETE_PREPARED", s -> s.prepare(STMT_INSERT_PK_CONCRETE).bind(2));
-    statements.put("INSERT_CK_CONCRETE_PREPARED", s -> s.prepare(STMT_INSERT_CK_CONCRETE).bind(2));
-    statements.put(
-        "INSERT_LWT_IF_NOT_EXISTS", s -> s.prepare(STMT_INSERT_LWT_IF_NOT_EXISTS).bind(2, 2));
-    statements.put("UPDATE_CONCRETE", s -> STMT_UPDATE_CONCRETE);
-    statements.put("UPDATE_PREPARED", s -> s.prepare(STMT_UPDATE).bind(2, 2));
-    statements.put("UPDATE_NO_KS_PREPARED", s -> s.prepare(STMT_UPDATE_NO_KS).bind(2, 2));
-    statements.put("UPDATE_CONCRETE_PREPARED", s -> s.prepare(STMT_UPDATE_CONCRETE).bind());
-    statements.put("UPDATE_PK_CONCRETE_PREPARED", s -> s.prepare(STMT_UPDATE_PK_CONCRETE).bind(2));
-    statements.put("UPDATE_CK_CONCRETE_PREPARED", s -> s.prepare(STMT_UPDATE_CK_CONCRETE).bind(2));
-    statements.put("UPDATE_LWT_IF_EXISTS", s -> s.prepare(STMT_UPDATE_LWT_IF_EXISTS).bind(2, 2));
-    statements.put("STMT_UPDATE_LWT_IF_VAL", s -> s.prepare(STMT_UPDATE_LWT_IF_VAL).bind(2, 2));
-    statements.put("DELETE_CONCRETE", s -> STMT_DELETE_CONCRETE);
-    statements.put("DELETE_PREPARED", s -> s.prepare(STMT_DELETE).bind(2, 2));
-    statements.put("DELETE_NO_KS_PREPARED", s -> s.prepare(STMT_DELETE_NO_KS).bind(2, 2));
-    statements.put("DELETE_CONCRETE_PREPARED", s -> s.prepare(STMT_DELETE_CONCRETE).bind());
-    statements.put("DELETE_PK_CONCRETE_PREPARED", s -> s.prepare(STMT_DELETE_PK_CONCRETE).bind(2));
-    statements.put("DELETE_CK_CONCRETE_PREPARED", s -> s.prepare(STMT_DELETE_CK_CONCRETE).bind(2));
-    statements.put("DELETE_LWT_IF_EXISTS", s -> s.prepare(STMT_DELETE_IF_EXISTS).bind(2, 2));
-
-    List<String> testErrors = new ArrayList<>();
-    for (Map.Entry<String, Supplier<CqlSession>> sessionEntry : sessions.entrySet()) {
-      for (Map.Entry<String, Function<CqlSession, Statement>> stmtEntry : statements.entrySet()) {
-        if (stmtEntry.getKey().contains("CONCRETE")
-            && !stmtEntry.getKey().contains("CK_CONCRETE")) {
-          // Scylla does not return tablet info for queries with PK built into query
-          continue;
-        }
-        if (stmtEntry.getKey().contains("LWT")) {
-          // LWT is not yet supported by scylla on tables with tablets
-          continue;
-        }
-        if (sessionEntry.getKey().equals("REGULAR") && stmtEntry.getKey().contains("NO_KS")) {
-          // Preparation of the statements without KS will fail on the session with no ks specified
-          continue;
-        }
-        CqlSession session = sessionEntry.getValue().get();
-        // Empty out tablets information
-        if (session.getMetadata().getTabletMap().isPresent()) {
-          session
-              .getMetadata()
-              .getTabletMap()
-              .get()
-              .removeByKeyspace(CqlIdentifier.fromCql(KEYSPACE_NAME));
-        }
-        Statement stmt;
-        try {
-          stmt = stmtEntry.getValue().apply(session);
-        } catch (Exception e) {
-          RuntimeException ex =
-              new RuntimeException(
+                .build();
+        CqlSession withUseKsSession =
+            CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build()) {
+      withUseKsSession.execute("USE " + KEYSPACE_NAME);
+      Map<String, Supplier<CqlSession>> sessions = new HashMap<>();
+      sessions.put("REGULAR", () -> regularSession);
+      sessions.put("WITH_KEYSPACE", () -> withKsSession);
+      sessions.put("USE_KEYSPACE", () -> withUseKsSession);
+
+      Map<String, Function<CqlSession, Statement>> statements = new HashMap<>();
+      statements.put("SELECT_CONCRETE", s -> STMT_SELECT_CONCRETE);
+      statements.put("SELECT_PREPARED", s -> s.prepare(STMT_SELECT).bind(2, 2));
+      statements.put("SELECT_NO_KS_PREPARED", s -> s.prepare(STMT_SELECT_NO_KS).bind(2, 2));
+      statements.put("SELECT_CONCRETE_PREPARED", s -> s.prepare(STMT_SELECT_CONCRETE).bind());
+      statements.put(
+          "SELECT_PK_CONCRETE_PREPARED", s -> s.prepare(STMT_SELECT_PK_CONCRETE).bind(2));
+      statements.put(
+          "SELECT_CK_CONCRETE_PREPARED", s -> s.prepare(STMT_SELECT_CK_CONCRETE).bind(2));
+      statements.put("INSERT_CONCRETE", s -> STMT_INSERT_CONCRETE);
+      statements.put("INSERT_PREPARED", s -> s.prepare(STMT_INSERT).bind(2, 2));
+      statements.put("INSERT_NO_KS_PREPARED", s -> s.prepare(STMT_INSERT_NO_KS).bind(2, 2));
+      statements.put("INSERT_CONCRETE_PREPARED", s -> s.prepare(STMT_INSERT_CONCRETE).bind());
+      statements.put(
+          "INSERT_PK_CONCRETE_PREPARED", s -> s.prepare(STMT_INSERT_PK_CONCRETE).bind(2));
+      statements.put(
+          "INSERT_CK_CONCRETE_PREPARED", s -> s.prepare(STMT_INSERT_CK_CONCRETE).bind(2));
+      statements.put(
+          "INSERT_LWT_IF_NOT_EXISTS", s -> s.prepare(STMT_INSERT_LWT_IF_NOT_EXISTS).bind(2, 2));
+      statements.put("UPDATE_CONCRETE", s -> STMT_UPDATE_CONCRETE);
+      statements.put("UPDATE_PREPARED", s -> s.prepare(STMT_UPDATE).bind(2, 2));
+      statements.put("UPDATE_NO_KS_PREPARED", s -> s.prepare(STMT_UPDATE_NO_KS).bind(2, 2));
+      statements.put("UPDATE_CONCRETE_PREPARED", s -> s.prepare(STMT_UPDATE_CONCRETE).bind());
+      statements.put(
+          "UPDATE_PK_CONCRETE_PREPARED", s -> s.prepare(STMT_UPDATE_PK_CONCRETE).bind(2));
+      statements.put(
+          "UPDATE_CK_CONCRETE_PREPARED", s -> s.prepare(STMT_UPDATE_CK_CONCRETE).bind(2));
+      statements.put("UPDATE_LWT_IF_EXISTS", s -> s.prepare(STMT_UPDATE_LWT_IF_EXISTS).bind(2, 2));
+      statements.put("STMT_UPDATE_LWT_IF_VAL", s -> s.prepare(STMT_UPDATE_LWT_IF_VAL).bind(2, 2));
+      statements.put("DELETE_CONCRETE", s -> STMT_DELETE_CONCRETE);
+      statements.put("DELETE_PREPARED", s -> s.prepare(STMT_DELETE).bind(2, 2));
+      statements.put("DELETE_NO_KS_PREPARED", s -> s.prepare(STMT_DELETE_NO_KS).bind(2, 2));
+      statements.put("DELETE_CONCRETE_PREPARED", s -> s.prepare(STMT_DELETE_CONCRETE).bind());
+      statements.put(
+          "DELETE_PK_CONCRETE_PREPARED", s -> s.prepare(STMT_DELETE_PK_CONCRETE).bind(2));
+      statements.put(
+          "DELETE_CK_CONCRETE_PREPARED", s -> s.prepare(STMT_DELETE_CK_CONCRETE).bind(2));
+      statements.put("DELETE_LWT_IF_EXISTS", s -> s.prepare(STMT_DELETE_IF_EXISTS).bind(2, 2));
+
+      List<String> testErrors = new ArrayList<>();
+      for (Map.Entry<String, Supplier<CqlSession>> sessionEntry : sessions.entrySet()) {
+        for (Map.Entry<String, Function<CqlSession, Statement>> stmtEntry : statements.entrySet()) {
+          if (stmtEntry.getKey().contains("CONCRETE")
+              && !stmtEntry.getKey().contains("CK_CONCRETE")) {
+            // Scylla does not return tablet info for queries with PK built into query
+            continue;
+          }
+          if (stmtEntry.getKey().contains("LWT")) {
+            // LWT is not yet supported by scylla on tables with tablets
+            continue;
+          }
+          if (sessionEntry.getKey().equals("REGULAR") && stmtEntry.getKey().contains("NO_KS")) {
+            // Preparation of the statements without KS will fail on the session with no ks
+            // specified
+            continue;
+          }
+          CqlSession session = sessionEntry.getValue().get();
+          // Empty out tablets information
+          if (session.getMetadata().getTabletMap().isPresent()) {
+            session
+                .getMetadata()
+                .getTabletMap()
+                .get()
+                .removeByKeyspace(CqlIdentifier.fromCql(KEYSPACE_NAME));
+          }
+          Statement stmt;
+          try {
+            stmt = stmtEntry.getValue().apply(session);
+          } catch (Exception e) {
+            RuntimeException ex =
+                new RuntimeException(
+                    String.format(
+                        "Failed to build statement %s on session %s",
+                        stmtEntry.getKey(), sessionEntry.getKey()));
+            ex.addSuppressed(e);
+            throw ex;
+          }
+          try {
+            if (!executeOnAllHostsAndReturnIfResultHasTabletsInfo(session, stmt)) {
+              testErrors.add(
                   String.format(
-                      "Failed to build statement %s on session %s",
+                      "Statement %s on session %s got no tablet info",
                       stmtEntry.getKey(), sessionEntry.getKey()));
-          ex.addSuppressed(e);
-          throw ex;
-        }
-        try {
-          if (!executeOnAllHostsAndReturnIfResultHasTabletsInfo(session, stmt)) {
+              continue;
+            }
+          } catch (Exception e) {
+            testErrors.add(
+                String.format(
+                    "Failed to execute statement %s on session %s: %s",
+                    stmtEntry.getKey(), sessionEntry.getKey(), e));
+            continue;
+          }
+          if (!waitSessionLearnedTabletInfo(session)) {
             testErrors.add(
                 String.format(
-                    "Statement %s on session %s got no tablet info",
+                    "Statement %s on session %s did not trigger session tablets update",
                     stmtEntry.getKey(), sessionEntry.getKey()));
             continue;
           }
-        } catch (Exception e) {
-          testErrors.add(
-              String.format(
-                  "Failed to execute statement %s on session %s: %s",
-                  stmtEntry.getKey(), sessionEntry.getKey(), e));
-          continue;
-        }
-        if (!waitSessionLearnedTabletInfo(session)) {
-          testErrors.add(
-              String.format(
-                  "Statement %s on session %s did not trigger session tablets update",
-                  stmtEntry.getKey(), sessionEntry.getKey()));
-          continue;
-        }
-        if (!checkIfRoutedProperly(session, stmt)) {
-          testErrors.add(
-              String.format(
-                  "Statement %s on session %s was routed to different nodes",
-                  stmtEntry.getKey(), sessionEntry.getKey()));
+          if (!checkIfRoutedProperly(session, stmt)) {
+            testErrors.add(
+                String.format(
+                    "Statement %s on session %s was routed to different nodes",
+                    stmtEntry.getKey(), sessionEntry.getKey()));
+          }
         }
       }
-    }
 
-    if (!testErrors.isEmpty()) {
-      throw new AssertionError(
-          String.format(
-              "Found queries that got no tablet info: \n%s", String.join("\n", testErrors)));
+      if (!testErrors.isEmpty()) {
+        throw new AssertionError(
+            String.format(
+                "Found queries that got no tablet info: \n%s", String.join("\n", testErrors)));
+      }
     }
   }
 
   @Test
   public void should_receive_each_tablet_exactly_once() {
-    CqlSession session =
-        CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build();
     int counter = 0;
-    PreparedStatement preparedStatement = session.prepare(STMT_INSERT);
-    for (int i = 1; i <= QUERIES; i++) {
-      if (executeAndReturnIfResultHasTabletsInfo(session, preparedStatement.bind(i, i))) {
-        counter++;
+    try (CqlSession session =
+        CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build(); ) {
+      PreparedStatement preparedStatement = session.prepare(STMT_INSERT);
+      for (int i = 1; i <= QUERIES; i++) {
+        if (executeAndReturnIfResultHasTabletsInfo(session, preparedStatement.bind(i, i))) {
+          counter++;
+        }
       }
+      Assert.assertEquals(INITIAL_TABLETS, counter);
+      assertSessionTabletMapIsFilled(session);
     }
-    Assert.assertEquals(INITIAL_TABLETS, counter);
-    assertSessionTabletMapIsFilled(session);
-    session.close();
-
-    session = CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build();
-    counter = 0;
-    preparedStatement = session.prepare(STMT_SELECT);
-    for (int i = 1; i <= QUERIES; i++) {
-      if (executeAndReturnIfResultHasTabletsInfo(session, preparedStatement.bind(i, i))) {
-        counter++;
+
+    try (CqlSession session =
+        CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build()) {
+      counter = 0;
+      PreparedStatement preparedStatement = session.prepare(STMT_SELECT);
+      for (int i = 1; i <= QUERIES; i++) {
+        if (executeAndReturnIfResultHasTabletsInfo(session, preparedStatement.bind(i, i))) {
+          counter++;
+        }
       }
-    }
 
-    LOG.debug("Ran first set of queries");
+      LOG.debug("Ran first set of queries");
 
-    // With enough queries we should hit a wrong node for each tablet exactly once.
-    Assert.assertEquals(INITIAL_TABLETS, counter);
-    assertSessionTabletMapIsFilled(session);
+      // With enough queries we should hit a wrong node for each tablet exactly once.
+      Assert.assertEquals(INITIAL_TABLETS, counter);
+      assertSessionTabletMapIsFilled(session);
 
-    // All tablet information should be available by now (unless for some reason cluster did sth on
-    // its own)
-    // We should not receive any tablet payloads now, since they are sent only on mismatch.
-    for (int i = 1; i <= QUERIES; i++) {
+      // All tablet information should be available by now (unless for some reason cluster did sth
+      // on
+      // its own)
+      // We should not receive any tablet payloads now, since they are sent only on mismatch.
+      for (int i = 1; i <= QUERIES; i++) {
 
-      ResultSet rs = session.execute(preparedStatement.bind(i, i));
-      Map<String, ByteBuffer> payload = rs.getExecutionInfo().getIncomingPayload();
+        ResultSet rs = session.execute(preparedStatement.bind(i, i));
+        Map<String, ByteBuffer> payload = rs.getExecutionInfo().getIncomingPayload();
 
-      if (payload.containsKey(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY)) {
-        throw new RuntimeException(
-            "Received non empty payload with tablets routing information: " + payload);
+        if (payload.containsKey(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY)) {
+          throw new RuntimeException(
+              "Received non empty payload with tablets routing information: " + payload);
+        }
       }
     }
   }