[ISSUE #9789] LitePullConsumer supports manually adding subscription reported in Heartbeat (#9790)

Author: qianye1001

client/src/main/java/org/apache/rocketmq/client/consumer/DefaultLitePullConsumer.java

@@ -17,6 +17,7 @@
 package org.apache.rocketmq.client.consumer;
 
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -33,11 +34,13 @@
 import org.apache.rocketmq.common.consumer.ConsumeFromWhere;
 import org.apache.rocketmq.common.message.MessageExt;
 import org.apache.rocketmq.common.message.MessageQueue;
+import org.apache.rocketmq.logging.org.slf4j.Logger;
+import org.apache.rocketmq.logging.org.slf4j.LoggerFactory;
 import org.apache.rocketmq.remoting.RPCHook;
 import org.apache.rocketmq.remoting.protocol.NamespaceUtil;
+import org.apache.rocketmq.remoting.protocol.filter.FilterAPI;
 import org.apache.rocketmq.remoting.protocol.heartbeat.MessageModel;
-import org.apache.rocketmq.logging.org.slf4j.Logger;
-import org.apache.rocketmq.logging.org.slf4j.LoggerFactory;
+import org.apache.rocketmq.remoting.protocol.heartbeat.SubscriptionData;
 
 import static org.apache.rocketmq.remoting.protocol.heartbeat.SubscriptionData.SUB_ALL;
 
@@ -171,6 +174,8 @@ public class DefaultLitePullConsumer extends ClientConfig implements LitePullCon
 
     private RPCHook rpcHook;
 
+    private final Set<SubscriptionData> subscriptionsForHeartbeat = new HashSet<>();
+
     /**
      * Default constructor.
      */
@@ -618,4 +623,17 @@ public boolean isEnableMsgTrace() {
     public void setEnableMsgTrace(boolean enableMsgTrace) {
         this.enableTrace = enableMsgTrace;
     }
+
+    public Set<SubscriptionData> getSubscriptionsForHeartbeat() {
+        return this.subscriptionsForHeartbeat;
+    }
+
+    public synchronized void buildSubscriptionsForHeartbeat(Map<String, MessageSelector> messageSelectorMap) throws Exception {
+        this.subscriptionsForHeartbeat.clear();
+        for (Map.Entry<String, MessageSelector> entry : messageSelectorMap.entrySet()) {
+            SubscriptionData subscriptionData = FilterAPI.build(entry.getKey(),
+                entry.getValue().getExpression(), entry.getValue().getExpressionType());
+            this.subscriptionsForHeartbeat.add(subscriptionData);
+        }
+    }
 }

client/src/main/java/org/apache/rocketmq/client/consumer/LitePullConsumer.java

@@ -16,14 +16,13 @@
  */
 package org.apache.rocketmq.client.consumer;
 
-import org.apache.rocketmq.client.exception.MQClientException;
-import org.apache.rocketmq.common.message.MessageExt;
-import org.apache.rocketmq.common.message.MessageQueue;
-
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import org.apache.rocketmq.client.exception.MQClientException;
+import org.apache.rocketmq.common.message.MessageExt;
+import org.apache.rocketmq.common.message.MessageQueue;
 
 public interface LitePullConsumer {
 
@@ -107,6 +106,8 @@ public interface LitePullConsumer {
      */
     void setSubExpressionForAssign(final String topic, final String subExpression);
 
+    void buildSubscriptionsForHeartbeat(Map<String, MessageSelector> subExpressionMap) throws Exception;
+
     /**
      * Fetch data for the topics or partitions specified using assign API
      *

client/src/main/java/org/apache/rocketmq/client/impl/consumer/DefaultLitePullConsumerImpl.java

@@ -63,6 +63,8 @@
 import org.apache.rocketmq.common.message.MessageExt;
 import org.apache.rocketmq.common.message.MessageQueue;
 import org.apache.rocketmq.common.sysflag.PullSysFlag;
+import org.apache.rocketmq.logging.org.slf4j.Logger;
+import org.apache.rocketmq.logging.org.slf4j.LoggerFactory;
 import org.apache.rocketmq.remoting.RPCHook;
 import org.apache.rocketmq.remoting.exception.RemotingException;
 import org.apache.rocketmq.remoting.protocol.NamespaceUtil;
@@ -73,8 +75,6 @@
 import org.apache.rocketmq.remoting.protocol.heartbeat.ConsumeType;
 import org.apache.rocketmq.remoting.protocol.heartbeat.MessageModel;
 import org.apache.rocketmq.remoting.protocol.heartbeat.SubscriptionData;
-import org.apache.rocketmq.logging.org.slf4j.Logger;
-import org.apache.rocketmq.logging.org.slf4j.LoggerFactory;
 
 public class DefaultLitePullConsumerImpl implements MQConsumerInner {
 
@@ -1122,7 +1122,8 @@ public Set<SubscriptionData> subscriptions() {
         Set<SubscriptionData> subSet = new HashSet<>();
 
         subSet.addAll(this.rebalanceImpl.getSubscriptionInner().values());
-
+        subSet.addAll(this.defaultLitePullConsumer.getSubscriptionsForHeartbeat());
+        
         return subSet;
     }