Skip to content

fix: enforce Lite subscription quota #10896

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Aias00 wants to merge 1 commit into from
+76 −21
Closed

fix: enforce Lite subscription quota #10896

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 40 additions & 21 deletions broker/src/main/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,24 +79,24 @@ public void updateClientChannel(String clientId, Channel channel) {
}

@Override
public void addPartialSubscription(String clientId, String group, String topic, Set<String> lmqNameSet,
public synchronized void addPartialSubscription(String clientId, String group, String topic, Set<String> 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<String> 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)) {
Expand All @@ -123,32 +123,37 @@ public void removePartialSubscription(String clientId, String group, String topi
}

@Override
public void addCompleteSubscription(String clientId, String group, String topic, Set<String> lmqNameAll, long version) {
public synchronized void addCompleteSubscription(String clientId, String group, String topic, Set<String> lmqNameAll, long version) {
Set<String> 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<String> lmqNamePrev = thisSub.getLiteTopicSet();
// Find topics to remove (in current set but not in new set)
LiteSubscription existingSubscription = client2Subscription.get(clientId);
Set<String> lmqNamePrev = existingSubscription == null ? Collections.emptySet() : existingSubscription.getLiteTopicSet();
ClientGroup clientGroup = new ClientGroup(clientId, group);
Set<String> 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)) {
Expand Down Expand Up @@ -281,6 +286,20 @@ protected void addTopicGroup(ClientGroup clientGroup, String lmqName) {
}
}

private void ensureQuota(ClientGroup clientGroup, Set<String> lmqNameRemove, Set<String> 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<ClientGroup> clientGroups = liteTopic2Group.get(lmqName);
return clientGroups != null && clientGroups.contains(clientGroup);
}

protected void removeTopicGroup(ClientGroup clientGroup, String lmqName, boolean resetOffset) {
Set<ClientGroup> topicGroupSet = liteTopic2Group.get(lmqName);
if (topicGroupSet == null) {
Expand Down Expand Up @@ -502,4 +521,4 @@ public boolean hasExclusiveEvictionTombstone(String clientId, String lmqName) {
return exclusiveEvictionTombstones.contains(clientId, lmqName);
}

}
}
36 changes: 36 additions & 0 deletions broker/src/test/java/org/apache/rocketmq/broker/lite/LiteSubscriptionRegistryImplTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<String> 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
*/
Expand Down
Loading