wip cf consumer groups

application/src/main/java/org/thingsboard/server/service/cf/AbstractCalculatedFieldStateService.java

@@ -30,7 +30,9 @@ import org.thingsboard.server.service.cf.ctx.CalculatedFieldEntityCtxId;
 import org.thingsboard.server.service.cf.ctx.state.CalculatedFieldState;
 
 import java.util.Collection;
+import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
 import static org.thingsboard.server.utils.CalculatedFieldUtils.fromProto;
@@ -41,7 +43,7 @@ public abstract class AbstractCalculatedFieldStateService implements CalculatedF
     @Autowired
     private ActorSystemContext actorSystemContext;
 
-    protected QueueStateService<TbProtoQueueMsg<ToCalculatedFieldMsg>, TbProtoQueueMsg<CalculatedFieldStateProto>> stateService;
+    protected Map<QueueKey, QueueStateService<TbProtoQueueMsg<ToCalculatedFieldMsg>, TbProtoQueueMsg<CalculatedFieldStateProto>>> stateServices = new ConcurrentHashMap<>();
 
     @Override
     public final void persistState(CalculatedFieldEntityCtxId stateId, CalculatedFieldState state, TbCallback callback) {
@@ -72,22 +74,22 @@ public abstract class AbstractCalculatedFieldStateService implements CalculatedF
 
     @Override
     public void restore(QueueKey queueKey, Set<TopicPartitionInfo> partitions) {
-        stateService.update(queueKey, partitions);
+        stateServices.get(queueKey).update(queueKey, partitions);
     }
 
     @Override
-    public void delete(Set<TopicPartitionInfo> partitions) {
-        stateService.delete(partitions);
+    public void delete(QueueKey queueKey, Set<TopicPartitionInfo> partitions) {
+        stateServices.get(queueKey).delete(partitions);
     }
 
     @Override
-    public Set<TopicPartitionInfo> getPartitions() {
-        return stateService.getPartitions().values().stream().flatMap(Collection::stream).collect(Collectors.toSet());
+    public Set<TopicPartitionInfo> getPartitions(QueueKey queueKey) {
+        return stateServices.get(queueKey).getPartitions().values().stream().flatMap(Collection::stream).collect(Collectors.toSet());
     }
 
     @Override
-    public void stop() {
-        stateService.stop();
+    public void stop(QueueKey queueKey) {
+        stateServices.get(queueKey).stop();
     }
 
 }

application/src/main/java/org/thingsboard/server/service/cf/CalculatedFieldStateService.java

@@ -37,10 +37,10 @@ public interface CalculatedFieldStateService {
 
     void restore(QueueKey queueKey, Set<TopicPartitionInfo> partitions);
 
-    void delete(Set<TopicPartitionInfo> partitions);
+    void delete(QueueKey queueKey, Set<TopicPartitionInfo> partitions);
 
-    Set<TopicPartitionInfo> getPartitions();
+    Set<TopicPartitionInfo> getPartitions(QueueKey queueKey);
 
-    void stop();
+    void stop(QueueKey queueKey);
 
 }

application/src/main/java/org/thingsboard/server/service/cf/ctx/state/KafkaCalculatedFieldStateService.java

@@ -97,7 +97,7 @@ public class KafkaCalculatedFieldStateService extends AbstractCalculatedFieldSta
                 .scheduler(eventConsumer.getScheduler())
                 .taskExecutor(eventConsumer.getTaskExecutor())
                 .build();
-        super.stateService = new KafkaQueueStateService<>(eventConsumer, stateConsumer);
+        super.stateServices.put(queueKey, new KafkaQueueStateService<>(eventConsumer, stateConsumer));
         this.stateProducer = (TbKafkaProducerTemplate<TbProtoQueueMsg<CalculatedFieldStateProto>>) queueFactory.createCalculatedFieldStateProducer();
     }
 
@@ -145,8 +145,8 @@ public class KafkaCalculatedFieldStateService extends AbstractCalculatedFieldSta
     }
 
     @Override
-    public void stop() {
-        super.stop();
+    public void stop(QueueKey queueKey) {
+        super.stop(queueKey);
         stateProducer.stop();
     }
 

application/src/main/java/org/thingsboard/server/service/cf/ctx/state/RocksDBCalculatedFieldStateService.java

@@ -20,13 +20,15 @@ import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnExpression;
 import org.springframework.stereotype.Service;
+import org.thingsboard.server.common.data.DataConstants;
+import org.thingsboard.server.common.msg.queue.ServiceType;
 import org.thingsboard.server.common.msg.queue.TbCallback;
 import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
-import org.thingsboard.server.queue.common.state.DefaultQueueStateService;
 import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldStateProto;
 import org.thingsboard.server.gen.transport.TransportProtos.ToCalculatedFieldMsg;
 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
 import org.thingsboard.server.queue.common.consumer.PartitionedQueueConsumerManager;
+import org.thingsboard.server.queue.common.state.DefaultQueueStateService;
 import org.thingsboard.server.queue.discovery.QueueKey;
 import org.thingsboard.server.service.cf.AbstractCalculatedFieldStateService;
 import org.thingsboard.server.service.cf.CfRocksDb;
@@ -44,7 +46,7 @@ public class RocksDBCalculatedFieldStateService extends AbstractCalculatedFieldS
 
     @Override
     public void init(PartitionedQueueConsumerManager<TbProtoQueueMsg<ToCalculatedFieldMsg>> eventConsumer) {
-        super.stateService = new DefaultQueueStateService<>(eventConsumer);
+        super.stateServices.put(new QueueKey(ServiceType.TB_RULE_ENGINE, DataConstants.CF_STATES_QUEUE_NAME), new DefaultQueueStateService<>(eventConsumer));
     }
 
     @Override
@@ -61,7 +63,7 @@ public class RocksDBCalculatedFieldStateService extends AbstractCalculatedFieldS
 
     @Override
     public void restore(QueueKey queueKey, Set<TopicPartitionInfo> partitions) {
-        if (stateService.getPartitions().isEmpty()) {
+        if (stateServices.get(queueKey).getPartitions().isEmpty()) {
             cfRocksDb.forEach((key, value) -> {
                 try {
                     processRestoredState(CalculatedFieldStateProto.parseFrom(value));

application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCalculatedFieldConsumerService.java

@@ -30,12 +30,14 @@ import org.thingsboard.server.common.data.EntityType;
 import org.thingsboard.server.common.data.id.EntityIdFactory;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.plugin.ComponentLifecycleEvent;
+import org.thingsboard.server.common.data.queue.Queue;
 import org.thingsboard.server.common.data.queue.QueueConfig;
 import org.thingsboard.server.common.msg.cf.CalculatedFieldPartitionChangeMsg;
 import org.thingsboard.server.common.msg.plugin.ComponentLifecycleMsg;
 import org.thingsboard.server.common.msg.queue.ServiceType;
 import org.thingsboard.server.common.msg.queue.TbCallback;
 import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
+import org.thingsboard.server.dao.queue.QueueService;
 import org.thingsboard.server.dao.tenant.TbTenantProfileCache;
 import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldLinkedTelemetryMsgProto;
 import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldTelemetryMsgProto;
@@ -79,7 +81,10 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractConsumerPar
     private long packProcessingTimeout;
 
     private final TbRuleEngineQueueFactory queueFactory;
+    private final QueueService queueService;
+
     private final CalculatedFieldStateService stateService;
+    private final ConcurrentMap<QueueKey, PartitionedQueueConsumerManager<TbProtoQueueMsg<ToCalculatedFieldMsg>>> consumers = new ConcurrentHashMap<>();
 
     public DefaultTbCalculatedFieldConsumerService(TbRuleEngineQueueFactory tbQueueFactory,
                                                    ActorSystemContext actorContext,
@@ -91,28 +96,40 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractConsumerPar
                                                    ApplicationEventPublisher eventPublisher,
                                                    JwtSettingsService jwtSettingsService,
                                                    CalculatedFieldCache calculatedFieldCache,
+                                                   QueueService queueService,
                                                    CalculatedFieldStateService stateService) {
         super(actorContext, tenantProfileCache, deviceProfileCache, assetProfileCache, calculatedFieldCache, apiUsageStateService, partitionService,
                 eventPublisher, jwtSettingsService);
         this.queueFactory = tbQueueFactory;
+        this.queueService = queueService;
         this.stateService = stateService;
     }
 
     @Override
     protected void doAfterStartUp() {
-        var queueKey = new QueueKey(ServiceType.TB_RULE_ENGINE, DataConstants.CF_QUEUE_NAME);
-        PartitionedQueueConsumerManager<TbProtoQueueMsg<ToCalculatedFieldMsg>> eventConsumer = PartitionedQueueConsumerManager.<TbProtoQueueMsg<ToCalculatedFieldMsg>>create()
+        List<Queue> queues = queueService.findAllQueues();
+        for (Queue configuration : queues) {
+            if (partitionService.isManagedByCurrentService(configuration.getTenantId())) {
+                stateService.init(createConsumer(configuration));
+            }
+        }
+    }
+
+    private PartitionedQueueConsumerManager<TbProtoQueueMsg<ToCalculatedFieldMsg>> createConsumer(Queue queue) {
+        QueueKey queueKey = new QueueKey(ServiceType.TB_RULE_ENGINE, DataConstants.CF_QUEUE_NAME, queue.getTenantId());
+        var eventConsumer = PartitionedQueueConsumerManager.<TbProtoQueueMsg<ToCalculatedFieldMsg>>create()
                 .queueKey(queueKey)
                 .topic(partitionService.getTopic(queueKey))
                 .pollInterval(pollInterval)
                 .msgPackProcessor(this::processMsgs)
-                .consumerCreator((config, partitionId) -> queueFactory.createToCalculatedFieldMsgConsumer())
+                .consumerCreator((config, partitionId) -> queueFactory.createToCalculatedFieldMsgConsumer(queue, partitionId))
                 .queueAdmin(queueFactory.getCalculatedFieldQueueAdmin())
                 .consumerExecutor(consumersExecutor)
                 .scheduler(scheduler)
                 .taskExecutor(mgmtExecutor)
                 .build();
-        stateService.init(eventConsumer);
+        consumers.put(queueKey, eventConsumer);
+        return eventConsumer;
     }
 
     @PreDestroy
@@ -227,13 +244,20 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractConsumerPar
     public void handleComponentLifecycleEvent(ComponentLifecycleMsg event) {
         if (event.getEntityId().getEntityType() == EntityType.TENANT) {
             if (event.getEvent() == ComponentLifecycleEvent.DELETED) {
-                Set<TopicPartitionInfo> partitions = stateService.getPartitions();
-                if (CollectionUtils.isEmpty(partitions)) {
-                    return;
-                }
-                stateService.delete(partitions.stream()
-                        .filter(tpi -> tpi.getTenantId().isPresent() && tpi.getTenantId().get().equals(event.getTenantId()))
-                        .collect(Collectors.toSet()));
+                consumers.keySet().removeIf(queueKey -> {
+                    boolean toRemove = queueKey.getTenantId().equals(event.getTenantId());
+                    if (toRemove) {
+                        Set<TopicPartitionInfo> partitions = stateService.getPartitions(queueKey);
+                        if (!CollectionUtils.isEmpty(partitions)) {
+                            stateService.delete(queueKey, partitions);
+                        }
+                        var consumer = consumers.remove(queueKey);
+                        if (consumer != null) {
+                            consumer.stop();
+                        }
+                    }
+                    return toRemove;
+                });
             }
         }
     }
@@ -258,7 +282,7 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractConsumerPar
     @Override
     protected void stopConsumers() {
         super.stopConsumers();
-        stateService.stop(); // eventConsumer will be stopped by stateService
+        consumers.keySet().forEach(stateService::stop); // eventConsumer will be stopped by stateService
     }
 
 }

common/queue/src/main/java/org/thingsboard/server/queue/provider/InMemoryMonolithQueueFactory.java

@@ -134,7 +134,7 @@ public class InMemoryMonolithQueueFactory implements TbCoreQueueFactory, TbRuleE
     }
 
     @Override
-    public TbQueueConsumer<TbProtoQueueMsg<TransportProtos.ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer() {
+    public TbQueueConsumer<TbProtoQueueMsg<TransportProtos.ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer(Queue configuration, Integer partitionId) {
         return new InMemoryTbQueueConsumer<>(storage, topicService.buildTopicName(calculatedFieldSettings.getEventTopic()));
     }
 

common/queue/src/main/java/org/thingsboard/server/queue/provider/KafkaMonolithQueueFactory.java

@@ -104,7 +104,7 @@ public class KafkaMonolithQueueFactory implements TbCoreQueueFactory, TbRuleEngi
     private final TbQueueAdmin housekeeperReprocessingAdmin;
     private final TbQueueAdmin edgeAdmin;
     private final TbQueueAdmin edgeEventAdmin;
-    private final TbQueueAdmin cfAdmin;
+    private final TbKafkaAdmin cfAdmin;
     private final TbQueueAdmin cfStateAdmin;
     private final TbQueueAdmin edqsEventsAdmin;
     private final TbKafkaAdmin edqsRequestsAdmin;
@@ -514,15 +514,22 @@ public class KafkaMonolithQueueFactory implements TbCoreQueueFactory, TbRuleEngi
     }
 
     @Override
-    public TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer() {
+    public TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer(Queue configuration, Integer partitionId) {
+        String queueName = configuration.getName();
+        String groupId = topicService.buildConsumerGroupId("cf-", configuration.getTenantId(), queueName, partitionId);
+
+        cfAdmin.syncOffsets(topicService.buildConsumerGroupId("cf-", configuration.getTenantId(), queueName, null), // the fat groupId
+                groupId, partitionId);
+
         TbKafkaConsumerTemplate.TbKafkaConsumerTemplateBuilder<TbProtoQueueMsg<ToCalculatedFieldMsg>> consumerBuilder = TbKafkaConsumerTemplate.builder();
         consumerBuilder.settings(kafkaSettings);
         consumerBuilder.topic(topicService.buildTopicName(calculatedFieldSettings.getEventTopic()));
-        consumerBuilder.clientId("monolith-calculated-field-consumer-" + serviceInfoProvider.getServiceId() + "-" + consumerCount.incrementAndGet());
-        consumerBuilder.groupId(topicService.buildTopicName("monolith-calculated-field-consumer"));
+        consumerBuilder.clientId("cf-" + queueName + "-consumer-" + serviceInfoProvider.getServiceId() + "-" + consumerCount.incrementAndGet());
+        consumerBuilder.groupId(groupId);
         consumerBuilder.decoder(msg -> new TbProtoQueueMsg<>(msg.getKey(), ToCalculatedFieldMsg.parseFrom(msg.getData()), msg.getHeaders()));
         consumerBuilder.admin(cfAdmin);
         consumerBuilder.statsService(consumerStatsService);
+
         return consumerBuilder.build();
     }
 

common/queue/src/main/java/org/thingsboard/server/queue/provider/KafkaTbRuleEngineQueueFactory.java

@@ -89,7 +89,7 @@ public class KafkaTbRuleEngineQueueFactory implements TbRuleEngineQueueFactory {
     private final TbQueueAdmin housekeeperAdmin;
     private final TbQueueAdmin edgeAdmin;
     private final TbQueueAdmin edgeEventAdmin;
-    private final TbQueueAdmin cfAdmin;
+    private final TbKafkaAdmin cfAdmin;
     private final TbQueueAdmin cfStateAdmin;
     private final TbQueueAdmin edqsEventsAdmin;
     private final AtomicLong consumerCount = new AtomicLong();
@@ -309,12 +309,18 @@ public class KafkaTbRuleEngineQueueFactory implements TbRuleEngineQueueFactory {
     }
 
     @Override
-    public TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer() {
+    public TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer(Queue configuration, Integer partitionId) {
+        String queueName = configuration.getName();
+        String groupId = topicService.buildConsumerGroupId("cf-", configuration.getTenantId(), queueName, partitionId);
+
+        cfAdmin.syncOffsets(topicService.buildConsumerGroupId("cf-", configuration.getTenantId(), queueName, null), // the fat groupId
+                groupId, partitionId);
+
         TbKafkaConsumerTemplate.TbKafkaConsumerTemplateBuilder<TbProtoQueueMsg<ToCalculatedFieldMsg>> consumerBuilder = TbKafkaConsumerTemplate.builder();
         consumerBuilder.settings(kafkaSettings);
         consumerBuilder.topic(topicService.buildTopicName(calculatedFieldSettings.getEventTopic()));
-        consumerBuilder.clientId("tb-rule-engine-calculated-field-consumer-" + serviceInfoProvider.getServiceId() + "-" + consumerCount.incrementAndGet());
-        consumerBuilder.groupId(topicService.buildTopicName("tb-rule-engine-calculated-field-consumer"));
+        consumerBuilder.clientId("cf-" + queueName + "-consumer-" + serviceInfoProvider.getServiceId() + "-" + consumerCount.incrementAndGet());
+        consumerBuilder.groupId(groupId);
         consumerBuilder.decoder(msg -> new TbProtoQueueMsg<>(msg.getKey(), ToCalculatedFieldMsg.parseFrom(msg.getData()), msg.getHeaders()));
         consumerBuilder.admin(cfAdmin);
         consumerBuilder.statsService(consumerStatsService);

common/queue/src/main/java/org/thingsboard/server/queue/provider/TbRuleEngineQueueFactory.java

@@ -121,7 +121,7 @@ public interface TbRuleEngineQueueFactory extends TbUsageStatsClientQueueFactory
 
     TbQueueRequestTemplate<TbProtoJsQueueMsg<JsInvokeProtos.RemoteJsRequest>, TbProtoQueueMsg<JsInvokeProtos.RemoteJsResponse>> createRemoteJsRequestTemplate();
 
-    TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer();
+    TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer(Queue configuration, Integer partitionId);
 
     TbQueueAdmin getCalculatedFieldQueueAdmin();