Kafka states restore improvements

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

@@ -42,7 +42,7 @@ import org.thingsboard.server.queue.provider.TbRuleEngineQueueFactory;
 import org.thingsboard.server.service.cf.AbstractCalculatedFieldStateService;
 import org.thingsboard.server.service.cf.ctx.CalculatedFieldEntityCtxId;
 import org.thingsboard.server.service.queue.DefaultTbCalculatedFieldConsumerService.CalculatedFieldQueueConfig;
-import org.thingsboard.server.service.queue.consumer.MainQueueConsumerManager;
+import org.thingsboard.server.queue.common.consumer.MainQueueConsumerManager;
 
 import java.util.Set;
 import java.util.concurrent.ExecutorService;

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

@@ -55,7 +55,7 @@ import org.thingsboard.server.service.cf.CalculatedFieldCache;
 import org.thingsboard.server.service.cf.CalculatedFieldStateService;
 import org.thingsboard.server.service.profile.TbAssetProfileCache;
 import org.thingsboard.server.service.profile.TbDeviceProfileCache;
-import org.thingsboard.server.service.queue.consumer.MainQueueConsumerManager;
+import org.thingsboard.server.queue.common.consumer.MainQueueConsumerManager;
 import org.thingsboard.server.service.queue.processing.AbstractConsumerService;
 import org.thingsboard.server.service.queue.processing.IdMsgPair;
 import org.thingsboard.server.service.security.auth.jwt.settings.JwtSettingsService;

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

@@ -21,8 +21,6 @@ import com.google.common.util.concurrent.MoreExecutors;
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
 import lombok.Data;
-import lombok.Getter;
-import lombok.Setter;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.factory.annotation.Value;
@@ -90,7 +88,7 @@ import org.thingsboard.server.service.notification.NotificationSchedulerService;
 import org.thingsboard.server.service.ota.OtaPackageStateService;
 import org.thingsboard.server.service.profile.TbAssetProfileCache;
 import org.thingsboard.server.service.profile.TbDeviceProfileCache;
-import org.thingsboard.server.service.queue.consumer.MainQueueConsumerManager;
+import org.thingsboard.server.queue.common.consumer.MainQueueConsumerManager;
 import org.thingsboard.server.service.queue.processing.AbstractConsumerService;
 import org.thingsboard.server.service.queue.processing.IdMsgPair;
 import org.thingsboard.server.service.resource.TbImageService;

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

@@ -20,8 +20,6 @@ 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.Getter;
-import lombok.Setter;
 import lombok.extern.slf4j.Slf4j;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.jetbrains.annotations.NotNull;
@@ -52,7 +50,7 @@ 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.service.queue.consumer.MainQueueConsumerManager;
+import org.thingsboard.server.queue.common.consumer.MainQueueConsumerManager;
 import org.thingsboard.server.service.queue.processing.AbstractConsumerService;
 import org.thingsboard.server.service.queue.processing.IdMsgPair;
 

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

@@ -1,48 +0,0 @@
-/**
- * Copyright © 2016-2024 The Thingsboard Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.thingsboard.server.service.queue.ruleengine;
-
-import lombok.AllArgsConstructor;
-import lombok.Getter;
-import lombok.ToString;
-import org.thingsboard.server.common.data.queue.QueueConfig;
-import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
-
-import java.util.Set;
-
-@Getter
-@ToString
-@AllArgsConstructor
-public class TbQueueConsumerManagerTask {
-
-    private final QueueEvent event;
-    private QueueConfig config;
-    private Set<TopicPartitionInfo> partitions;
-    private boolean drainQueue;
-
-    public static TbQueueConsumerManagerTask delete(boolean drainQueue) {
-        return new TbQueueConsumerManagerTask(QueueEvent.DELETE, null, null, drainQueue);
-    }
-
-    public static TbQueueConsumerManagerTask configUpdate(QueueConfig config) {
-        return new TbQueueConsumerManagerTask(QueueEvent.CONFIG_UPDATE, config, null, false);
-    }
-
-    public static TbQueueConsumerManagerTask partitionChange(Set<TopicPartitionInfo> partitions) {
-        return new TbQueueConsumerManagerTask(QueueEvent.PARTITION_CHANGE, null, partitions, false);
-    }
-
-}

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

@@ -33,11 +33,14 @@ import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
 import org.thingsboard.server.gen.transport.TransportProtos.ToRuleEngineMsg;
 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.TbQueueConsumerManagerTask;
+import org.thingsboard.server.queue.common.consumer.TbQueueConsumerManagerTask.DeleteQueueTask;
+import org.thingsboard.server.queue.common.consumer.TbQueueConsumerTask;
 import org.thingsboard.server.queue.discovery.QueueKey;
 import org.thingsboard.server.service.queue.TbMsgPackCallback;
 import org.thingsboard.server.service.queue.TbMsgPackProcessingContext;
 import org.thingsboard.server.service.queue.TbRuleEngineConsumerStats;
-import org.thingsboard.server.service.queue.consumer.MainQueueConsumerManager;
 import org.thingsboard.server.service.queue.processing.TbRuleEngineProcessingDecision;
 import org.thingsboard.server.service.queue.processing.TbRuleEngineProcessingResult;
 import org.thingsboard.server.service.queue.processing.TbRuleEngineProcessingStrategy;
@@ -75,13 +78,13 @@ public class TbRuleEngineQueueConsumerManager extends MainQueueConsumerManager<T
     }
 
     public void delete(boolean drainQueue) {
-        addTask(TbQueueConsumerManagerTask.delete(drainQueue));
+        addTask(new DeleteQueueTask(drainQueue));
     }
 
     @Override
     protected void processTask(TbQueueConsumerManagerTask task) {
-        if (task.getEvent() == QueueEvent.DELETE) {
+        if (task instanceof DeleteQueueTask deleteQueueTask) {
-            doDelete(task.isDrainQueue());
+            doDelete(deleteQueueTask.drainQueue());
         }
     }
 

common/data/src/main/java/org/thingsboard/server/common/data/queue/QueueConfig.java

@@ -15,10 +15,22 @@
  */
 package org.thingsboard.server.common.data.queue;
 
+import lombok.Data;
+
 public interface QueueConfig {
 
     boolean isConsumerPerPartition();
 
     int getPollInterval();
 
+    static QueueConfig of(boolean consumerPerPartition, long pollInterval) {
+        return new BasicQueueConfig(consumerPerPartition, (int) pollInterval);
+    }
+
+    @Data
+    class BasicQueueConfig implements QueueConfig {
+        private final boolean consumerPerPartition;
+        private final int pollInterval;
+    }
+
 }

common/queue/src/main/java/org/thingsboard/server/queue/common/AbstractTbQueueConsumerTemplate.java

@@ -149,6 +149,9 @@ public abstract class AbstractTbQueueConsumerTemplate<R, T extends TbQueueMsg> i
     @Override
     public void commit() {
         if (consumerLock.isLocked()) {
+            if (stopped) {
+                return;
+            }
             log.error("commit. consumerLock is locked. will wait with no timeout. it looks like a race conditions or deadlock topic " + topic, new RuntimeException("stacktrace"));
         }
         consumerLock.lock();

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

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.thingsboard.server.service.queue.consumer;
+package org.thingsboard.server.queue.common.consumer;
 
 import lombok.Builder;
 import lombok.Getter;
@@ -23,10 +23,9 @@ import org.thingsboard.server.common.data.queue.QueueConfig;
 import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
 import org.thingsboard.server.queue.TbQueueConsumer;
 import org.thingsboard.server.queue.TbQueueMsg;
+import org.thingsboard.server.queue.common.consumer.TbQueueConsumerManagerTask.UpdateConfigTask;
+import org.thingsboard.server.queue.common.consumer.TbQueueConsumerManagerTask.UpdatePartitionsTask;
 import org.thingsboard.server.queue.discovery.QueueKey;
-import org.thingsboard.server.service.queue.ruleengine.QueueEvent;
-import org.thingsboard.server.service.queue.ruleengine.TbQueueConsumerManagerTask;
-import org.thingsboard.server.service.queue.ruleengine.TbQueueConsumerTask;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -34,6 +33,7 @@ import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ExecutorService;
@@ -42,6 +42,7 @@ import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 @Slf4j
@@ -85,20 +86,24 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
 
     public void init(C config) {
         this.config = config;
+        this.consumerWrapper = createConsumerWrapper(config);
+        log.debug("[{}] Initialized consumer for queue: {}", queueKey, config);
+    }
+
+    protected ConsumerWrapper<M> createConsumerWrapper(C config) {
         if (config.isConsumerPerPartition()) {
-            this.consumerWrapper = new ConsumerPerPartitionWrapper();
+            return new ConsumerPerPartitionWrapper();
         } else {
-            this.consumerWrapper = new SingleConsumerWrapper();
+            return new SingleConsumerWrapper();
         }
-        log.debug("[{}] Initialized consumer for queue: {}", queueKey, config);
     }
 
     public void update(C config) {
-        addTask(TbQueueConsumerManagerTask.configUpdate(config));
+        addTask(new UpdateConfigTask(config));
     }
 
     public void update(Set<TopicPartitionInfo> partitions) {
-        addTask(TbQueueConsumerManagerTask.partitionChange(partitions));
+        addTask(new UpdatePartitionsTask(partitions));
     }
 
     protected void addTask(TbQueueConsumerManagerTask todo) {
@@ -123,10 +128,10 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
                         }
                         log.trace("[{}] Processing task: {}", queueKey, task);
 
-                        if (task.getEvent() == QueueEvent.PARTITION_CHANGE) {
+                        if (task instanceof UpdatePartitionsTask updatePartitionsTask) {
-                            newPartitions = task.getPartitions();
+                            newPartitions = updatePartitionsTask.partitions();
-                        } else if (task.getEvent() == QueueEvent.CONFIG_UPDATE) {
+                        } else if (task instanceof UpdateConfigTask updateConfigTask) {
-                            newConfig = (C) task.getConfig();
+                            newConfig = (C) updateConfigTask.config();
                         } else {
                             processTask(task);
                         }
@@ -182,7 +187,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
         }
     }
 
-    public void doUpdate(Set<TopicPartitionInfo> partitions) {
+    private void doUpdate(Set<TopicPartitionInfo> partitions) {
         this.partitions = partitions;
         consumerWrapper.updatePartitions(partitions);
     }
@@ -197,6 +202,15 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
                 log.error("Failure in consumer loop", e);
             }
             log.info("[{}] Consumer stopped", consumerTask.getKey());
+
+            try {
+                Runnable callback = consumerTask.getCallback();
+                if (callback != null) {
+                    callback.run();
+                }
+            } catch (Throwable t) {
+                log.error("Failed to execute finish callback", t);
+            }
         });
         consumerTask.setTask(consumerLoop);
     }
@@ -241,13 +255,13 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
         awaitStop(30);
     }
 
-    public void awaitStop(long timeoutSec) {
+    private void awaitStop(int timeoutSec) {
         log.debug("[{}] Waiting for consumers to stop", queueKey);
         consumerWrapper.getConsumers().forEach(consumerTask -> consumerTask.awaitCompletion(timeoutSec));
         log.debug("[{}] Unsubscribed and stopped consumers", queueKey);
     }
 
-    private static String partitionsToString(Collection<TopicPartitionInfo> partitions) {
+    static String partitionsToString(Collection<TopicPartitionInfo> partitions) {
         return partitions.stream().map(TopicPartitionInfo::getFullTopicName).collect(Collectors.joining(", ", "[", "]"));
     }
 
@@ -273,15 +287,24 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
 
             Set<TopicPartitionInfo> removedPartitions = new HashSet<>(consumers.keySet());
             removedPartitions.removeAll(partitions);
+
             log.info("[{}] Added partitions: {}, removed partitions: {}", queueKey, partitionsToString(addedPartitions), partitionsToString(removedPartitions));
+            removePartitions(removedPartitions);
+            addPartitions(addedPartitions, null);
+        }
 
-            removedPartitions.forEach((tpi) -> consumers.get(tpi).initiateStop());
+        protected void removePartitions(Set<TopicPartitionInfo> removedPartitions) {
-            removedPartitions.forEach((tpi) -> consumers.remove(tpi).awaitCompletion());
+            removedPartitions.forEach((tpi) -> Optional.ofNullable(consumers.get(tpi)).ifPresent(TbQueueConsumerTask::initiateStop));
+            removedPartitions.forEach((tpi) -> Optional.ofNullable(consumers.remove(tpi)).ifPresent(TbQueueConsumerTask::awaitCompletion));
+        }
 
-            addedPartitions.forEach((tpi) -> {
+        protected void addPartitions(Set<TopicPartitionInfo> partitions, Consumer<TopicPartitionInfo> onStop) {
+            partitions.forEach(tpi -> {
                 Integer partitionId = tpi.getPartition().orElse(-1);
                 String key = queueKey + "-" + partitionId;
-                TbQueueConsumerTask<M> consumer = new TbQueueConsumerTask<>(key, () -> consumerCreator.apply(config, partitionId));
+                Runnable callback = onStop != null ? () -> onStop.accept(tpi) : null;
+
+                TbQueueConsumerTask<M> consumer = new TbQueueConsumerTask<>(key, () -> consumerCreator.apply(config, partitionId), callback);
                 consumers.put(tpi, consumer);
                 consumer.subscribe(Set.of(tpi));
                 launchConsumer(consumer);
@@ -310,7 +333,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
             }
 
             if (consumer == null) {
-                consumer = new TbQueueConsumerTask<>(queueKey, () -> consumerCreator.apply(config, null)); // no partitionId passed
+                consumer = new TbQueueConsumerTask<>(queueKey, () -> consumerCreator.apply(config, null), null); // no partitionId passed
             }
             consumer.subscribe(partitions);
             if (!consumer.isRunning()) {

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

@@ -0,0 +1,80 @@
+/**
+ * Copyright © 2016-2024 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.queue.common.consumer;
+
+import lombok.Builder;
+import lombok.Getter;
+import lombok.extern.slf4j.Slf4j;
+import org.thingsboard.server.common.data.queue.QueueConfig;
+import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
+import org.thingsboard.server.queue.TbQueueConsumer;
+import org.thingsboard.server.queue.TbQueueMsg;
+import org.thingsboard.server.queue.common.consumer.TbQueueConsumerManagerTask.AddPartitionsTask;
+import org.thingsboard.server.queue.common.consumer.TbQueueConsumerManagerTask.RemovePartitionsTask;
+import org.thingsboard.server.queue.discovery.QueueKey;
+
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.function.BiFunction;
+import java.util.function.Consumer;
+
+@Slf4j
+public class PartitionedQueueConsumerManager<M extends TbQueueMsg> extends MainQueueConsumerManager<M, QueueConfig> {
+
+    private final ConsumerPerPartitionWrapper consumerWrapper;
+    @Getter
+    private final String topic;
+
+    @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,
+                                           ExecutorService consumerExecutor, ScheduledExecutorService scheduler,
+                                           ExecutorService taskExecutor) {
+        super(queueKey, QueueConfig.of(true, pollInterval), msgPackProcessor, consumerCreator, consumerExecutor, scheduler, taskExecutor);
+        this.topic = topic;
+        this.consumerWrapper = (ConsumerPerPartitionWrapper) super.consumerWrapper;
+    }
+
+    @Override
+    public void update(Set<TopicPartitionInfo> partitions) {
+        throw new UnsupportedOperationException("Use manual addPartitions and removePartitions");
+    }
+
+    @Override
+    protected void processTask(TbQueueConsumerManagerTask task) {
+        if (task instanceof AddPartitionsTask addPartitionsTask) {
+            log.info("[{}] Added partitions: {}", queueKey, partitionsToString(addPartitionsTask.partitions()));
+            consumerWrapper.addPartitions(addPartitionsTask.partitions(), addPartitionsTask.onStop());
+        } else if (task instanceof RemovePartitionsTask removePartitionsTask) {
+            log.info("[{}] Removed partitions: {}", queueKey, partitionsToString(removePartitionsTask.partitions()));
+            consumerWrapper.removePartitions(removePartitionsTask.partitions());
+        }
+    }
+
+    public void addPartitions(Set<TopicPartitionInfo> partitions) {
+        addPartitions(partitions, null);
+    }
+
+    public void addPartitions(Set<TopicPartitionInfo> partitions, Consumer<TopicPartitionInfo> onStop) {
+        addTask(new AddPartitionsTask(partitions, onStop));
+    }
+
+    public void removePartitions(Set<TopicPartitionInfo> partitions) {
+        addTask(new RemovePartitionsTask(partitions));
+    }
+
+}

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

@@ -0,0 +1,76 @@
+/**
+ * Copyright © 2016-2024 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.queue.common.consumer;
+
+import lombok.Getter;
+import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
+import org.thingsboard.server.queue.TbQueueMsg;
+
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.stream.Collectors;
+
+public class QueueStateService<E extends TbQueueMsg, S extends TbQueueMsg> {
+
+    private PartitionedQueueConsumerManager<S> stateConsumer;
+    private PartitionedQueueConsumerManager<E> eventConsumer;
+
+    @Getter
+    private Set<TopicPartitionInfo> partitions;
+    private final Lock lock = new ReentrantLock();
+
+    public void init(PartitionedQueueConsumerManager<S> stateConsumer, PartitionedQueueConsumerManager<E> eventConsumer) {
+        this.stateConsumer = stateConsumer;
+        this.eventConsumer = eventConsumer;
+    }
+
+    public void update(Set<TopicPartitionInfo> newPartitions) {
+        lock.lock();
+        Set<TopicPartitionInfo> oldPartitions = this.partitions != null ? this.partitions : Collections.emptySet();
+        Set<TopicPartitionInfo> addedPartitions;
+        Set<TopicPartitionInfo> removedPartitions;
+        try {
+            addedPartitions = new HashSet<>(newPartitions);
+            addedPartitions.removeAll(oldPartitions);
+            removedPartitions = new HashSet<>(oldPartitions);
+            removedPartitions.removeAll(newPartitions);
+            this.partitions = newPartitions;
+        } finally {
+            lock.unlock();
+        }
+        if (!removedPartitions.isEmpty()) {
+            stateConsumer.removePartitions(removedPartitions);
+            eventConsumer.removePartitions(removedPartitions.stream().map(tpi -> tpi.newByTopic(eventConsumer.getTopic())).collect(Collectors.toSet()));
+        }
+
+        if (!addedPartitions.isEmpty()) {
+            stateConsumer.addPartitions(addedPartitions, partition -> {
+                lock.lock();
+                try {
+                    if (this.partitions.contains(partition)) {
+                        eventConsumer.addPartitions(Set.of(partition.newByTopic(eventConsumer.getTopic())));
+                    }
+                } finally {
+                    lock.unlock();
+                }
+            });
+        }
+    }
+
+}

application/src/main/java/org/thingsboard/server/service/queue/ruleengine/QueueEvent.java → common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/QueueTaskType.java

@@ -13,12 +13,13 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.thingsboard.server.service.queue.ruleengine;
+package org.thingsboard.server.queue.common.consumer;
 
 import java.io.Serializable;
 
-public enum QueueEvent implements Serializable {
+public enum QueueTaskType implements Serializable {
 
-    PARTITION_CHANGE, CONFIG_UPDATE, DELETE
+    UPDATE_PARTITIONS, UPDATE_CONFIG, DELETE,
+    ADD_PARTITIONS, REMOVE_PARTITIONS
 
 }

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

@@ -0,0 +1,63 @@
+/**
+ * Copyright © 2016-2024 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.queue.common.consumer;
+
+import org.thingsboard.server.common.data.queue.QueueConfig;
+import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
+
+import java.util.Set;
+import java.util.function.Consumer;
+
+public interface TbQueueConsumerManagerTask {
+
+    QueueTaskType getType();
+
+    record DeleteQueueTask(boolean drainQueue) implements TbQueueConsumerManagerTask {
+        @Override
+        public QueueTaskType getType() {
+            return QueueTaskType.DELETE;
+        }
+    }
+
+    record UpdateConfigTask(QueueConfig config) implements TbQueueConsumerManagerTask {
+        @Override
+        public QueueTaskType getType() {
+            return QueueTaskType.UPDATE_CONFIG;
+        }
+    }
+
+    record UpdatePartitionsTask(Set<TopicPartitionInfo> partitions) implements TbQueueConsumerManagerTask {
+        @Override
+        public QueueTaskType getType() {
+            return QueueTaskType.UPDATE_PARTITIONS;
+        }
+    }
+
+    record AddPartitionsTask(Set<TopicPartitionInfo> partitions, Consumer<TopicPartitionInfo> onStop) implements TbQueueConsumerManagerTask {
+        @Override
+        public QueueTaskType getType() {
+            return QueueTaskType.ADD_PARTITIONS;
+        }
+    }
+
+    record RemovePartitionsTask(Set<TopicPartitionInfo> partitions) implements TbQueueConsumerManagerTask {
+        @Override
+        public QueueTaskType getType() {
+            return QueueTaskType.REMOVE_PARTITIONS;
+        }
+    }
+
+}

application/src/main/java/org/thingsboard/server/service/queue/ruleengine/TbQueueConsumerTask.java → common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/TbQueueConsumerTask.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.thingsboard.server.service.queue.ruleengine;
+package org.thingsboard.server.queue.common.consumer;
 
 import lombok.Getter;
 import lombok.Setter;
@@ -35,14 +35,17 @@ public class TbQueueConsumerTask<M extends TbQueueMsg> {
     private final Object key;
     private volatile TbQueueConsumer<M> consumer;
     private volatile Supplier<TbQueueConsumer<M>> consumerSupplier;
+    @Getter
+    private final Runnable callback;
 
     @Setter
     private Future<?> task;
 
-    public TbQueueConsumerTask(Object key, Supplier<TbQueueConsumer<M>> consumerSupplier) {
+    public TbQueueConsumerTask(Object key, Supplier<TbQueueConsumer<M>> consumerSupplier, Runnable callback) {
         this.key = key;
         this.consumer = null;
         this.consumerSupplier = consumerSupplier;
+        this.callback = callback;
     }
 
     public TbQueueConsumer<M> getConsumer() {