Rely on snapshotted session infos on StreamResultFuture.maybeComplete…

… to avoid races patch by Massimiliano Tomassi; reviewed by Sergio Bossa and by ZhaoYang for CASSANDRA-15667
apache · May 25, 2020 · a4b6deb · a4b6deb
1 parent ea202ce
commit a4b6deb
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 9 deletions.
diff --git a/CHANGES.txt b/CHANGES.txt
@@ -1,4 +1,5 @@
 3.0.21
+ * Rely on snapshotted session infos on StreamResultFuture.maybeComplete to avoid race conditions (CASSANDRA-15667)
  * EmptyType doesn't override writeValue so could attempt to write bytes when expected not to (CASSANDRA-15790)
  * Fix index queries on partition key columns when some partitions contains only static data (CASSANDRA-13666)
  * Avoid creating duplicate rows during major upgrades (CASSANDRA-15789)

diff --git a/src/java/org/apache/cassandra/streaming/StreamCoordinator.java b/src/java/org/apache/cassandra/streaming/StreamCoordinator.java
@@ -232,7 +232,7 @@ public boolean hasActiveSessions()
             for (StreamSession session : streamSessions.values())
             {
                 StreamSession.State state = session.state();
-                if (state != StreamSession.State.COMPLETE && state != StreamSession.State.FAILED)
+                if (!state.isFinalState())
                     return true;
             }
             return false;
@@ -245,6 +245,7 @@ public StreamSession getOrCreateNextSession(InetAddress peer, InetAddress connec
             {
                 StreamSession session = new StreamSession(peer, connecting, factory, streamSessions.size(), keepSSTableLevel, isIncremental);
                 streamSessions.put(++lastReturned, session);
+                sessionInfos.put(lastReturned, session.getSessionInfo());
                 return session;
             }
             // get
@@ -277,6 +278,7 @@ public StreamSession getOrCreateSessionById(InetAddress peer, int id, InetAddres
             {
                 session = new StreamSession(peer, connecting, factory, id, keepSSTableLevel, isIncremental);
                 streamSessions.put(id, session);
+                sessionInfos.put(id, session.getSessionInfo());
             }
             return session;
         }

diff --git a/src/java/org/apache/cassandra/streaming/StreamResultFuture.java b/src/java/org/apache/cassandra/streaming/StreamResultFuture.java
@@ -202,7 +202,7 @@ synchronized void fireStreamEvent(StreamEvent event)
 
     private synchronized void maybeComplete()
     {
-        if (!coordinator.hasActiveSessions())
+        if (finishedAllSessions())
         {
             StreamState finalState = getCurrentState();
             if (finalState.hasFailedSession())
@@ -217,4 +217,14 @@ private synchronized void maybeComplete()
             }
         }
     }
+
+    /**
+     * We can't use {@link StreamCoordinator#hasActiveSessions()} directly because {@link this#maybeComplete()}
+     * relies on the snapshotted state from {@link StreamCoordinator} and not the {@link StreamSession} state
+     * directly (CASSANDRA-15667), otherwise inconsistent snapshotted states may lead to completion races.
+     */
+    private boolean finishedAllSessions()
+    {
+        return coordinator.getAllSessionInfo().stream().allMatch(s -> s.state.isFinalState());
+    }
 }
diff --git a/src/java/org/apache/cassandra/streaming/StreamSession.java b/src/java/org/apache/cassandra/streaming/StreamSession.java
@@ -152,12 +152,27 @@ public class StreamSession implements IEndpointStateChangeSubscriber
 
     public static enum State
     {
-        INITIALIZED,
-        PREPARING,
-        STREAMING,
-        WAIT_COMPLETE,
-        COMPLETE,
-        FAILED,
+        INITIALIZED(false),
+        PREPARING(false),
+        STREAMING(false),
+        WAIT_COMPLETE(false),
+        COMPLETE(true),
+        FAILED(true);
+
+        private final boolean finalState;
+
+        State(boolean finalState)
+        {
+            this.finalState = finalState;
+        }
+
+        /**
+         * @return true if current state is final, either COMPLETE OR FAILED.
+         */
+        public boolean isFinalState()
+        {
+            return finalState;
+        }
     }
 
     private volatile State state = State.INITIALIZED;
@@ -300,7 +315,7 @@ public synchronized void addTransferRanges(String keyspace, Collection<Range<Tok
 
     private void failIfFinished()
     {
-        if (state() == State.COMPLETE || state() == State.FAILED)
+        if (state().isFinalState())
             throw new RuntimeException(String.format("Stream %s is finished with state %s", planId(), state().name()));
     }