Merge pull request #8775 from YevhenBondarenko/feature/recalculate-partitions-improvements

added recalculatePartitions delay for node restarting

application/src/main/resources/thingsboard.yml

@@ -96,6 +96,7 @@ zk:
   session_timeout_ms: "${ZOOKEEPER_SESSION_TIMEOUT_MS:3000}"
   # Name of the directory in zookeeper 'filesystem'
   zk_dir: "${ZOOKEEPER_NODES_DIR:/thingsboard}"
+  recalculate_delay: "${ZOOKEEPER_RECALCULATE_DELAY_MS:60000}"
 
 cluster:
   stats:

common/queue/src/main/java/org/thingsboard/server/queue/discovery/ZkDiscoveryService.java

@@ -16,6 +16,7 @@
 package org.thingsboard.server.queue.discovery;
 
 import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.ProtocolStringList;
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.curator.framework.CuratorFramework;
@@ -44,8 +45,10 @@ import javax.annotation.PostConstruct;
 import javax.annotation.PreDestroy;
 import java.util.List;
 import java.util.NoSuchElementException;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
@@ -66,6 +69,10 @@ public class ZkDiscoveryService implements DiscoveryService, PathChildrenCacheLi
     private Integer zkSessionTimeout;
     @Value("${zk.zk_dir}")
     private String zkDir;
+    @Value("${zk.recalculate_delay:60000}")
+    private Long recalculateDelay;
+
+    protected final ConcurrentHashMap<String, ScheduledFuture<?>> delayedTasks;
 
     private final TbServiceInfoProvider serviceInfoProvider;
     private final PartitionService partitionService;
@@ -82,6 +89,7 @@ public class ZkDiscoveryService implements DiscoveryService, PathChildrenCacheLi
                               PartitionService partitionService) {
         this.serviceInfoProvider = serviceInfoProvider;
         this.partitionService = partitionService;
+        delayedTasks = new ConcurrentHashMap<>();
     }
 
     @PostConstruct
@@ -287,11 +295,39 @@ public class ZkDiscoveryService implements DiscoveryService, PathChildrenCacheLi
             log.error("Failed to decode server instance for node {}", data.getPath(), e);
             throw e;
         }
-        log.debug("Processing [{}] event for [{}]", pathChildrenCacheEvent.getType(), instance.getServiceId());
+
+        String serviceId = instance.getServiceId();
+        ProtocolStringList serviceTypesList = instance.getServiceTypesList();
+
+        log.trace("Processing [{}] event for [{}]", pathChildrenCacheEvent.getType(), serviceId);
         switch (pathChildrenCacheEvent.getType()) {
             case CHILD_ADDED:
+                ScheduledFuture<?> task = delayedTasks.remove(serviceId);
+                if (task != null) {
+                    if (task.cancel(false)) {
+                        log.debug("[{}] Recalculate partitions ignored. Service was restarted in time [{}].",
+                                serviceId, serviceTypesList);
+                    } else {
+                        log.debug("[{}] Going to recalculate partitions. Service was not restarted in time [{}]!",
+                                serviceId, serviceTypesList);
+                        recalculatePartitions();
+                    }
+                } else {
+                    log.trace("[{}] Going to recalculate partitions due to adding new node [{}].",
+                            serviceId, serviceTypesList);
+                    recalculatePartitions();
+                }
+                break;
             case CHILD_REMOVED:
-                recalculatePartitions();
+                ScheduledFuture<?> future = zkExecutorService.schedule(() -> {
+                    log.debug("[{}] Going to recalculate partitions due to removed node [{}]",
+                            serviceId, serviceTypesList);
+                    ScheduledFuture<?> removedTask = delayedTasks.remove(serviceId);
+                    if (removedTask != null) {
+                        recalculatePartitions();
+                    }
+                }, recalculateDelay, TimeUnit.MILLISECONDS);
+                delayedTasks.put(serviceId, future);
                 break;
             default:
                 break;
@@ -303,6 +339,8 @@ public class ZkDiscoveryService implements DiscoveryService, PathChildrenCacheLi
      * Synchronized to ensure that other servers info is up to date
      * */
     synchronized void recalculatePartitions() {
+        delayedTasks.values().forEach(future -> future.cancel(false));
+        delayedTasks.clear();
         partitionService.recalculatePartitions(serviceInfoProvider.getServiceInfo(), getOtherServers());
     }
 

common/queue/src/test/java/org/thingsboard/server/queue/discovery/ZkDiscoveryServiceTest.java

@@ -0,0 +1,189 @@
+/**
+ * Copyright © 2016-2023 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.discovery;
+
+import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.framework.imps.CuratorFrameworkState;
+import org.apache.curator.framework.recipes.cache.ChildData;
+import org.apache.curator.framework.recipes.cache.PathChildrenCache;
+import org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnitRunner;
+import org.springframework.test.util.ReflectionTestUtils;
+import org.thingsboard.common.util.ThingsBoardThreadFactory;
+import org.thingsboard.server.gen.transport.TransportProtos;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+
+import static org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent.Type.CHILD_ADDED;
+import static org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent.Type.CHILD_REMOVED;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@RunWith(MockitoJUnitRunner.class)
+public class ZkDiscoveryServiceTest {
+
+    @Mock
+    private TbServiceInfoProvider serviceInfoProvider;
+
+    @Mock
+    private PartitionService partitionService;
+
+    @Mock
+    private CuratorFramework client;
+
+    @Mock
+    private PathChildrenCache cache;
+
+    @Mock
+    private CuratorFramework curatorFramework;
+
+    private ZkDiscoveryService zkDiscoveryService;
+
+    private static final long RECALCULATE_DELAY = 100L;
+
+    final TransportProtos.ServiceInfo currentInfo = TransportProtos.ServiceInfo.newBuilder().setServiceId("tb-rule-engine-0").build();
+    final ChildData currentData = new ChildData("/thingsboard/nodes/0000000010", null, currentInfo.toByteArray());
+    final TransportProtos.ServiceInfo childInfo = TransportProtos.ServiceInfo.newBuilder().setServiceId("tb-rule-engine-1").build();
+    final ChildData childData = new ChildData("/thingsboard/nodes/0000000020", null, childInfo.toByteArray());
+
+    @Before
+    public void setup() {
+        zkDiscoveryService = Mockito.spy(new ZkDiscoveryService(serviceInfoProvider, partitionService));
+        ScheduledExecutorService zkExecutorService = Executors.newSingleThreadScheduledExecutor(ThingsBoardThreadFactory.forName("zk-discovery"));
+        when(client.getState()).thenReturn(CuratorFrameworkState.STARTED);
+        ReflectionTestUtils.setField(zkDiscoveryService, "stopped", false);
+        ReflectionTestUtils.setField(zkDiscoveryService, "client", client);
+        ReflectionTestUtils.setField(zkDiscoveryService, "cache", cache);
+        ReflectionTestUtils.setField(zkDiscoveryService, "nodePath", "/thingsboard/nodes/0000000010");
+        ReflectionTestUtils.setField(zkDiscoveryService, "zkExecutorService", zkExecutorService);
+        ReflectionTestUtils.setField(zkDiscoveryService, "recalculateDelay", RECALCULATE_DELAY);
+        ReflectionTestUtils.setField(zkDiscoveryService, "zkDir", "/thingsboard");
+
+        when(serviceInfoProvider.getServiceInfo()).thenReturn(currentInfo);
+
+        List<ChildData> dataList = new ArrayList<>();
+        dataList.add(currentData);
+        when(cache.getCurrentData()).thenReturn(dataList);
+    }
+
+    @Test
+    public void restartNodeInTimeTest() throws Exception {
+        startNode(childData);
+
+        verify(partitionService, times(1)).recalculatePartitions(eq(currentInfo), eq(List.of(childInfo)));
+
+        reset(partitionService);
+
+        stopNode(childData);
+
+        assertEquals(1, zkDiscoveryService.delayedTasks.size());
+
+        verify(partitionService, never()).recalculatePartitions(any(), any());
+
+        startNode(childData);
+
+        verify(partitionService, never()).recalculatePartitions(any(), any());
+
+        Thread.sleep(RECALCULATE_DELAY * 2);
+
+        verify(partitionService, never()).recalculatePartitions(any(), any());
+
+        assertTrue(zkDiscoveryService.delayedTasks.isEmpty());
+    }
+
+    @Test
+    public void restartNodeNotInTimeTest() throws Exception {
+        startNode(childData);
+
+        verify(partitionService, times(1)).recalculatePartitions(eq(currentInfo), eq(List.of(childInfo)));
+
+        reset(partitionService);
+
+        stopNode(childData);
+
+        assertEquals(1, zkDiscoveryService.delayedTasks.size());
+
+        Thread.sleep(RECALCULATE_DELAY * 2);
+
+        assertTrue(zkDiscoveryService.delayedTasks.isEmpty());
+
+        startNode(childData);
+
+        verify(partitionService, times(1)).recalculatePartitions(eq(currentInfo), eq(Collections.emptyList()));
+
+        verify(partitionService, times(1)).recalculatePartitions(eq(currentInfo), eq(List.of(childInfo)));
+
+        reset(partitionService);
+    }
+
+    @Test
+    public void startAnotherNodeDuringRestartTest() throws Exception {
+        var anotherInfo = TransportProtos.ServiceInfo.newBuilder().setServiceId("tb-transport").build();
+        var anotherData = new ChildData("/thingsboard/nodes/0000000030", null, anotherInfo.toByteArray());
+
+        startNode(childData);
+
+        verify(partitionService, times(1)).recalculatePartitions(eq(currentInfo), eq(List.of(childInfo)));
+
+        reset(partitionService);
+
+        stopNode(childData);
+
+        assertEquals(1, zkDiscoveryService.delayedTasks.size());
+
+        startNode(anotherData);
+
+        assertTrue(zkDiscoveryService.delayedTasks.isEmpty());
+
+        verify(partitionService, times(1)).recalculatePartitions(eq(currentInfo), eq(List.of(anotherInfo)));
+        reset(partitionService);
+
+        Thread.sleep(RECALCULATE_DELAY * 2);
+
+        verify(partitionService, never()).recalculatePartitions(any(), any());
+
+        startNode(childData);
+
+        verify(partitionService, times(1)).recalculatePartitions(eq(currentInfo), eq(List.of(anotherInfo, childInfo)));
+    }
+
+    private void startNode(ChildData data) throws Exception {
+        cache.getCurrentData().add(data);
+        zkDiscoveryService.childEvent(curatorFramework, new PathChildrenCacheEvent(CHILD_ADDED, data));
+    }
+
+    private void stopNode(ChildData data) throws Exception {
+        cache.getCurrentData().remove(data);
+        zkDiscoveryService.childEvent(curatorFramework, new PathChildrenCacheEvent(CHILD_REMOVED, data));
+    }
+
+}

msa/vc-executor/src/main/resources/tb-vc-executor.yml

@@ -41,6 +41,7 @@ zk:
   session_timeout_ms: "${ZOOKEEPER_SESSION_TIMEOUT_MS:3000}"
   # Name of the directory in zookeeper 'filesystem'
   zk_dir: "${ZOOKEEPER_NODES_DIR:/thingsboard}"
+  recalculate_delay: "${ZOOKEEPER_RECALCULATE_DELAY_MS:60000}"
 
 queue:
   type: "${TB_QUEUE_TYPE:kafka}" # in-memory or kafka (Apache Kafka) or aws-sqs (AWS SQS) or pubsub (PubSub) or service-bus (Azure Service Bus) or rabbitmq (RabbitMQ)

transport/coap/src/main/resources/tb-coap-transport.yml

@@ -41,6 +41,7 @@ zk:
   session_timeout_ms: "${ZOOKEEPER_SESSION_TIMEOUT_MS:3000}"
   # Name of the directory in zookeeper 'filesystem'
   zk_dir: "${ZOOKEEPER_NODES_DIR:/thingsboard}"
+  recalculate_delay: "${ZOOKEEPER_RECALCULATE_DELAY_MS:60000}"
 
 cache:
   type: "${CACHE_TYPE:redis}"

transport/http/src/main/resources/tb-http-transport.yml

@@ -68,6 +68,7 @@ zk:
   session_timeout_ms: "${ZOOKEEPER_SESSION_TIMEOUT_MS:3000}"
   # Name of the directory in zookeeper 'filesystem'
   zk_dir: "${ZOOKEEPER_NODES_DIR:/thingsboard}"
+  recalculate_delay: "${ZOOKEEPER_RECALCULATE_DELAY_MS:60000}"
 
 cache:
   type: "${CACHE_TYPE:redis}"

transport/lwm2m/src/main/resources/tb-lwm2m-transport.yml

@@ -41,6 +41,7 @@ zk:
   session_timeout_ms: "${ZOOKEEPER_SESSION_TIMEOUT_MS:3000}"
   # Name of the directory in zookeeper 'filesystem'
   zk_dir: "${ZOOKEEPER_NODES_DIR:/thingsboard}"
+  recalculate_delay: "${ZOOKEEPER_RECALCULATE_DELAY_MS:60000}"
 
 cache:
   type: "${CACHE_TYPE:redis}"

transport/mqtt/src/main/resources/tb-mqtt-transport.yml

@@ -41,6 +41,7 @@ zk:
   session_timeout_ms: "${ZOOKEEPER_SESSION_TIMEOUT_MS:3000}"
   # Name of the directory in zookeeper 'filesystem'
   zk_dir: "${ZOOKEEPER_NODES_DIR:/thingsboard}"
+  recalculate_delay: "${ZOOKEEPER_RECALCULATE_DELAY_MS:60000}"
 
 cache:
   type: "${CACHE_TYPE:redis}"

transport/snmp/src/main/resources/tb-snmp-transport.yml

@@ -41,6 +41,7 @@ zk:
   session_timeout_ms: "${ZOOKEEPER_SESSION_TIMEOUT_MS:3000}"
   # Name of the directory in zookeeper 'filesystem'
   zk_dir: "${ZOOKEEPER_NODES_DIR:/thingsboard}"
+  recalculate_delay: "${ZOOKEEPER_RECALCULATE_DELAY_MS:60000}"
 
 cache:
   type: "${CACHE_TYPE:redis}"