diff --git a/broker/src/main/java/org/apache/rocketmq/broker/longpolling/LmqPullRequestHoldService.java b/broker/src/main/java/org/apache/rocketmq/broker/longpolling/LmqPullRequestHoldService.java index eddaee706a9..a481903b7ff 100644 --- a/broker/src/main/java/org/apache/rocketmq/broker/longpolling/LmqPullRequestHoldService.java +++ b/broker/src/main/java/org/apache/rocketmq/broker/longpolling/LmqPullRequestHoldService.java @@ -54,10 +54,7 @@ public void checkHoldRequest() { LOGGER.error("check hold request failed. topic={}, queueId={}", topic, queueId, e); } if (MixAll.isLmq(topic)) { - ManyPullRequest mpr = pullRequestTable.get(key); - if (mpr == null || mpr.getPullRequestList() == null || mpr.getPullRequestList().isEmpty()) { - pullRequestTable.remove(key); - } + pullRequestTable.computeIfPresent(key, (k, mpr) -> mpr.isEmpty() ? null : mpr); } } } diff --git a/broker/src/main/java/org/apache/rocketmq/broker/longpolling/PullRequestHoldService.java b/broker/src/main/java/org/apache/rocketmq/broker/longpolling/PullRequestHoldService.java index 7dbc9e4fd86..545f7f84278 100644 --- a/broker/src/main/java/org/apache/rocketmq/broker/longpolling/PullRequestHoldService.java +++ b/broker/src/main/java/org/apache/rocketmq/broker/longpolling/PullRequestHoldService.java @@ -21,6 +21,7 @@ import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.function.Consumer; import org.apache.rocketmq.broker.BrokerController; import org.apache.rocketmq.common.ServiceThread; import org.apache.rocketmq.common.SystemClock; @@ -44,17 +45,16 @@ public PullRequestHoldService(final BrokerController brokerController) { public void suspendPullRequest(final String topic, final int queueId, final PullRequest pullRequest) { String key = this.buildKey(topic, queueId); - ManyPullRequest mpr = this.pullRequestTable.get(key); - if (null == mpr) { - mpr = new ManyPullRequest(); - ManyPullRequest prev = this.pullRequestTable.putIfAbsent(key, mpr); - if (prev != null) { - mpr = prev; - } - } - pullRequest.getRequestCommand().setSuspended(true); - mpr.addPullRequest(pullRequest); + this.addPullRequest(key, mpr -> mpr.addPullRequest(pullRequest)); + } + + private void addPullRequest(final String key, final Consumer addOperation) { + this.pullRequestTable.compute(key, (k, current) -> { + ManyPullRequest mpr = current == null ? new ManyPullRequest() : current; + addOperation.accept(mpr); + return mpr; + }); } private String buildKey(final String topic, final int queueId) { @@ -177,7 +177,7 @@ public void notifyMessageArriving(final String topic, final int queueId, final l } if (!replayList.isEmpty()) { - mpr.addPullRequest(replayList); + this.addPullRequest(key, current -> current.addPullRequest(replayList)); } } } diff --git a/broker/src/test/java/org/apache/rocketmq/broker/longpolling/LmqPullRequestHoldServiceTest.java b/broker/src/test/java/org/apache/rocketmq/broker/longpolling/LmqPullRequestHoldServiceTest.java new file mode 100644 index 00000000000..8f87ddf6b9e --- /dev/null +++ b/broker/src/test/java/org/apache/rocketmq/broker/longpolling/LmqPullRequestHoldServiceTest.java @@ -0,0 +1,244 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.rocketmq.broker.longpolling; + +import java.util.concurrent.ConcurrentHashMap; +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 java.util.function.BiFunction; +import org.apache.rocketmq.broker.BrokerController; +import org.apache.rocketmq.common.MixAll; +import org.apache.rocketmq.remoting.protocol.RemotingCommand; +import org.apache.rocketmq.store.MessageStore; +import org.junit.Test; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class LmqPullRequestHoldServiceTest { + private static final String TOPIC = MixAll.LMQ_PREFIX + "cleanup-race"; + private static final int QUEUE_ID = 0; + private static final String KEY = TOPIC + "@" + QUEUE_ID; + + @Test + public void testConcurrentSuspendRemainsReachableDuringEmptyBucketCleanup() throws Exception { + BrokerController brokerController = mock(BrokerController.class); + MessageStore messageStore = mock(MessageStore.class); + when(brokerController.getMessageStore()).thenReturn(messageStore); + when(messageStore.getMaxOffsetInQueue(TOPIC, QUEUE_ID)).thenReturn(0L); + + BlockingCleanupMap pullRequestTable = new BlockingCleanupMap(KEY); + pullRequestTable.put(KEY, new ManyPullRequest()); + LmqPullRequestHoldService service = new LmqPullRequestHoldService(brokerController); + service.pullRequestTable = pullRequestTable; + + PullRequest pullRequest = mock(PullRequest.class); + when(pullRequest.getRequestCommand()).thenReturn(mock(RemotingCommand.class)); + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future suspendFuture = executor.submit(() -> { + try { + if (!pullRequestTable.cleanupEntered.await(5, TimeUnit.SECONDS)) { + throw new AssertionError("cleanup did not reach the remove point"); + } + service.suspendPullRequest(TOPIC, QUEUE_ID, pullRequest); + ManyPullRequest mappedBucket = service.pullRequestTable.get(KEY); + assertNotNull("suspend should keep a reachable bucket before cleanup", mappedBucket); + assertTrue("suspend should add the request before cleanup continues", + mappedBucket.getPullRequestList().contains(pullRequest)); + } finally { + pullRequestTable.allowCleanup.countDown(); + } + return null; + }); + + try { + service.checkHoldRequest(); + suspendFuture.get(5, TimeUnit.SECONDS); + + ManyPullRequest mappedBucket = service.pullRequestTable.get(KEY); + assertNotNull("concurrently suspended request should retain a reachable bucket", mappedBucket); + assertTrue("concurrently suspended request should remain reachable", + mappedBucket.getPullRequestList().contains(pullRequest)); + } finally { + pullRequestTable.allowCleanup.countDown(); + executor.shutdownNow(); + } + } + + @Test + public void testSuspendDoesNotAppendToDetachedBucket() throws Exception { + BrokerController brokerController = mock(BrokerController.class); + BlockingLookupMap pullRequestTable = new BlockingLookupMap(KEY); + pullRequestTable.put(KEY, new ManyPullRequest()); + LmqPullRequestHoldService service = new LmqPullRequestHoldService(brokerController); + service.pullRequestTable = pullRequestTable; + + PullRequest pullRequest = mock(PullRequest.class); + when(pullRequest.getRequestCommand()).thenReturn(mock(RemotingCommand.class)); + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future suspendFuture = executor.submit( + () -> service.suspendPullRequest(TOPIC, QUEUE_ID, pullRequest)); + + try { + assertTrue("suspend should select the mapped bucket", + pullRequestTable.bucketSelected.await(5, TimeUnit.SECONDS)); + pullRequestTable.computeIfPresent(KEY, (key, current) -> current.isEmpty() ? null : current); + pullRequestTable.allowAppend.countDown(); + suspendFuture.get(5, TimeUnit.SECONDS); + + ManyPullRequest mappedBucket = pullRequestTable.get(KEY); + assertNotNull("suspended request should retain a reachable bucket", mappedBucket); + assertTrue("suspended request should remain reachable from the table", + mappedBucket.getPullRequestList().contains(pullRequest)); + } finally { + pullRequestTable.allowAppend.countDown(); + executor.shutdownNow(); + } + } + + @Test + public void testReplayRemainsReachableAfterConcurrentEmptyBucketCleanup() throws Exception { + BrokerController brokerController = mock(BrokerController.class); + MessageStore messageStore = mock(MessageStore.class); + when(brokerController.getMessageStore()).thenReturn(messageStore); + when(messageStore.getMaxOffsetInQueue(TOPIC, QUEUE_ID)).thenReturn(0L); + + LmqPullRequestHoldService service = new LmqPullRequestHoldService(brokerController); + CountDownLatch replayEvaluationStarted = new CountDownLatch(1); + CountDownLatch allowReplay = new CountDownLatch(1); + PullRequest pullRequest = mock(PullRequest.class); + when(pullRequest.getRequestCommand()).thenReturn(mock(RemotingCommand.class)); + when(pullRequest.getPullFromThisOffset()).thenAnswer(invocation -> { + replayEvaluationStarted.countDown(); + if (!allowReplay.await(5, TimeUnit.SECONDS)) { + throw new AssertionError("timed out waiting to replay the pull request"); + } + return 0L; + }); + when(pullRequest.getSuspendTimestamp()).thenReturn(System.currentTimeMillis()); + when(pullRequest.getTimeoutMillis()).thenReturn(TimeUnit.HOURS.toMillis(1)); + service.suspendPullRequest(TOPIC, QUEUE_ID, pullRequest); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future notificationFuture = executor.submit( + () -> service.notifyMessageArriving(TOPIC, QUEUE_ID, 0L)); + + try { + assertTrue("notification should start evaluating the detached request", + replayEvaluationStarted.await(5, TimeUnit.SECONDS)); + assertTrue("notification should have cleared the mapped bucket before evaluation", + service.pullRequestTable.get(KEY).isEmpty()); + + service.checkHoldRequest(); + assertFalse("cleanup should remove the empty mapped bucket", + service.pullRequestTable.containsKey(KEY)); + + allowReplay.countDown(); + notificationFuture.get(5, TimeUnit.SECONDS); + + ManyPullRequest replayBucket = service.pullRequestTable.get(KEY); + assertNotNull("replayed request should restore a reachable bucket", replayBucket); + assertTrue("replayed request should remain reachable from the table", + replayBucket.getPullRequestList().contains(pullRequest)); + } finally { + allowReplay.countDown(); + executor.shutdownNow(); + } + } + + private static class BlockingLookupMap extends ConcurrentHashMap { + private final String blockedKey; + private final CountDownLatch bucketSelected = new CountDownLatch(1); + private final CountDownLatch allowAppend = new CountDownLatch(1); + + BlockingLookupMap(String blockedKey) { + this.blockedKey = blockedKey; + } + + @Override + public ManyPullRequest get(Object key) { + ManyPullRequest current = super.get(key); + if (blockedKey.equals(key) && current != null) { + bucketSelected.countDown(); + try { + if (!allowAppend.await(5, TimeUnit.SECONDS)) { + throw new AssertionError("timed out waiting to append to the selected bucket"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + return current; + } + + @Override + public ManyPullRequest compute(String key, + BiFunction remappingFunction) { + ManyPullRequest result = super.compute(key, remappingFunction); + if (blockedKey.equals(key)) { + bucketSelected.countDown(); + } + return result; + } + } + + private static class BlockingCleanupMap extends ConcurrentHashMap { + private final String blockedKey; + private final CountDownLatch cleanupEntered = new CountDownLatch(1); + private final CountDownLatch allowCleanup = new CountDownLatch(1); + + BlockingCleanupMap(String blockedKey) { + this.blockedKey = blockedKey; + } + + @Override + public ManyPullRequest remove(Object key) { + if (blockedKey.equals(key)) { + awaitCleanup(); + } + return super.remove(key); + } + + @Override + public ManyPullRequest computeIfPresent(String key, + BiFunction remappingFunction) { + if (blockedKey.equals(key)) { + awaitCleanup(); + } + return super.computeIfPresent(key, remappingFunction); + } + + private void awaitCleanup() { + cleanupEntered.countDown(); + try { + if (!allowCleanup.await(5, TimeUnit.SECONDS)) { + throw new AssertionError("timed out waiting to clean up the empty bucket"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + } +}