diff --git a/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/ConnectionManager.java b/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/ConnectionManager.java index 962e8ad8ec00..c2fb2f56d4eb 100644 --- a/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/ConnectionManager.java +++ b/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/ConnectionManager.java @@ -58,10 +58,23 @@ private static String readToSessionHash(Read read) { return readToClusterHash(read) + read.keyspace().get(); } - static Session getSession(Read read) { + static synchronized Session getSession(Read read) { + String clusterHash = readToClusterHash(read); + String sessionHash = readToSessionHash(read); + + Cluster cachedCluster = clusterMap.get(clusterHash); + + if (cachedCluster != null && cachedCluster.isClosed()) { + Session brokenSession = sessionMap.get(sessionHash); + if (brokenSession != null) { + sessionMap.remove(sessionHash, brokenSession); + } + // Removing broken cluster object + clusterMap.remove(clusterHash, cachedCluster); + } Cluster cluster = clusterMap.computeIfAbsent( - readToClusterHash(read), + clusterHash, k -> CassandraIO.getCluster( Objects.requireNonNull(read.hosts()), diff --git a/sdks/java/io/cassandra/src/test/java/org/apache/beam/sdk/io/cassandra/CassandraIOTest.java b/sdks/java/io/cassandra/src/test/java/org/apache/beam/sdk/io/cassandra/CassandraIOTest.java index f63c819d4202..93ba98af6745 100644 --- a/sdks/java/io/cassandra/src/test/java/org/apache/beam/sdk/io/cassandra/CassandraIOTest.java +++ b/sdks/java/io/cassandra/src/test/java/org/apache/beam/sdk/io/cassandra/CassandraIOTest.java @@ -19,6 +19,9 @@ import static junit.framework.TestCase.assertTrue; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNotSame; import static org.junit.Assert.assertNull; import com.datastax.driver.core.Cluster; @@ -1218,4 +1221,31 @@ public int hashCode() { return Objects.hashCode(tableColumn, indexColumn, valueColumn, data); } } + + @Test + public void testSessionEvictionOnClosedCluster() { + CassandraIO.Read readConfig = + CassandraIO.read() + .withHosts(Collections.singletonList(CASSANDRA_HOST)) + .withPort(cassandraPort) + .withKeyspace(CASSANDRA_KEYSPACE) + .withTable(CASSANDRA_TABLE); + + Session initialSession = ConnectionManager.getSession(readConfig); + Cluster initialCluster = initialSession.getCluster(); + + initialCluster.close(); + assertTrue("Cluster should be closed", initialCluster.isClosed()); + + Session newSession = ConnectionManager.getSession(readConfig); + Cluster newCluster = newSession.getCluster(); + + assertNotNull("New session should not be null", newSession); + assertFalse("New cluster should be open", newCluster.isClosed()); + + assertNotSame( + "ConnectionManager should create a new Session instance", initialSession, newSession); + assertNotSame( + "ConnectionManager should create a new Cluster instance", initialCluster, newCluster); + } }