Skip to content

Fix stale broken cluster connection in CassandraIO (#39788) #39803

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sharan-malyala wants to merge 2 commits into
base: release-2.76
Choose a base branch
from
+45 −2
Open

Fix stale broken cluster connection in CassandraIO (#39788) #39803

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
eea67a5
Fix stale broken cluster connection in CassandraIO (#39788)
sharan-malyala Aug 18, 2026
09280a0
Set getSession to synchronized to avoid race conditions
sharan-malyala Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/ConnectionManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()) {

@Abacn Abacn Aug 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inside this if clause it modified sessionMap and clusterMap. Even though they are concurrent hash map, is there a risk of racing?

@sharan-malyala sharan-malyala Aug 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we dont cleanup the session and just re-create cluster object, the broken session will be returned and used by the work items leading to failure.

@Abacn Abacn Aug 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was thinking about the following scenario:

  1. A session corrupted and cluster closed

  2. Caller A runs getSession(), found cachedCluster.isClosed

  3. Caller B also runs getSession(), found cachedCluster.isClosed

  4. Caller A cleaned up session and cluster cache, recreated cluster connection

  5. Caller B clean up session again, because it's already inside the if clause. But it should not do so as the cache is current valid, may cause session leak.

Basically it seems the whole getSession now needs to be synchronized.

@Abacn Abacn Aug 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, this is already merged in #39788 however a race appears possible

@sharan-malyala sharan-malyala Aug 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will check further.

@stankiewicz stankiewicz Aug 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

apologies, you are correct, getSession should be synchronized in such change. Sharan, can you apply fix to master as well?

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()),
Expand Down
30 changes: 30 additions & 0 deletions sdks/java/io/cassandra/src/test/java/org/apache/beam/sdk/io/cassandra/CassandraIOTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -1218,4 +1221,31 @@ public int hashCode() {
return Objects.hashCode(tableColumn, indexColumn, valueColumn, data);
}
}

@Test
public void testSessionEvictionOnClosedCluster() {
CassandraIO.Read<String> readConfig =
CassandraIO.<String>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);
}
}
Loading