diff --git a/broker/src/main/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImpl.java b/broker/src/main/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImpl.java index b487b8757f4..9a3e5d2ca3b 100644 --- a/broker/src/main/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImpl.java +++ b/broker/src/main/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImpl.java @@ -79,24 +79,24 @@ public void updateClientChannel(String clientId, Channel channel) { } @Override - public void addPartialSubscription(String clientId, String group, String topic, Set lmqNameSet, + public synchronized void addPartialSubscription(String clientId, String group, String topic, Set lmqNameSet, OffsetOption offsetOption) { - long maxCount = brokerController.getBrokerConfig().getMaxLiteSubscriptionCount(); - if (getActiveSubscriptionNum() >= maxCount) { - // No need to check existence, if reach here, it must be new. - throw new LiteQuotaException("lite subscription quota exceeded " + maxCount); - } if (LiteMetadataUtil.isWildcardGroup(group, brokerController)) { throw new IllegalStateException("subscribe lite operation is not supported for this group"); } + if (brokerController.getBrokerConfig().getMaxLiteSubscriptionCount() <= 0 + && !lmqNameSet.isEmpty()) { + throw new LiteQuotaException("lite subscription quota exceeded 0"); + } + + Set activeLmqNames = lmqNameSet.stream() + .filter(lmqName -> liteLifecycleManager.isSubscriptionActive(topic, lmqName)) + .collect(Collectors.toSet()); + ClientGroup clientGroup = new ClientGroup(clientId, group); + ensureQuota(clientGroup, Collections.emptySet(), activeLmqNames); LiteSubscription thisSub = getOrCreateLiteSubscription(clientId, group, topic); - // Utilize existing string object - final ClientGroup clientGroup = new ClientGroup(clientId, thisSub.getGroup()); - for (String lmqName : lmqNameSet) { - if (!liteLifecycleManager.isSubscriptionActive(topic, lmqName)) { - continue; - } + for (String lmqName : activeLmqNames) { thisSub.addLiteTopic(lmqName); // First remove the old subscription if (LiteMetadataUtil.isSubLiteExclusive(group, brokerController)) { @@ -123,32 +123,37 @@ public void removePartialSubscription(String clientId, String group, String topi } @Override - public void addCompleteSubscription(String clientId, String group, String topic, Set lmqNameAll, long version) { + public synchronized void addCompleteSubscription(String clientId, String group, String topic, Set lmqNameAll, long version) { Set lmqNameNew; if (LiteMetadataUtil.isWildcardGroup(group, brokerController)) { lmqNameNew = Collections.singleton(mockLmqNameForWildcardGroup(topic, group)); - markWildcardGroup(topic, group); } else { lmqNameNew = lmqNameAll.stream() .filter(lmqName -> liteLifecycleManager.isSubscriptionActive(topic, lmqName)) .collect(Collectors.toSet()); } - LiteSubscription thisSub = getOrCreateLiteSubscription(clientId, group, topic); - Set lmqNamePrev = thisSub.getLiteTopicSet(); - // Find topics to remove (in current set but not in new set) + LiteSubscription existingSubscription = client2Subscription.get(clientId); + Set lmqNamePrev = existingSubscription == null ? Collections.emptySet() : existingSubscription.getLiteTopicSet(); + ClientGroup clientGroup = new ClientGroup(clientId, group); Set lmqNameRemove = lmqNamePrev.stream() .filter(lmqName -> !lmqNameNew.contains(lmqName)) .collect(Collectors.toSet()); + ensureQuota(clientGroup, lmqNameRemove, lmqNameNew); - ClientGroup clientGroup = new ClientGroup(clientId, thisSub.getGroup()); + if (LiteMetadataUtil.isWildcardGroup(group, brokerController)) { + markWildcardGroup(topic, group); + } + LiteSubscription thisSub = getOrCreateLiteSubscription(clientId, group, topic); + + ClientGroup subscriptionClientGroup = new ClientGroup(clientId, thisSub.getGroup()); lmqNameRemove.forEach(lmqName -> { thisSub.removeLiteTopic(lmqName); - removeTopicGroup(clientGroup, lmqName, false); + removeTopicGroup(subscriptionClientGroup, lmqName, false); }); lmqNameNew.forEach(lmqName -> { thisSub.addLiteTopic(lmqName); - addTopicGroup(clientGroup, lmqName); + addTopicGroup(subscriptionClientGroup, lmqName); }); // Tombstone operations only apply to exclusive groups. if (LiteMetadataUtil.isSubLiteExclusive(group, brokerController)) { @@ -281,6 +286,20 @@ protected void addTopicGroup(ClientGroup clientGroup, String lmqName) { } } + private void ensureQuota(ClientGroup clientGroup, Set lmqNameRemove, Set lmqNameNew) { + long maxCount = brokerController.getBrokerConfig().getMaxLiteSubscriptionCount(); + long removedCount = lmqNameRemove.stream().filter(lmqName -> containsClientGroup(lmqName, clientGroup)).count(); + long addedCount = lmqNameNew.stream().filter(lmqName -> !containsClientGroup(lmqName, clientGroup)).count(); + if ((long) getActiveSubscriptionNum() - removedCount + addedCount > maxCount) { + throw new LiteQuotaException("lite subscription quota exceeded " + maxCount); + } + } + + private boolean containsClientGroup(String lmqName, ClientGroup clientGroup) { + Set clientGroups = liteTopic2Group.get(lmqName); + return clientGroups != null && clientGroups.contains(clientGroup); + } + protected void removeTopicGroup(ClientGroup clientGroup, String lmqName, boolean resetOffset) { Set topicGroupSet = liteTopic2Group.get(lmqName); if (topicGroupSet == null) { @@ -502,4 +521,4 @@ public boolean hasExclusiveEvictionTombstone(String clientId, String lmqName) { return exclusiveEvictionTombstones.contains(clientId, lmqName); } -} \ No newline at end of file +} diff --git a/broker/src/test/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImplTest.java b/broker/src/test/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImplTest.java index 7645a470962..05570e72dcf 100644 --- a/broker/src/test/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImplTest.java +++ b/broker/src/test/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImplTest.java @@ -18,6 +18,7 @@ package org.apache.rocketmq.broker.lite; import io.netty.channel.Channel; +import java.util.Arrays; import java.util.Collections; import java.util.HashSet; import java.util.List; @@ -125,6 +126,41 @@ public void testAddPartialSubscription_QuotaExceeded() { }); } + @Test + public void testAddPartialSubscription_RejectsBatchThatExceedsQuota() { + when(mockBrokerConfig.getMaxLiteSubscriptionCount()).thenReturn(1L); + when(mockLifecycleManager.isSubscriptionActive(anyString(), anyString())).thenReturn(true); + + assertThrows(LiteQuotaException.class, () -> registry.addPartialSubscription( + "testClient", "testGroup", "testTopic", new HashSet<>(Arrays.asList("lmq1", "lmq2")), null)); + + assertEquals(0, registry.getActiveSubscriptionNum()); + } + + @Test + public void testAddCompleteSubscription_RejectsQuotaOverflow() { + when(mockBrokerConfig.getMaxLiteSubscriptionCount()).thenReturn(1L); + when(mockLifecycleManager.isSubscriptionActive(anyString(), anyString())).thenReturn(true); + + assertThrows(LiteQuotaException.class, () -> registry.addCompleteSubscription( + "testClient", "testGroup", "testTopic", new HashSet<>(Arrays.asList("lmq1", "lmq2")), 1L)); + + assertEquals(0, registry.getActiveSubscriptionNum()); + } + + @Test + public void testAddPartialSubscription_AllowsExistingSubscriptionAtQuota() { + when(mockBrokerConfig.getMaxLiteSubscriptionCount()).thenReturn(1L); + when(mockLifecycleManager.isSubscriptionActive(anyString(), anyString())).thenReturn(true); + Set subscriptions = Collections.singleton("lmq1"); + + registry.addPartialSubscription("testClient", "testGroup", "testTopic", subscriptions, null); + + registry.addPartialSubscription("testClient", "testGroup", "testTopic", subscriptions, null); + + assertEquals(1, registry.getActiveSubscriptionNum()); + } + /** * Test addPartialSubscription throws exception for wildcard group */