From 134db4249472f094523e35d2bb3efeb6f64be26f Mon Sep 17 00:00:00 2001
From: ViacheslavKlimov <viacheslavklimov11@gmail.com>
Date: Tue, 19 Mar 2024 13:49:46 +0200
Subject: [PATCH] Configurable consumer per partition for tb-core queue

---
 .../server/service/queue/DefaultTbCoreConsumerService.java   | 5 ++++-
 .../service/queue/processing/AbstractConsumerService.java    | 3 +--
 application/src/main/resources/thingsboard.yml               | 2 ++
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCoreConsumerService.java b/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCoreConsumerService.java
index c0ccaea9d0..474efc5ed7 100644
--- a/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCoreConsumerService.java
+++ b/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCoreConsumerService.java
@@ -126,6 +126,8 @@ public class DefaultTbCoreConsumerService extends AbstractConsumerService<ToCore
     private long pollInterval;
     @Value("${queue.core.pack-processing-timeout}")
     private long packProcessingTimeout;
+    @Value("${queue.core.consumer-per-partition-enabled:true}")
+    private boolean consumerPerPartitionEnabled;
     @Value("${queue.core.stats.enabled:false}")
     private boolean statsEnabled;
 
@@ -199,7 +201,7 @@ public class DefaultTbCoreConsumerService extends AbstractConsumerService<ToCore
 
         this.mainConsumer = QueueConsumerManager.<TbProtoQueueMsg<ToCoreMsg>, CoreQueueConfig>builder()
                 .queueKey(new QueueKey(ServiceType.TB_CORE))
-                .config(CoreQueueConfig.of(true, (int) pollInterval))
+                .config(CoreQueueConfig.of(consumerPerPartitionEnabled, (int) pollInterval))
                 .msgPackProcessor(this::processMsgs)
                 .consumerCreator(config -> queueFactory.createToCoreMsgConsumer())
                 .consumerExecutor(consumersExecutor)
@@ -764,6 +766,7 @@ public class DefaultTbCoreConsumerService extends AbstractConsumerService<ToCore
 
     @Override
     protected void stopConsumers() {
+        super.stopConsumers();
         mainConsumer.stop();
         mainConsumer.awaitStop();
         usageStatsConsumer.stop();
diff --git a/application/src/main/java/org/thingsboard/server/service/queue/processing/AbstractConsumerService.java b/application/src/main/java/org/thingsboard/server/service/queue/processing/AbstractConsumerService.java
index e2a6a3f117..747e53a0bf 100644
--- a/application/src/main/java/org/thingsboard/server/service/queue/processing/AbstractConsumerService.java
+++ b/application/src/main/java/org/thingsboard/server/service/queue/processing/AbstractConsumerService.java
@@ -16,7 +16,6 @@
 package org.thingsboard.server.service.queue.processing;
 
 import lombok.extern.slf4j.Slf4j;
-import org.springframework.boot.context.event.ApplicationReadyEvent;
 import org.springframework.context.ApplicationEventPublisher;
 import org.thingsboard.common.util.ThingsBoardExecutors;
 import org.thingsboard.common.util.ThingsBoardThreadFactory;
@@ -112,7 +111,7 @@ public abstract class AbstractConsumerService<N extends com.google.protobuf.Gene
     }
 
     @AfterStartUp(order = AfterStartUp.REGULAR_SERVICE)
-    public void onApplicationEvent(ApplicationReadyEvent event) {
+    public void afterStartUp() {
         startConsumers();
     }
 
diff --git a/application/src/main/resources/thingsboard.yml b/application/src/main/resources/thingsboard.yml
index 8664eaa054..a0709d8bc7 100644
--- a/application/src/main/resources/thingsboard.yml
+++ b/application/src/main/resources/thingsboard.yml
@@ -1562,6 +1562,8 @@ queue:
     partitions: "${TB_QUEUE_CORE_PARTITIONS:10}"
     # Timeout for processing a message pack by Core microservices
     pack-processing-timeout: "${TB_QUEUE_CORE_PACK_PROCESSING_TIMEOUT_MS:2000}"
+    # Enable/disable a separate consumer per partition for Core queue
+    consumer-per-partition-enabled: "${TB_QUEUE_CORE_CONSUMER_PER_PARTITION_ENABLED:true}"
     ota:
       # Default topic name for OTA updates
       topic: "${TB_QUEUE_CORE_OTA_TOPIC:tb_ota_package}"