consumer creation refactored after merge

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

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

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

@@ -41,7 +41,7 @@ import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
-import java.util.function.Function;
+import java.util.function.BiFunction;
 import java.util.stream.Collectors;
 
 @Slf4j
@@ -51,7 +51,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
     @Getter
     protected C config;
     protected final MsgPackProcessor<M, C> msgPackProcessor;
-    protected final Function<C, TbQueueConsumer<M>> consumerCreator;
+    protected final BiFunction<C, Integer, TbQueueConsumer<M>> consumerCreator;
     protected final ExecutorService consumerExecutor;
     protected final ScheduledExecutorService scheduler;
     protected final ExecutorService taskExecutor;
@@ -67,7 +67,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
     @Builder
     public MainQueueConsumerManager(QueueKey queueKey, C config,
                                     MsgPackProcessor<M, C> msgPackProcessor,
-                                    Function<C, TbQueueConsumer<M>> consumerCreator,
+                                    BiFunction<C, Integer, TbQueueConsumer<M>> consumerCreator,
                                     ExecutorService consumerExecutor,
                                     ScheduledExecutorService scheduler,
                                     ExecutorService taskExecutor) {
@@ -273,8 +273,9 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
             removedPartitions.forEach((tpi) -> consumers.remove(tpi).awaitCompletion());
 
             addedPartitions.forEach((tpi) -> {
-                String key = queueKey + "-" + tpi.getPartition().orElse(-1);
-                TbQueueConsumerTask<M> consumer = new TbQueueConsumerTask<>(key, consumerCreator.apply(config));
+                Integer partitionId = tpi.getPartition().orElse(-1);
+                String key = queueKey + "-" + partitionId;
+                TbQueueConsumerTask<M> consumer = new TbQueueConsumerTask<>(key, () -> consumerCreator.apply(config, partitionId));
                 consumers.put(tpi, consumer);
                 consumer.subscribe(Set.of(tpi));
                 launchConsumer(consumer);
@@ -303,7 +304,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
             }
 
             if (consumer == null) {
-                consumer = new TbQueueConsumerTask<>(queueKey, consumerCreator.apply(config));
+                consumer = new TbQueueConsumerTask<>(queueKey, () -> consumerCreator.apply(config, null)); // no partitionId passed
             }
             consumer.subscribe(partitions);
             if (!consumer.isRunning()) {

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

@@ -249,87 +249,4 @@ public class TbRuleEngineQueueConsumerManager extends MainQueueConsumerManager<T
         }
     }
 
-    private static String partitionsToString(Collection<TopicPartitionInfo> partitions) {
-        return partitions.stream().map(TopicPartitionInfo::getFullTopicName).collect(Collectors.joining(", ", "[", "]"));
-    }
-
-    interface ConsumerWrapper {
-
-        void updatePartitions(Set<TopicPartitionInfo> partitions);
-
-        Collection<TbQueueConsumerTask> getConsumers();
-
-    }
-
-    class ConsumerPerPartitionWrapper implements ConsumerWrapper {
-        private final Map<TopicPartitionInfo, TbQueueConsumerTask> consumers = new HashMap<>();
-
-        @Override
-        public void updatePartitions(Set<TopicPartitionInfo> partitions) {
-            Set<TopicPartitionInfo> addedPartitions = new HashSet<>(partitions);
-            addedPartitions.removeAll(consumers.keySet());
-
-            Set<TopicPartitionInfo> removedPartitions = new HashSet<>(consumers.keySet());
-            removedPartitions.removeAll(partitions);
-            log.info("[{}] Added partitions: {}, removed partitions: {}", queueKey, partitionsToString(addedPartitions), partitionsToString(removedPartitions));
-
-            removedPartitions.forEach((tpi) -> {
-                consumers.get(tpi).initiateStop();
-            });
-
-            List<TbQueueConsumerTask> removedTasks = removedPartitions.stream()
-                    .map(consumers::remove).toList();
-
-//            ctx.getQueueAdmin().createTopicIfNotExists();
-            addedPartitions.forEach((tpi) -> {
-                int partitionId = tpi.getPartition().orElse(-999999);
-                String key = queueKey + "-" + partitionId;
-                TbQueueConsumerTask consumer = new TbQueueConsumerTask(key, () -> ctx.getQueueFactory().createToRuleEngineMsgConsumer(queue, partitionId));
-                consumers.put(tpi, consumer);
-                consumer.subscribe(Set.of(tpi));
-                launchConsumer(consumer);
-            });
-
-            removedTasks.forEach(TbQueueConsumerTask::awaitCompletion);
-        }
-
-        @Override
-        public Collection<TbQueueConsumerTask> getConsumers() {
-            return consumers.values();
-        }
-    }
-
-    class SingleConsumerWrapper implements ConsumerWrapper {
-        private TbQueueConsumerTask consumer;
-
-        @Override
-        public void updatePartitions(Set<TopicPartitionInfo> partitions) {
-            log.info("[{}] New partitions: {}", queueKey, partitionsToString(partitions));
-            if (partitions.isEmpty()) {
-                if (consumer != null && consumer.isRunning()) {
-                    consumer.initiateStop();
-                    consumer.awaitCompletion();
-                }
-                consumer = null;
-                return;
-            }
-
-            if (consumer == null) {
-                consumer = new TbQueueConsumerTask(queueKey, () -> ctx.getQueueFactory().createToRuleEngineMsgConsumer(queue, null));
-            }
-            consumer.subscribe(partitions);
-            if (!consumer.isRunning()) {
-                launchConsumer(consumer);
-            }
-        }
-
-        @Override
-        public Collection<TbQueueConsumerTask> getConsumers() {
-            if (consumer == null) {
-                return Collections.emptyList();
-            }
-            return List.of(consumer);
-        }
-    }
-
 }

application/src/test/java/org/thingsboard/server/service/queue/ruleengine/TbRuleEngineQueueConsumerManagerTest.java

@@ -21,10 +21,10 @@ import lombok.extern.slf4j.Slf4j;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.Mock;
 import org.mockito.Mockito;
-import org.mockito.junit.jupiter.MockitoExtension;
 import org.mockito.junit.jupiter.MockitoSettings;
 import org.mockito.quality.Strictness;
 import org.testcontainers.shaded.org.apache.commons.lang3.RandomUtils;