used tpi in main consumer manager

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

@@ -91,7 +91,7 @@ public class KafkaCalculatedFieldStateService extends AbstractCalculatedFieldSta
                         }
                     }
                 })
-                .consumerCreator((queueConfig, partitionId) -> queueFactory.createCalculatedFieldStateConsumer())
+                .consumerCreator((queueConfig, tpi) -> queueFactory.createCalculatedFieldStateConsumer())
                 .queueAdmin(queueFactory.getCalculatedFieldQueueAdmin())
                 .consumerExecutor(eventConsumer.getConsumerExecutor())
                 .scheduler(eventConsumer.getScheduler())

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

@@ -17,6 +17,7 @@ package org.thingsboard.server.service.queue;
 
 import jakarta.annotation.PreDestroy;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.collections4.CollectionUtils;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.context.event.EventListener;
@@ -35,7 +36,6 @@ 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;
@@ -43,7 +43,6 @@ import org.thingsboard.server.gen.transport.TransportProtos.ToCalculatedFieldMsg
 import org.thingsboard.server.gen.transport.TransportProtos.ToCalculatedFieldNotificationMsg;
 import org.thingsboard.server.queue.TbQueueConsumer;
 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
-import org.thingsboard.server.queue.common.consumer.MainQueueConsumerManager;
 import org.thingsboard.server.queue.common.consumer.PartitionedQueueConsumerManager;
 import org.thingsboard.server.queue.discovery.PartitionService;
 import org.thingsboard.server.queue.discovery.QueueKey;
@@ -61,7 +60,6 @@ import org.thingsboard.server.service.queue.processing.IdMsgPair;
 import org.thingsboard.server.service.security.auth.jwt.settings.JwtSettingsService;
 
 import java.util.List;
-import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
@@ -85,8 +83,6 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractPartitionBa
     private final CalculatedFieldStateService stateService;
     private final CalculatedFieldEntityProfileCache entityProfileCache;
 
-    private final ConcurrentMap<QueueKey, PartitionedQueueConsumerManager<TbProtoQueueMsg<ToCalculatedFieldMsg>>> consumers = new ConcurrentHashMap<>();
-
     public DefaultTbCalculatedFieldConsumerService(TbRuleEngineQueueFactory tbQueueFactory,
                                                    ActorSystemContext actorContext,
                                                    TbDeviceProfileCache deviceProfileCache,
@@ -109,32 +105,18 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractPartitionBa
     @Override
     protected void onStartUp() {
         var queueKey = new QueueKey(ServiceType.TB_RULE_ENGINE, DataConstants.CF_QUEUE_NAME);
-        createConsumer(queueKey);
-    }
-
-    private PartitionedQueueConsumerManager<TbProtoQueueMsg<ToCalculatedFieldMsg>> createConsumer(QueueKey queueKey) {
-        String topic = partitionService.getTopic(queueKey);
         var eventConsumer = PartitionedQueueConsumerManager.<TbProtoQueueMsg<ToCalculatedFieldMsg>>create()
                 .queueKey(queueKey)
-                .topic(topic)
+                .topic(partitionService.getTopic(queueKey))
                 .pollInterval(pollInterval)
                 .msgPackProcessor(this::processMsgs)
-                .consumerCreator((queueConfig, partitionId) -> {
-                    TopicPartitionInfo tpi = TopicPartitionInfo.builder()
-                            .tenantId(queueKey.getTenantId())
-                            .topic(partitionService.getTopic(queueKey))
-                            .partition(partitionId)
-                            .build();
-                    return queueFactory.createToCalculatedFieldMsgConsumer(tpi, partitionId);
-                })
+                .consumerCreator((queueConfig, tpi) -> queueFactory.createToCalculatedFieldMsgConsumer(tpi))
                 .queueAdmin(queueFactory.getCalculatedFieldQueueAdmin())
                 .consumerExecutor(consumersExecutor)
                 .scheduler(scheduler)
                 .taskExecutor(mgmtExecutor)
                 .build();
         stateService.init(eventConsumer);
-        consumers.put(queueKey, eventConsumer);
-        return eventConsumer;
     }
 
     @PreDestroy
@@ -151,26 +133,11 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractPartitionBa
     protected void onPartitionChangeEvent(PartitionChangeEvent event) {
         try {
             event.getNewPartitions().forEach((queueKey, partitions) -> {
-                if (DataConstants.CF_QUEUE_NAME.equals(queueKey.getQueueName()) || DataConstants.CF_STATES_QUEUE_NAME.equals(queueKey.getQueueName())) {
-                    if (partitionService.isManagedByCurrentService(queueKey.getTenantId())) {
-                        var consumer = Optional.ofNullable(consumers.get(queueKey)).orElseGet(() -> createConsumer(queueKey));
-                        if (consumer != null) {
-                            stateService.restore(queueKey, partitions);
-                        }
-                    }
+                if (DataConstants.CF_QUEUE_NAME.equals(queueKey.getQueueName())) {
+                    stateService.restore(queueKey, partitions);
                 }
             });
 
-            consumers.keySet().stream()
-                    .collect(Collectors.groupingBy(QueueKey::getTenantId))
-                    .forEach((tenantId, queueKeys) -> {
-                        if (!partitionService.isManagedByCurrentService(tenantId)) {
-                            queueKeys.forEach(queueKey -> {
-                                Optional.ofNullable(consumers.remove(queueKey)).ifPresent(MainQueueConsumerManager::stop);
-                            });
-                        }
-                    });
-
             // Cleanup old entities after corresponding consumers are stopped.
             // Any periodic tasks need to check that the entity is still managed by the current server before processing.
             actorContext.tell(new CalculatedFieldPartitionChangeMsg());
@@ -265,19 +232,13 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractPartitionBa
             if (event.getEvent() == ComponentLifecycleEvent.DELETED) {
                 entityProfileCache.removeTenant(event.getTenantId());
 
-                List<QueueKey> toRemove = consumers.keySet().stream()
-                        .filter(queueKey -> queueKey.getTenantId().equals(event.getTenantId()))
-                        .toList();
-                toRemove.forEach(queueKey -> {
-                    Optional.ofNullable(consumers.remove(queueKey)).ifPresent(consumer -> {
-                        Set<TopicPartitionInfo> partitions = stateService.getPartitions().stream()
-                                .filter(tpi -> tpi.getTenantId().isPresent() && tpi.getTenantId().get().equals(event.getTenantId()))
-                                .collect(Collectors.toSet());
-                        if (!partitions.isEmpty()) {
-                            consumer.delete(partitions);
-                        }
-                    });
-                });
+                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()));
             }
         } else if (event.getEntityId().getEntityType() == EntityType.ASSET_PROFILE) {
             if (event.getEvent() == ComponentLifecycleEvent.DELETED) {

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

@@ -206,7 +206,7 @@ public class DefaultTbCoreConsumerService extends AbstractConsumerService<ToCore
                 .queueKey(new QueueKey(ServiceType.TB_CORE))
                 .config(QueueConfig.of(consumerPerPartition, pollInterval))
                 .msgPackProcessor(this::processMsgs)
-                .consumerCreator((config, partitionId) -> queueFactory.createToCoreMsgConsumer())
+                .consumerCreator((config, tpi) -> queueFactory.createToCoreMsgConsumer())
                 .consumerExecutor(consumersExecutor)
                 .scheduler(scheduler)
                 .taskExecutor(mgmtExecutor)

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

@@ -19,7 +19,6 @@ import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
-import lombok.Data;
 import lombok.extern.slf4j.Slf4j;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.jetbrains.annotations.NotNull;
@@ -45,12 +44,12 @@ import org.thingsboard.server.gen.transport.TransportProtos.ToEdgeMsg;
 import org.thingsboard.server.gen.transport.TransportProtos.ToEdgeNotificationMsg;
 import org.thingsboard.server.queue.TbQueueConsumer;
 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
+import org.thingsboard.server.queue.common.consumer.MainQueueConsumerManager;
 import org.thingsboard.server.queue.discovery.QueueKey;
 import org.thingsboard.server.queue.discovery.event.PartitionChangeEvent;
 import org.thingsboard.server.queue.provider.TbCoreQueueFactory;
 import org.thingsboard.server.queue.util.TbCoreComponent;
 import org.thingsboard.server.service.edge.EdgeContextComponent;
-import org.thingsboard.server.queue.common.consumer.MainQueueConsumerManager;
 import org.thingsboard.server.service.queue.processing.AbstractConsumerService;
 import org.thingsboard.server.service.queue.processing.IdMsgPair;
 
@@ -104,7 +103,7 @@ public class DefaultTbEdgeConsumerService extends AbstractConsumerService<ToEdge
                 .queueKey(new QueueKey(ServiceType.TB_CORE).withQueueName(DataConstants.EDGE_QUEUE_NAME))
                 .config(QueueConfig.of(consumerPerPartition, pollInterval))
                 .msgPackProcessor(this::processMsgs)
-                .consumerCreator((config, partitionId) -> queueFactory.createEdgeMsgConsumer())
+                .consumerCreator((config, tpi) -> queueFactory.createEdgeMsgConsumer())
                 .consumerExecutor(consumersExecutor)
                 .scheduler(scheduler)
                 .taskExecutor(mgmtExecutor)

application/src/main/java/org/thingsboard/server/service/queue/ruleengine/TbRuleEngineQueueConsumerManager.java

@@ -72,7 +72,12 @@ public class TbRuleEngineQueueConsumerManager extends MainQueueConsumerManager<T
                                             ExecutorService consumerExecutor,
                                             ScheduledExecutorService scheduler,
                                             ExecutorService taskExecutor) {
-        super(queueKey, null, null, ctx.getQueueFactory()::createToRuleEngineMsgConsumer, consumerExecutor, scheduler, taskExecutor, null);
+        super(queueKey, null, null,
+                (queueConfig, topicPartitionInfo) -> {
+                    Integer partitionId = topicPartitionInfo.getPartition().orElse(-1);
+                    return ctx.getQueueFactory().createToRuleEngineMsgConsumer(queueConfig, partitionId);
+                },
+                consumerExecutor, scheduler, taskExecutor, null);
         this.ctx = ctx;
         this.stats = new TbRuleEngineConsumerStats(queueKey, ctx.getStatsFactory());
     }

common/edqs/src/main/java/org/thingsboard/server/edqs/processor/EdqsProcessor.java

@@ -139,7 +139,7 @@ public class EdqsProcessor implements TbQueueHandler<TbProtoQueueMsg<ToEdqsMsg>,
                     }
                     consumer.commit();
                 })
-                .consumerCreator((config, partitionId) -> queueFactory.createEdqsEventsConsumer())
+                .consumerCreator((config, tpi) -> queueFactory.createEdqsEventsConsumer())
                 .queueAdmin(queueFactory.getEdqsQueueAdmin())
                 .consumerExecutor(consumersExecutor)
                 .taskExecutor(taskExecutor)

common/edqs/src/main/java/org/thingsboard/server/edqs/state/KafkaEdqsStateService.java

@@ -59,7 +59,8 @@ public class KafkaEdqsStateService implements EdqsStateService {
     private final EdqsConfig config;
     private final EdqsPartitionService partitionService;
     private final KafkaEdqsQueueFactory queueFactory;
-    @Autowired @Lazy
+    @Autowired
+    @Lazy
     private EdqsProcessor edqsProcessor;
 
     private PartitionedQueueConsumerManager<TbProtoQueueMsg<ToEdqsMsg>> stateConsumer;
@@ -93,7 +94,7 @@ public class KafkaEdqsStateService implements EdqsStateService {
                     }
                     consumer.commit();
                 })
-                .consumerCreator((config, partitionId) -> queueFactory.createEdqsStateConsumer())
+                .consumerCreator((config, tpi) -> queueFactory.createEdqsStateConsumer())
                 .queueAdmin(queueAdmin)
                 .consumerExecutor(eventConsumer.getConsumerExecutor())
                 .taskExecutor(eventConsumer.getTaskExecutor())

common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/MainQueueConsumerManager.java

@@ -54,7 +54,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
     @Getter
     protected C config;
     protected final MsgPackProcessor<M, C> msgPackProcessor;
-    protected final BiFunction<C, Integer, TbQueueConsumer<M>> consumerCreator;
+    protected final BiFunction<C, TopicPartitionInfo, TbQueueConsumer<M>> consumerCreator;
     @Getter
     protected final ExecutorService consumerExecutor;
     @Getter
@@ -74,7 +74,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
     @Builder
     public MainQueueConsumerManager(QueueKey queueKey, C config,
                                     MsgPackProcessor<M, C> msgPackProcessor,
-                                    BiFunction<C, Integer, TbQueueConsumer<M>> consumerCreator,
+                                    BiFunction<C, TopicPartitionInfo, TbQueueConsumer<M>> consumerCreator,
                                     ExecutorService consumerExecutor,
                                     ScheduledExecutorService scheduler,
                                     ExecutorService taskExecutor,
@@ -313,7 +313,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
                 Runnable callback = onStop != null ? () -> onStop.accept(tpi) : null;
 
                 TbQueueConsumerTask<M> consumer = new TbQueueConsumerTask<>(key, () -> {
-                    TbQueueConsumer<M> queueConsumer = consumerCreator.apply(config, partitionId);
+                    TbQueueConsumer<M> queueConsumer = consumerCreator.apply(config, tpi);
                     if (startOffsetProvider != null && queueConsumer instanceof TbKafkaConsumerTemplate<M> kafkaConsumer) {
                         kafkaConsumer.setStartOffsetProvider(startOffsetProvider);
                     }

common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/PartitionedQueueConsumerManager.java

@@ -45,7 +45,7 @@ public class PartitionedQueueConsumerManager<M extends TbQueueMsg> extends MainQ
 
     @Builder(builderMethodName = "create") // not to conflict with super.builder()
     public PartitionedQueueConsumerManager(QueueKey queueKey, String topic, long pollInterval, MsgPackProcessor<M, QueueConfig> msgPackProcessor,
-                                           BiFunction<QueueConfig, Integer, TbQueueConsumer<M>> consumerCreator, TbQueueAdmin queueAdmin,
+                                           BiFunction<QueueConfig, TopicPartitionInfo, TbQueueConsumer<M>> consumerCreator, TbQueueAdmin queueAdmin,
                                            ExecutorService consumerExecutor, ScheduledExecutorService scheduler,
                                            ExecutorService taskExecutor, Consumer<Throwable> uncaughtErrorHandler) {
         super(queueKey, QueueConfig.of(true, pollInterval), msgPackProcessor, consumerCreator, consumerExecutor, scheduler, taskExecutor, uncaughtErrorHandler);

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(TopicPartitionInfo tpi, Integer partitionId) {
+    public TbQueueConsumer<TbProtoQueueMsg<TransportProtos.ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer(TopicPartitionInfo tpi) {
         return new InMemoryTbQueueConsumer<>(storage, topicService.buildTopicName(calculatedFieldSettings.getEventTopic()));
     }
 

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

@@ -105,7 +105,7 @@ public class KafkaMonolithQueueFactory implements TbCoreQueueFactory, TbRuleEngi
     private final TbQueueAdmin housekeeperReprocessingAdmin;
     private final TbQueueAdmin edgeAdmin;
     private final TbQueueAdmin edgeEventAdmin;
-    private final TbKafkaAdmin cfAdmin;
+    private final TbQueueAdmin cfAdmin;
     private final TbQueueAdmin cfStateAdmin;
     private final TbQueueAdmin edqsEventsAdmin;
     private final TbKafkaAdmin edqsRequestsAdmin;
@@ -515,9 +515,10 @@ public class KafkaMonolithQueueFactory implements TbCoreQueueFactory, TbRuleEngi
     }
 
     @Override
-    public TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer(TopicPartitionInfo tpi, Integer partitionId) {
+    public TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer(TopicPartitionInfo tpi) {
         String queueName = DataConstants.CF_QUEUE_NAME;
         TenantId tenantId = tpi.getTenantId().orElse(TenantId.SYS_TENANT_ID);
+        Integer partitionId = tpi.getPartition().orElseThrow(() -> new IllegalArgumentException("PartitionId is required."));
         String groupId = topicService.buildConsumerGroupId("cf-", tenantId, queueName, partitionId);
 
         TbKafkaConsumerTemplate.TbKafkaConsumerTemplateBuilder<TbProtoQueueMsg<ToCalculatedFieldMsg>> consumerBuilder = TbKafkaConsumerTemplate.builder();

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

@@ -93,7 +93,7 @@ public class KafkaTbRuleEngineQueueFactory implements TbRuleEngineQueueFactory {
     private final TbQueueAdmin housekeeperAdmin;
     private final TbQueueAdmin edgeAdmin;
     private final TbQueueAdmin edgeEventAdmin;
-    private final TbKafkaAdmin cfAdmin;
+    private final TbQueueAdmin cfAdmin;
     private final TbQueueAdmin cfStateAdmin;
     private final TbQueueAdmin edqsEventsAdmin;
     private final AtomicLong consumerCount = new AtomicLong();
@@ -316,9 +316,10 @@ public class KafkaTbRuleEngineQueueFactory implements TbRuleEngineQueueFactory {
     }
 
     @Override
-    public TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer(TopicPartitionInfo tpi, Integer partitionId) {
+    public TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> createToCalculatedFieldMsgConsumer(TopicPartitionInfo tpi) {
         String queueName = DataConstants.CF_QUEUE_NAME;
         TenantId tenantId = tpi.getTenantId().orElse(TenantId.SYS_TENANT_ID);
+        Integer partitionId = tpi.getPartition().orElseThrow(() -> new IllegalArgumentException("PartitionId is required."));
         String groupId = topicService.buildConsumerGroupId("cf-", tenantId, queueName, partitionId);
 
         TbKafkaConsumerTemplate.TbKafkaConsumerTemplateBuilder<TbProtoQueueMsg<ToCalculatedFieldMsg>> consumerBuilder = TbKafkaConsumerTemplate.builder();

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

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