This is an automated email from the ASF dual-hosted git repository.

iamaleksey pushed a commit to branch 20386
in repository https://gitbox.apache.org/repos/asf/cassandra.git

commit 876b4e171a72a3a98739fcfc9d41e54d9e430d70
Author: Aleksey Yeshchenko <[email protected]>
AuthorDate: Tue Aug 11 16:56:52 2026 +0100

    Improve shard drain logic to make double-release safe
---
 .../cassandra/replication/ForwardedWrite.java      | 50 ++++++++++++----------
 .../replication/MutationTrackingService.java       |  2 +-
 .../org/apache/cassandra/replication/Shard.java    | 26 ++++++-----
 3 files changed, 43 insertions(+), 35 deletions(-)

diff --git a/src/java/org/apache/cassandra/replication/ForwardedWrite.java 
b/src/java/org/apache/cassandra/replication/ForwardedWrite.java
index 27cb1c9eee..74eca34d48 100644
--- a/src/java/org/apache/cassandra/replication/ForwardedWrite.java
+++ b/src/java/org/apache/cassandra/replication/ForwardedWrite.java
@@ -224,23 +224,9 @@ public class ForwardedWrite
              * the coordinator, we need to abort here and now.
              */
             MutationId id = 
MutationTrackingService.instance().nextMutationId(keyspaceName, token);
-            Participants shardParticipants = 
MutationTrackingService.instance().getLogParticipants(id.asLogId());
-            Participants liveAndDownParticipants = 
Participants.merge(liveReplicas, downReplicas);
-            if (!shardParticipants.equals(liveAndDownParticipants))
-            {
-                MutationTrackingService.instance().completeLocalWrite(id);
-                TCMMetrics.instance.coordinatorBehindPlacements.mark();
-                String msg =
-                    format("Mutation id %s: shard participants %s disagree 
with plan replicas %s; coordinator must refresh and retry",
-                           id, shardParticipants, liveAndDownParticipants);
-                throw new CoordinatorBehindException(msg);
-            }
-            Mutation mutation = this.mutation.withMutationId(id);
-
-            // Do not wait for handler completion, since the coordinator is 
already waiting and we don't want to block the stage
-            LeaderCallback handler = new LeaderCallback(id, ackTo);
 
-            boolean applyLocally = false;
+            Mutation mutation;
+            LeaderCallback handler;
 
             // this DC replicas
             List<Replica> localDCReplicas = null;
@@ -251,14 +237,32 @@ public class ForwardedWrite
             // only need to create a Message for non-local writes
             Message<Mutation> message = null;
 
-            // Expensive, but easier to work with Replica than 
InetAddressAndPort for now
-            Int2ObjectHashMap<Replica> replicas = new 
Int2ObjectHashMap<>(liveReplicas.size(), 0.65f);
-            EndpointsForToken endpoints = writePlacements.get();
-            for (Replica replica : endpoints)
-                replicas.put(cm.directory.peerId(replica.endpoint()).id(), 
replica);
-
             try
             {
+                Participants shardParticipants = 
MutationTrackingService.instance().getLogParticipants(id.asLogId());
+                Participants liveAndDownParticipants = 
Participants.merge(liveReplicas, downReplicas);
+                if (!shardParticipants.equals(liveAndDownParticipants))
+                {
+                    TCMMetrics.instance.coordinatorBehindPlacements.mark();
+                    String msg =
+                        format("Mutation id %s: shard participants %s disagree 
with plan replicas %s; coordinator must refresh and retry",
+                               id, shardParticipants, liveAndDownParticipants);
+                    throw new CoordinatorBehindException(msg);
+                }
+
+                mutation = this.mutation.withMutationId(id);
+
+                // Do not wait for handler completion, since the coordinator 
is already waiting and we don't want to block the stage
+                handler = new LeaderCallback(id, ackTo);
+
+                boolean applyLocally = false;
+
+                // Expensive, but easier to work with Replica than 
InetAddressAndPort for now
+                Int2ObjectHashMap<Replica> replicas = new 
Int2ObjectHashMap<>(liveReplicas.size(), 0.65f);
+                EndpointsForToken endpoints = writePlacements.get();
+                for (Replica replica : endpoints)
+                    replicas.put(cm.directory.peerId(replica.endpoint()).id(), 
replica);
+
                 // For performance, Mutation caches serialized buffers that 
are computed lazily in serializedBuffer(). That
                 // computation is not synchronized however, and we will 
potentially call that method concurrently for each
                 // dispatched message (not that concurrent calls to 
serializedBuffer() are "unsafe" per se, just that they
@@ -307,7 +311,7 @@ public class ForwardedWrite
             }
             catch (Throwable t)
             {
-                
MutationTrackingService.instance().completeLocalWrite(mutation.id());
+                MutationTrackingService.instance().completeLocalWrite(id);
                 throw t;
             }
 
diff --git 
a/src/java/org/apache/cassandra/replication/MutationTrackingService.java 
b/src/java/org/apache/cassandra/replication/MutationTrackingService.java
index 56d9e1842e..22fb0dce96 100644
--- a/src/java/org/apache/cassandra/replication/MutationTrackingService.java
+++ b/src/java/org/apache/cassandra/replication/MutationTrackingService.java
@@ -463,7 +463,7 @@ public class MutationTrackingService implements 
MutationTrackingServiceMBean
         Shard shard = getShardNullable(id.asLogId());
         if (null == shard)
             throw new IllegalStateException(format("Shard for log %s was not 
found in log2ShardMap", id.asLogId()));
-        shard.completeLocalWrite();
+        shard.completeLocalWrite(id);
     }
 
     /**
diff --git a/src/java/org/apache/cassandra/replication/Shard.java 
b/src/java/org/apache/cassandra/replication/Shard.java
index 4da657e192..e2b22f653c 100644
--- a/src/java/org/apache/cassandra/replication/Shard.java
+++ b/src/java/org/apache/cassandra/replication/Shard.java
@@ -35,6 +35,7 @@ import javax.annotation.Nonnull;
 import com.google.common.base.Preconditions;
 
 import org.agrona.collections.IntArrayList;
+import org.jctools.maps.NonBlockingHashMap;
 import org.jctools.maps.NonBlockingHashMapLong;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -159,23 +160,25 @@ public class Shard
 
     /**
      * Incremented before this shard allocates a MutationId.
-     * Decremented once the mutation has applied or failed locally.
+     * Decremented once the id has been allocated and added to the pending 
writes set.
      * Used by shard sealing logic for drain() step.
      */
-    private final AtomicInteger pendingLocalWrites = new AtomicInteger();
+    private final AtomicInteger pendingIdAllocations = new AtomicInteger();
+    private final NonBlockingHashMap<MutationId, Boolean> pendingLocalWrites = 
new NonBlockingHashMap<>();
 
     @Nonnull
     MutationId nextMutationId()
     {
-        pendingLocalWrites.incrementAndGet();
+        pendingIdAllocations.incrementAndGet();
         try
         {
-            return nextId();
+            MutationId id = nextId();
+            pendingLocalWrites.put(id, true);
+            return id;
         }
-        catch (Throwable t)
+        finally
         {
-            pendingLocalWrites.decrementAndGet();
-            throw t;
+            pendingIdAllocations.decrementAndGet();
         }
     }
 
@@ -213,11 +216,11 @@ public class Shard
     }
 
     /**
-     * Must be called exactly once per {@code nextId()} invocation.
+     * Must be called at least once per {@code nextId()} invocation.
      */
-    void completeLocalWrite()
+    void completeLocalWrite(MutationId mutationId)
     {
-        pendingLocalWrites.decrementAndGet();
+        pendingLocalWrites.remove(mutationId);
     }
 
     void receivedWriteResponse(ShortMutationId mutationId, InetAddressAndPort 
fromHost)
@@ -252,6 +255,7 @@ public class Shard
     void finishWriting(Mutation mutation)
     {
         getOrCreate(mutation).finishWriting(mutation);
+        pendingLocalWrites.remove(mutation.id());
     }
 
     void addSummaryForKey(Token token, boolean includePending, 
MutationSummary.Builder builder)
@@ -479,7 +483,7 @@ public class Shard
 
     boolean isDrained()
     {
-        return state != State.ACTIVE && pendingLocalWrites.get() == 0;
+        return state != State.ACTIVE && pendingIdAllocations.get() == 0 && 
pendingLocalWrites.isEmpty();
     }
 
     /**


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to