diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/processor/ClientProcessor.java b/proxy/src/main/java/org/apache/rocketmq/proxy/processor/ClientProcessor.java index c73e66416da..b1d56544546 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/processor/ClientProcessor.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/processor/ClientProcessor.java @@ -115,12 +115,14 @@ public CompletableFuture syncLiteSubscription(ProxyContext ctx, LiteSubscriptionDTO liteSubscriptionDTO, long timeoutMillis ) { try { + LiteSubscriptionAction action = liteSubscriptionDTO.getAction(); validateLiteBindTopic(ctx, liteSubscriptionDTO.getGroup(), liteSubscriptionDTO.getTopic()); - if (CollectionUtils.isNotEmpty(liteSubscriptionDTO.getLiteTopicSet())) { + if (CollectionUtils.isNotEmpty(liteSubscriptionDTO.getLiteTopicSet()) + && shouldValidateLiteSubscriptionQuota(action)) { validateLiteSubscriptionQuota(ctx, liteSubscriptionDTO.getGroup(), liteSubscriptionDTO.getLiteTopicSet().size()); } - if (LiteSubscriptionAction.PARTIAL_ADD == liteSubscriptionDTO.getAction()) { + if (LiteSubscriptionAction.PARTIAL_ADD == action) { if (!syncLiteSubscriptionRateLimiter.tryAcquire()) { String msg = String.format("Too many syncLiteSubscription requests, topic=%s, group=%s, clientId=%s", liteSubscriptionDTO.getTopic(), liteSubscriptionDTO.getGroup(), ctx.getClientID()); @@ -139,6 +141,10 @@ public CompletableFuture syncLiteSubscription(ProxyContext ctx, } } + private boolean shouldValidateLiteSubscriptionQuota(LiteSubscriptionAction action) { + return LiteSubscriptionAction.PARTIAL_ADD == action || LiteSubscriptionAction.COMPLETE_ADD == action; + } + public ClientChannelInfo findConsumerChannel( ProxyContext ctx, String consumerGroup, diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/processor/ClientProcessorTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/processor/ClientProcessorTest.java index 6644341e551..0568a7839fc 100644 --- a/proxy/src/test/java/org/apache/rocketmq/proxy/processor/ClientProcessorTest.java +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/processor/ClientProcessorTest.java @@ -19,10 +19,14 @@ import java.util.HashSet; import java.util.Set; +import java.util.concurrent.CompletableFuture; +import org.apache.rocketmq.common.lite.LiteSubscriptionAction; +import org.apache.rocketmq.common.lite.LiteSubscriptionDTO; import org.apache.rocketmq.proxy.common.ProxyContext; import org.apache.rocketmq.proxy.config.ConfigurationManager; import org.apache.rocketmq.proxy.grpc.v2.common.GrpcProxyException; import org.apache.rocketmq.proxy.service.ServiceManager; +import org.apache.rocketmq.proxy.service.lite.LiteSubscriptionService; import org.apache.rocketmq.remoting.protocol.heartbeat.MessageModel; import org.apache.rocketmq.remoting.protocol.heartbeat.SubscriptionData; import org.apache.rocketmq.remoting.protocol.subscription.SubscriptionGroupConfig; @@ -36,6 +40,10 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @RunWith(MockitoJUnitRunner.class) @@ -53,6 +61,9 @@ public class ClientProcessorTest { @Mock private SubscriptionGroupConfig groupConfig; + @Mock + private LiteSubscriptionService liteSubscriptionService; + private ClientProcessor clientProcessor; @Before @@ -195,6 +206,31 @@ public void testValidateLiteSubscriptionQuota_exceedsQuota_throwsException() { assertTrue(exception.getMessage().contains("lite subscription quota exceeded")); } + @Test + public void testSyncLiteSubscription_partialRemoveSkipsQuotaValidation() { + String group = "group"; + String topic = "topic"; + Set liteTopicSet = new HashSet<>(); + for (int i = 0; i < 400; i++) { + liteTopicSet.add("liteTopic-" + i); + } + LiteSubscriptionDTO dto = new LiteSubscriptionDTO() + .setAction(LiteSubscriptionAction.PARTIAL_REMOVE) + .setGroup(group) + .setTopic(topic) + .setLiteTopicSet(liteTopicSet); + + when(groupConfig.getLiteBindTopic()).thenReturn(topic); + when(messagingProcessor.getSubscriptionGroupConfig(ctx, group)).thenReturn(groupConfig); + when(serviceManager.getLiteSubscriptionService()).thenReturn(liteSubscriptionService); + when(liteSubscriptionService.syncLiteSubscription(any(), any(), anyLong())) + .thenReturn(CompletableFuture.completedFuture(null)); + + assertDoesNotThrow(() -> clientProcessor.syncLiteSubscription(ctx, dto, 3000).join()); + verify(groupConfig, never()).getLiteSubClientQuota(); + verify(liteSubscriptionService).syncLiteSubscription(ctx, dto, 3000); + } + @Test public void testGetGroupOrException_groupExists_returnsConfig() { String group = "group";