From dfb633b8e8a31a19b890c3158ed2067f172f0458 Mon Sep 17 00:00:00 2001
From: YevhenBondarenko <ybondarenko@thingsboard.io>
Date: Wed, 11 Sep 2024 11:26:04 +0200
Subject: [PATCH] improvements

---
 .../mqtt/gateway/GatewayLatencyService.java         |  6 +-----
 .../mqtt/gateway/latency/GatewayLatencyState.java   | 13 ++++++++-----
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/gateway/GatewayLatencyService.java b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/gateway/GatewayLatencyService.java
index 148a8819ee..acf10a03b8 100644
--- a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/gateway/GatewayLatencyService.java
+++ b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/gateway/GatewayLatencyService.java
@@ -67,10 +67,7 @@ public class GatewayLatencyService {
     }
 
     public void onDeviceDelete(DeviceId deviceId) {
-        var state = states.remove(deviceId);
-        if (state != null) {
-            state.clear();
-        }
+        states.remove(deviceId);
     }
 
     public void onDeviceDisconnect(DeviceId deviceId) {
@@ -100,7 +97,6 @@ public class GatewayLatencyService {
             return;
         }
         var result = state.getLatencyStateResult();
-        state.clear();
         var kvProto = TransportProtos.KeyValueProto.newBuilder()
                 .setKey("latencyCheck")
                 .setType(TransportProtos.KeyValueType.JSON_V)
diff --git a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/gateway/latency/GatewayLatencyState.java b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/gateway/latency/GatewayLatencyState.java
index c139bb3f19..3ad75fee75 100644
--- a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/gateway/latency/GatewayLatencyState.java
+++ b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/gateway/latency/GatewayLatencyState.java
@@ -54,13 +54,16 @@ public class GatewayLatencyState {
         }
     }
 
-    public void clear() {
-        connectors.clear();
-    }
-
     public Map<String, ConnectorLatencyResult> getLatencyStateResult() {
         Map<String, ConnectorLatencyResult> result = new HashMap<>();
-        connectors.forEach((name, state) -> result.put(name, state.getResult()));
+        updateLock.lock();
+        try {
+            connectors.forEach((name, state) -> result.put(name, state.getResult()));
+            connectors.clear();
+        } finally {
+            updateLock.unlock();
+        }
+
         return result;
     }