diff --git a/broker/src/main/java/org/apache/rocketmq/broker/pop/PopConsumerLockService.java b/broker/src/main/java/org/apache/rocketmq/broker/pop/PopConsumerLockService.java index 066db7192ae..addf2a3949c 100644 --- a/broker/src/main/java/org/apache/rocketmq/broker/pop/PopConsumerLockService.java +++ b/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> iterator = lockTable.entrySet().iterator(); - while (iterator.hasNext()) { - Map.Entry entry = iterator.next(); - if (System.currentTimeMillis() - entry.getValue().getLockTime() > timeout) { + for (Map.Entry 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; } } -} \ No newline at end of file +} diff --git a/broker/src/test/java/org/apache/rocketmq/broker/pop/PopConsumerLockServiceTest.java b/broker/src/test/java/org/apache/rocketmq/broker/pop/PopConsumerLockServiceTest.java index b5af2f31798..9da340144ce 100644 --- a/broker/src/test/java/org/apache/rocketmq/broker/pop/PopConsumerLockServiceTest.java +++ b/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()); } -} \ No newline at end of file + + @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 table = + (Map) 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(); + } + } +}