From d0fb42c2e15c8bb2d06bf72bbc628bc190589b99 Mon Sep 17 00:00:00 2001
From: Alexander Dejanovski <alex.dejanovski@datastax.com>
Date: Mon, 9 Dec 2024 14:37:25 +0100
Subject: [PATCH] Fix release lead on segments

---
 .../storage/MemoryStorageFacade.java             |  1 +
 .../memory/ReplicaLockManagerWithTtl.java        | 16 ++++++++++------
 .../service/RepairRunnerTest.java                |  2 +-
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/server/src/main/java/io/cassandrareaper/storage/MemoryStorageFacade.java b/src/server/src/main/java/io/cassandrareaper/storage/MemoryStorageFacade.java
index 93b7b3cb8..8be5b152b 100644
--- a/src/server/src/main/java/io/cassandrareaper/storage/MemoryStorageFacade.java
+++ b/src/server/src/main/java/io/cassandrareaper/storage/MemoryStorageFacade.java
@@ -322,6 +322,7 @@ public boolean renewRunningRepairsForNodes(UUID runId, UUID segmentId, Set<Strin
 
   @Override
   public boolean releaseRunningRepairsForNodes(UUID runId, UUID segmentId, Set<String> replicas) {
+    LOG.info("Releasing locks for runId: {}, segmentId: {}, replicas: {}", runId, segmentId, replicas);
     return repairRunLockManager.releaseRunningRepairsForNodes(runId, segmentId, replicas);
   }
 
diff --git a/src/server/src/main/java/io/cassandrareaper/storage/memory/ReplicaLockManagerWithTtl.java b/src/server/src/main/java/io/cassandrareaper/storage/memory/ReplicaLockManagerWithTtl.java
index bdb86f607..1485e1c11 100644
--- a/src/server/src/main/java/io/cassandrareaper/storage/memory/ReplicaLockManagerWithTtl.java
+++ b/src/server/src/main/java/io/cassandrareaper/storage/memory/ReplicaLockManagerWithTtl.java
@@ -43,13 +43,17 @@ public ReplicaLockManagerWithTtl(long ttlSeconds) {
     scheduler.scheduleAtFixedRate(this::cleanupExpiredLocks, 1, 1, TimeUnit.SECONDS);
   }
 
+  private String getReplicaLockKey(String replica, UUID runId) {
+    return replica + runId;
+  }
+
   public boolean lockRunningRepairsForNodes(UUID runId, UUID segmentId, Set<String> replicas) {
     lock.lock();
     try {
       long currentTime = System.currentTimeMillis();
       // Check if any replica is already locked by another runId
       for (String replica : replicas) {
-        LockInfo lockInfo = replicaLocks.get(replica + runId);
+        LockInfo lockInfo = replicaLocks.get(getReplicaLockKey(replica, runId));
         if (lockInfo != null && lockInfo.expirationTime > currentTime && lockInfo.runId.equals(runId)) {
           return false; // Replica is locked by another runId and not expired
         }
@@ -58,7 +62,7 @@ public boolean lockRunningRepairsForNodes(UUID runId, UUID segmentId, Set<String
       // Lock the replicas for the given runId and segmentId
       long expirationTime = currentTime + (ttlSeconds * 1000);
       for (String replica : replicas) {
-        replicaLocks.put(replica + runId, new LockInfo(runId, expirationTime));
+        replicaLocks.put(getReplicaLockKey(replica, runId), new LockInfo(runId, expirationTime));
       }
 
       // Update runId to segmentId mapping
@@ -76,7 +80,7 @@ public boolean renewRunningRepairsForNodes(UUID runId, UUID segmentId, Set<Strin
 
       // Check if all replicas are already locked by this runId
       for (String replica : replicas) {
-        LockInfo lockInfo = replicaLocks.get(replica + runId);
+        LockInfo lockInfo = replicaLocks.get(getReplicaLockKey(replica, runId));
         if (lockInfo == null || !lockInfo.runId.equals(runId) || lockInfo.expirationTime <= currentTime) {
           return false; // Some replica is not validly locked by this runId
         }
@@ -85,7 +89,7 @@ public boolean renewRunningRepairsForNodes(UUID runId, UUID segmentId, Set<Strin
       // Renew the lock by extending the expiration time
       long newExpirationTime = currentTime + (ttlSeconds * 1000);
       for (String replica : replicas) {
-        replicaLocks.put(replica + runId, new LockInfo(runId, newExpirationTime));
+        replicaLocks.put(getReplicaLockKey(replica, runId), new LockInfo(runId, newExpirationTime));
       }
 
       // Ensure the segmentId is linked to the runId
@@ -101,9 +105,9 @@ public boolean releaseRunningRepairsForNodes(UUID runId, UUID segmentId, Set<Str
     try {
       // Remove the lock for replicas
       for (String replica : replicas) {
-        LockInfo lockInfo = replicaLocks.get(replica);
+        LockInfo lockInfo = replicaLocks.get(getReplicaLockKey(replica, runId));
         if (lockInfo != null && lockInfo.runId.equals(runId)) {
-          replicaLocks.remove(replica);
+          replicaLocks.remove(getReplicaLockKey(replica, runId));
         }
       }
 
diff --git a/src/server/src/test/java/io/cassandrareaper/service/RepairRunnerTest.java b/src/server/src/test/java/io/cassandrareaper/service/RepairRunnerTest.java
index d10a2f2a1..09683eabd 100644
--- a/src/server/src/test/java/io/cassandrareaper/service/RepairRunnerTest.java
+++ b/src/server/src/test/java/io/cassandrareaper/service/RepairRunnerTest.java
@@ -805,7 +805,7 @@ protected JmxCassandraManagementProxy connectImpl(Node host) throws ReaperExcept
         run.with().runState(RepairRun.RunState.RUNNING).startTime(DateTime.now()).build(runId));
     context.repairManager.resumeRunningRepairRuns();
 
-    await().with().atMost(20, TimeUnit.SECONDS).until(() -> {
+    await().with().atMost(120, TimeUnit.SECONDS).until(() -> {
       return RepairRun.RunState.DONE == storage.getRepairRunDao().getRepairRun(runId).get().getRunState();
     });
   }