PaxosCleanupLocalCoordinator wait for transaction timeout before repairing

Patch by Ariel Weisberg; Reviewed by Benedict Elliott Smith for CASSANDRA-20585

Author: aweisberg

src/java/org/apache/cassandra/config/Config.java

@@ -1382,4 +1382,6 @@ public enum CQLStartTime
     // 3.x Cassandra Driver has its "read" timeout set to 12 seconds, default matches this.
     public DurationSpec.LongMillisecondsBound native_transport_timeout = new DurationSpec.LongMillisecondsBound("12s");
     public boolean enforce_native_deadline_for_hints = false;
+
+    public boolean paxos_repair_race_wait = true;
 }

src/java/org/apache/cassandra/config/DatabaseDescriptor.java

@@ -114,8 +114,8 @@
 import static org.apache.cassandra.config.CassandraRelevantProperties.ALLOCATE_TOKENS_FOR_KEYSPACE;
 import static org.apache.cassandra.config.CassandraRelevantProperties.ALLOW_UNLIMITED_CONCURRENT_VALIDATIONS;
 import static org.apache.cassandra.config.CassandraRelevantProperties.AUTO_BOOTSTRAP;
-import static org.apache.cassandra.config.CassandraRelevantProperties.CONFIG_LOADER;
 import static org.apache.cassandra.config.CassandraRelevantProperties.CHRONICLE_ANALYTICS_DISABLE;
+import static org.apache.cassandra.config.CassandraRelevantProperties.CONFIG_LOADER;
 import static org.apache.cassandra.config.CassandraRelevantProperties.DISABLE_STCS_IN_L0;
 import static org.apache.cassandra.config.CassandraRelevantProperties.INITIAL_TOKEN;
 import static org.apache.cassandra.config.CassandraRelevantProperties.IO_NETTY_TRANSPORT_ESTIMATE_SIZE_ON_SUBMIT;
@@ -5296,4 +5296,15 @@ public static void setRejectOutOfTokenRangeRequests(boolean enabled)
     {
         conf.reject_out_of_token_range_requests = enabled;
     }
+
+    public static boolean getPaxosRepairRaceWait()
+    {
+        return conf.paxos_repair_race_wait;
+    }
+
+    @VisibleForTesting
+    public static void setPaxosRepairRaceWait(boolean paxosRepairRaceWait)
+    {
+        conf.paxos_repair_race_wait = paxosRepairRaceWait;
+    }
 }

src/java/org/apache/cassandra/service/StorageService.java

@@ -351,7 +351,7 @@ private static int getSchemaDelay()
     @VisibleForTesting // this is used for dtests only, see CASSANDRA-18152
     public volatile boolean skipNotificationListeners = false;
 
-    private final java.util.function.Predicate<Keyspace> anyOutOfRangeOpsRecorded = 
+    private final java.util.function.Predicate<Keyspace> anyOutOfRangeOpsRecorded =
                 keyspace -> keyspace.metric.outOfRangeTokenReads.getCount() > 0
                             || keyspace.metric.outOfRangeTokenWrites.getCount() > 0
                             || keyspace.metric.outOfRangeTokenPaxosRequests.getCount() > 0;
@@ -7668,4 +7668,16 @@ public void setPrioritizeSAIOverLegacyIndex(boolean value)
     {
         DatabaseDescriptor.setPrioritizeSAIOverLegacyIndex(value);
     }
+
+    @Override
+    public void setPaxosRepairRaceWait(boolean paxosRepairRaceWait)
+    {
+        DatabaseDescriptor.setPaxosRepairRaceWait(paxosRepairRaceWait);
+    }
+
+    @Override
+    public boolean getPaxosRepairRaceWait()
+    {
+        return DatabaseDescriptor.getPaxosRepairRaceWait();
+    }
 }

src/java/org/apache/cassandra/service/StorageServiceMBean.java

@@ -1322,4 +1322,8 @@ public void enableAuditLog(String loggerName, String includedKeyspaces, String e
 
     boolean getPrioritizeSAIOverLegacyIndex();
     void setPrioritizeSAIOverLegacyIndex(boolean value);
+
+    void setPaxosRepairRaceWait(boolean paxosRepairCoordinatorWait);
+
+    boolean getPaxosRepairRaceWait();
 }

src/java/org/apache/cassandra/service/paxos/cleanup/PaxosCleanupLocalCoordinator.java

@@ -24,6 +24,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 
 import com.google.common.base.Preconditions;
+import com.google.common.util.concurrent.Uninterruptibles;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -40,9 +41,15 @@
 import org.apache.cassandra.service.paxos.PaxosRepair;
 import org.apache.cassandra.service.paxos.PaxosState;
 import org.apache.cassandra.service.paxos.uncommitted.UncommittedPaxosKey;
+import org.apache.cassandra.utils.Clock;
 import org.apache.cassandra.utils.CloseableIterator;
 import org.apache.cassandra.utils.concurrent.AsyncFuture;
 
+import static java.util.concurrent.TimeUnit.MICROSECONDS;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.apache.cassandra.config.DatabaseDescriptor.getCasContentionTimeout;
+import static org.apache.cassandra.config.DatabaseDescriptor.getWriteRpcTimeout;
 import static org.apache.cassandra.service.paxos.cleanup.PaxosCleanupSession.TIMEOUT_NANOS;
 
 public class PaxosCleanupLocalCoordinator extends AsyncFuture<PaxosCleanupResponse>
@@ -126,16 +133,18 @@ private void scheduleKeyRepairsOrFinish()
                 return;
             }
 
+            long txnTimeoutMicros = Math.max(getCasContentionTimeout(MICROSECONDS), getWriteRpcTimeout(MICROSECONDS));
+            boolean waitForCoordinator = DatabaseDescriptor.getPaxosRepairRaceWait();
             while (inflight.size() < parallelism && uncommittedIter.hasNext())
-                repairKey(uncommittedIter.next());
+                repairKey(uncommittedIter.next(), txnTimeoutMicros, waitForCoordinator);
 
         }
 
         if (inflight.isEmpty())
             finish();
     }
 
-    private boolean repairKey(UncommittedPaxosKey uncommitted)
+    private boolean repairKey(UncommittedPaxosKey uncommitted, long txnTimeoutMicros, boolean waitForCoordinator)
     {
         logger.trace("repairing {}", uncommitted);
         Preconditions.checkState(!inflight.containsKey(uncommitted.getKey()));
@@ -146,6 +155,9 @@ private boolean repairKey(UncommittedPaxosKey uncommitted)
         if (consistency == null)
             return false;
 
+        if (waitForCoordinator)
+            maybeWaitForOriginalCoordinator(uncommitted, txnTimeoutMicros);
+
         inflight.put(uncommitted.getKey(), tableRepairs.startOrGetOrQueue(uncommitted.getKey(), uncommitted.ballot(), uncommitted.getConsistencyLevel(), table, result -> {
             if (result.wasSuccessful())
                 onKeyFinish(uncommitted.getKey());
@@ -155,6 +167,24 @@ private boolean repairKey(UncommittedPaxosKey uncommitted)
         return true;
     }
 
+    /**
+     * Wait to repair things that are still potentially executing at the original coordinator to avoid
+     * causing timeouts. This should only have to happen at most a few times when the repair starts
+     */
+    private static void maybeWaitForOriginalCoordinator(UncommittedPaxosKey uncommitted, long txnTimeoutMicros)
+    {
+        long nowMicros = MILLISECONDS.toMicros(Clock.Global.currentTimeMillis());
+        long ballotElapsedMicros = nowMicros - uncommitted.ballot().unixMicros();
+        if (ballotElapsedMicros < 0 && Math.abs(ballotElapsedMicros) > SECONDS.toMicros(1))
+            logger.warn("Encountered ballot that is more than 1 second in the future, is there a clock sync issue? {}", uncommitted.ballot());
+        if (ballotElapsedMicros < txnTimeoutMicros)
+        {
+            long sleepMicros = txnTimeoutMicros - ballotElapsedMicros;
+            logger.info("Paxos auto repair encountered a potentially in progress ballot, sleeping {}us to allow the in flight operation to finish", sleepMicros);
+            Uninterruptibles.sleepUninterruptibly(sleepMicros, MICROSECONDS);
+        }
+    }
+
     private synchronized void onKeyFinish(DecoratedKey key)
     {
         if (!inflight.containsKey(key))