From 5d7dc19452dc95a1332d7955e369f3b38ed174ec Mon Sep 17 00:00:00 2001
From: Sergey Matvienko <smatvienko@thingsboard.io>
Date: Mon, 13 May 2024 17:21:36 +0200
Subject: [PATCH] TbQueueConsumerTask::awaitCompletion after launchConsumer to
 provide a fully async repartition experience

---
 .../queue/ruleengine/TbRuleEngineQueueConsumerManager.java    | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/application/src/main/java/org/thingsboard/server/service/queue/ruleengine/TbRuleEngineQueueConsumerManager.java b/application/src/main/java/org/thingsboard/server/service/queue/ruleengine/TbRuleEngineQueueConsumerManager.java
index 4df0a900d2..0fb005b928 100644
--- a/application/src/main/java/org/thingsboard/server/service/queue/ruleengine/TbRuleEngineQueueConsumerManager.java
+++ b/application/src/main/java/org/thingsboard/server/service/queue/ruleengine/TbRuleEngineQueueConsumerManager.java
@@ -434,6 +434,8 @@ public class TbRuleEngineQueueConsumerManager {
                 consumers.get(tpi).initiateStop();
             });
 
+            List<TbQueueConsumerTask> removedTasks = removedPartitions.stream()
+                    .map(consumers::remove).toList();
 
 //            ctx.getQueueAdmin().createTopicIfNotExists();
             addedPartitions.forEach((tpi) -> {
@@ -444,6 +446,8 @@ public class TbRuleEngineQueueConsumerManager {
                 consumer.subscribe(Set.of(tpi));
                 launchConsumer(consumer);
             });
+
+            removedTasks.forEach(TbQueueConsumerTask::awaitCompletion);
         }
 
         @Override