[ISSUE #10750] Fix POP lock cleanup race

broker/src/main/java/org/apache/rocketmq/broker/pop/PopConsumerLockService.java

@@ -16,16 +16,13 @@
  */
 package org.apache.rocketmq.broker.pop;
 
-import java.util.Iterator;
 import java.util.Map;
-import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.rocketmq.common.KeyBuilder;
 import org.apache.rocketmq.common.PopAckConstants;
 import org.apache.rocketmq.common.constant.LoggerName;
-import org.apache.rocketmq.common.utils.ConcurrentHashMapUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -42,8 +39,13 @@ public PopConsumerLockService(long timeout) {
     }
 
     public boolean tryLock(String key) {
-        return Objects.requireNonNull(ConcurrentHashMapUtils.computeIfAbsent(lockTable,
-            key, s -> new TimedLock())).tryLock();
+        AtomicBoolean locked = new AtomicBoolean(false);
+        lockTable.compute(key, (k, currentLock) -> {
+            TimedLock lock = currentLock == null ? new TimedLock() : currentLock;
+            locked.set(lock.tryLock());
+            return lock;
+        });
+        return locked.get();
     }
 
     public boolean tryLock(String groupId, String topicId) {
@@ -69,13 +71,22 @@ public boolean isLockTimeout(String groupId, String topicId) {
     }
 
     public void removeTimeout() {
-        Iterator<Map.Entry<String, TimedLock>> iterator = lockTable.entrySet().iterator();
-        while (iterator.hasNext()) {
-            Map.Entry<String, TimedLock> entry = iterator.next();
-            if (System.currentTimeMillis() - entry.getValue().getLockTime() > timeout) {
+        for (Map.Entry<String, TimedLock> entry : lockTable.entrySet()) {
+            if (System.currentTimeMillis() - entry.getValue().getLockTime() <= timeout) {
+                continue;
+            }
+
+            TimedLock[] removedLock = new TimedLock[1];
+            lockTable.computeIfPresent(entry.getKey(), (key, currentLock) -> {
+                if (System.currentTimeMillis() - currentLock.getLockTime() > timeout) {
+                    removedLock[0] = currentLock;
+                    return null;
+                }
+                return currentLock;
+            });
+            if (removedLock[0] != null) {
                 log.info("PopConsumerLockService remove timeout lock, " +
-                    "key={}, locked={}", entry.getKey(), entry.getValue().lock.get());
-                iterator.remove();
+                    "key={}, locked={}", entry.getKey(), removedLock[0].lock.get());
             }
         }
     }
@@ -105,4 +116,4 @@ public long getLockTime() {
             return lockTime;
         }
     }
-}
+}

broker/src/test/java/org/apache/rocketmq/broker/pop/PopConsumerLockServiceTest.java

@@ -18,6 +18,10 @@
 
 import java.lang.reflect.Field;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import org.apache.rocketmq.common.PopAckConstants;
 import org.junit.Assert;
@@ -57,4 +61,59 @@ public void consumerLockTest() throws NoSuchFieldException, IllegalAccessExcepti
 
         Assert.assertEquals(0, table.size());
     }
-}
+
+    @Test
+    @SuppressWarnings("unchecked")
+    public void removeTimeoutShouldNotRemoveReacquiredLock() throws Exception {
+        String key = "groupId" + PopAckConstants.SPLIT + "topicId";
+        PopConsumerLockService lockService =
+            new PopConsumerLockService(TimeUnit.MINUTES.toMillis(2));
+
+        Field tableField = PopConsumerLockService.class.getDeclaredField("lockTable");
+        tableField.setAccessible(true);
+        Map<String, PopConsumerLockService.TimedLock> table =
+            (Map<String, PopConsumerLockService.TimedLock>) tableField.get(lockService);
+
+        CountDownLatch timeoutObserved = new CountDownLatch(1);
+        CountDownLatch continueCleanup = new CountDownLatch(1);
+        PopConsumerLockService.TimedLock expiredLock = new PopConsumerLockService.TimedLock() {
+            @Override
+            public long getLockTime() {
+                long observedLockTime = super.getLockTime();
+                timeoutObserved.countDown();
+                try {
+                    if (!continueCleanup.await(5, TimeUnit.SECONDS)) {
+                        throw new AssertionError("Timed out waiting to continue lock cleanup");
+                    }
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new AssertionError(e);
+                }
+                return observedLockTime;
+            }
+        };
+
+        Field lockTimeField = PopConsumerLockService.TimedLock.class.getDeclaredField("lockTime");
+        lockTimeField.setAccessible(true);
+        lockTimeField.setLong(expiredLock, 0L);
+        table.put(key, expiredLock);
+
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        try {
+            Future<?> cleanup = executor.submit(lockService::removeTimeout);
+            Assert.assertTrue("Cleanup did not inspect the expired lock",
+                timeoutObserved.await(5, TimeUnit.SECONDS));
+
+            Assert.assertTrue("The expired lock should be reacquired before cleanup continues",
+                lockService.tryLock(key));
+            continueCleanup.countDown();
+            cleanup.get(5, TimeUnit.SECONDS);
+
+            Assert.assertFalse("Cleanup removed the reacquired lock and allowed a second holder",
+                lockService.tryLock(key));
+        } finally {
+            continueCleanup.countDown();
+            executor.shutdownNow();
+        }
+    }
+}