Merge pull request #8129 from smatvienko-tb/feature/web-service-handler-subscription-improvements

[3.5] WebSocket subscription flow fix, WS refactoring and test improvements

application/src/main/java/org/thingsboard/server/config/WebSocketConfiguration.java

@@ -15,6 +15,8 @@
  */
 package org.thingsboard.server.config;
 
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.http.HttpStatus;
@@ -40,11 +42,15 @@ import java.util.Map;
 @Configuration
 @TbCoreComponent
 @EnableWebSocket
+@RequiredArgsConstructor
+@Slf4j
 public class WebSocketConfiguration implements WebSocketConfigurer {
 
     public static final String WS_PLUGIN_PREFIX = "/api/ws/plugins/";
     private static final String WS_PLUGIN_MAPPING = WS_PLUGIN_PREFIX + "**";
 
+    private final WebSocketHandler wsHandler;
+
     @Bean
     public ServletServerContainerFactoryBean createWebSocketContainer() {
         ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();
@@ -55,7 +61,11 @@ public class WebSocketConfiguration implements WebSocketConfigurer {
 
     @Override
     public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
-        registry.addHandler(wsHandler(), WS_PLUGIN_MAPPING).setAllowedOriginPatterns("*")
+        if (!(wsHandler instanceof TbWebSocketHandler)) {
+            log.error("TbWebSocketHandler expected but [{}] provided", wsHandler);
+            throw new RuntimeException("TbWebSocketHandler expected but " + wsHandler + " provided");
+        }
+        registry.addHandler(wsHandler, WS_PLUGIN_MAPPING).setAllowedOriginPatterns("*")
                 .addInterceptors(new HttpSessionHandshakeInterceptor(), new HandshakeInterceptor() {
 
                     @Override
@@ -82,11 +92,6 @@ public class WebSocketConfiguration implements WebSocketConfigurer {
                 });
     }
 
-    @Bean
-    public WebSocketHandler wsHandler() {
-        return new TbWebSocketHandler();
-    }
-
     protected SecurityUser getCurrentUser() throws ThingsboardException {
         Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
         if (authentication != null && authentication.getPrincipal() instanceof SecurityUser) {

application/src/main/java/org/thingsboard/server/controller/plugin/TbWebSocketHandler.java

@@ -67,8 +67,8 @@ import static org.thingsboard.server.service.telemetry.DefaultTelemetryWebSocket
 @Slf4j
 public class TbWebSocketHandler extends TextWebSocketHandler implements TelemetryWebSocketMsgEndpoint {
 
-    private static final ConcurrentMap<String, SessionMetaData> internalSessionMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<String, SessionMetaData> internalSessionMap = new ConcurrentHashMap<>();
-    private static final ConcurrentMap<String, String> externalSessionMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<String, String> externalSessionMap = new ConcurrentHashMap<>();
 
 
     @Autowired
@@ -82,13 +82,13 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
     @Value("${server.ws.ping_timeout:30000}")
     private long pingTimeout;
 
-    private ConcurrentMap<String, TelemetryWebSocketSessionRef> blacklistedSessions = new ConcurrentHashMap<>();
+    private final ConcurrentMap<String, TelemetryWebSocketSessionRef> blacklistedSessions = new ConcurrentHashMap<>();
-    private ConcurrentMap<String, TbRateLimits> perSessionUpdateLimits = new ConcurrentHashMap<>();
+    private final ConcurrentMap<String, TbRateLimits> perSessionUpdateLimits = new ConcurrentHashMap<>();
 
-    private ConcurrentMap<TenantId, Set<String>> tenantSessionsMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<TenantId, Set<String>> tenantSessionsMap = new ConcurrentHashMap<>();
-    private ConcurrentMap<CustomerId, Set<String>> customerSessionsMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<CustomerId, Set<String>> customerSessionsMap = new ConcurrentHashMap<>();
-    private ConcurrentMap<UserId, Set<String>> regularUserSessionsMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<UserId, Set<String>> regularUserSessionsMap = new ConcurrentHashMap<>();
-    private ConcurrentMap<UserId, Set<String>> publicUserSessionsMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<UserId, Set<String>> publicUserSessionsMap = new ConcurrentHashMap<>();
 
     @Override
     public void handleTextMessage(WebSocketSession session, TextMessage message) {

application/src/main/java/org/thingsboard/server/service/telemetry/DefaultTelemetryWebSocketService.java

@@ -16,7 +16,6 @@
 package org.thingsboard.server.service.telemetry;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Function;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
@@ -27,6 +26,7 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Service;
 import org.springframework.web.socket.CloseStatus;
+import org.thingsboard.common.util.JacksonUtil;
 import org.thingsboard.common.util.ThingsBoardExecutors;
 import org.thingsboard.common.util.ThingsBoardThreadFactory;
 import org.thingsboard.server.common.data.DataConstants;
@@ -95,6 +95,8 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -112,7 +114,6 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
     private static final Aggregation DEFAULT_AGGREGATION = Aggregation.NONE;
     private static final int UNKNOWN_SUBSCRIPTION_ID = 0;
     private static final String PROCESSING_MSG = "[{}] Processing: {}";
-    private static final ObjectMapper jsonMapper = new ObjectMapper();
     private static final String FAILED_TO_FETCH_DATA = "Failed to fetch data!";
     private static final String FAILED_TO_FETCH_ATTRIBUTES = "Failed to fetch attributes!";
     private static final String SESSION_META_DATA_NOT_FOUND = "Session meta-data not found!";
@@ -147,10 +148,10 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
     @Value("${server.ws.ping_timeout:30000}")
     private long pingTimeout;
 
-    private ConcurrentMap<TenantId, Set<String>> tenantSubscriptionsMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<TenantId, Set<String>> tenantSubscriptionsMap = new ConcurrentHashMap<>();
-    private ConcurrentMap<CustomerId, Set<String>> customerSubscriptionsMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<CustomerId, Set<String>> customerSubscriptionsMap = new ConcurrentHashMap<>();
-    private ConcurrentMap<UserId, Set<String>> regularUserSubscriptionsMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<UserId, Set<String>> regularUserSubscriptionsMap = new ConcurrentHashMap<>();
-    private ConcurrentMap<UserId, Set<String>> publicUserSubscriptionsMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<UserId, Set<String>> publicUserSubscriptionsMap = new ConcurrentHashMap<>();
 
     private ExecutorService executor;
     private String serviceId;
@@ -204,7 +205,7 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
         }
 
         try {
-            TelemetryPluginCmdsWrapper cmdsWrapper = jsonMapper.readValue(msg, TelemetryPluginCmdsWrapper.class);
+            TelemetryPluginCmdsWrapper cmdsWrapper = JacksonUtil.OBJECT_MAPPER.readValue(msg, TelemetryPluginCmdsWrapper.class);
             if (cmdsWrapper != null) {
                 if (cmdsWrapper.getAttrSubCmds() != null) {
                     cmdsWrapper.getAttrSubCmds().forEach(cmd -> {
@@ -450,7 +451,6 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
             @Override
             public void onSuccess(List<AttributeKvEntry> data) {
                 List<TsKvEntry> attributesData = data.stream().map(d -> new BasicTsKvEntry(d.getLastUpdateTs(), d)).collect(Collectors.toList());
-                sendWsMsg(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), attributesData));
 
                 Map<String, Long> subState = new HashMap<>(keys.size());
                 keys.forEach(key -> subState.put(key, 0L));
@@ -458,6 +458,7 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
 
                 TbAttributeSubscriptionScope scope = StringUtils.isEmpty(cmd.getScope()) ? TbAttributeSubscriptionScope.ANY_SCOPE : TbAttributeSubscriptionScope.valueOf(cmd.getScope());
 
+                Lock subLock = new ReentrantLock();
                 TbAttributeSubscription sub = TbAttributeSubscription.builder()
                         .serviceId(serviceId)
                         .sessionId(sessionId)
@@ -467,9 +468,24 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
                         .allKeys(false)
                         .keyStates(subState)
                         .scope(scope)
-                        .updateConsumer(DefaultTelemetryWebSocketService.this::sendWsMsg)
+                        .updateConsumer((sessionId, update) -> {
+                            subLock.lock();
+                            try {
+                                sendWsMsg(sessionId, update);
+                            } finally {
+                                subLock.unlock();
+                            }
+                        })
                         .build();
-                oldSubService.addSubscription(sub);
+
+                subLock.lock();
+                try{
+                    oldSubService.addSubscription(sub);
+                    sendWsMsg(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), attributesData));
+                } finally {
+                    subLock.unlock();
+                }
+
             }
 
             @Override
@@ -550,13 +566,13 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
             @Override
             public void onSuccess(List<AttributeKvEntry> data) {
                 List<TsKvEntry> attributesData = data.stream().map(d -> new BasicTsKvEntry(d.getLastUpdateTs(), d)).collect(Collectors.toList());
-                sendWsMsg(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), attributesData));
 
                 Map<String, Long> subState = new HashMap<>(attributesData.size());
                 attributesData.forEach(v -> subState.put(v.getKey(), v.getTs()));
 
                 TbAttributeSubscriptionScope scope = StringUtils.isEmpty(cmd.getScope()) ? TbAttributeSubscriptionScope.ANY_SCOPE : TbAttributeSubscriptionScope.valueOf(cmd.getScope());
 
+                Lock subLock = new ReentrantLock();
                 TbAttributeSubscription sub = TbAttributeSubscription.builder()
                         .serviceId(serviceId)
                         .sessionId(sessionId)
@@ -565,9 +581,24 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
                         .entityId(entityId)
                         .allKeys(true)
                         .keyStates(subState)
-                        .updateConsumer(DefaultTelemetryWebSocketService.this::sendWsMsg)
+                        .updateConsumer((sessionId, update) -> {
-                        .scope(scope).build();
+                            subLock.lock();
-                oldSubService.addSubscription(sub);
+                            try {
+                                sendWsMsg(sessionId, update);
+                            } finally {
+                                subLock.unlock();
+                            }
+                        })
+                        .scope(scope)
+                        .build();
+
+                subLock.lock();
+                try {
+                    oldSubService.addSubscription(sub);
+                    sendWsMsg(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), attributesData));
+                } finally {
+                    subLock.unlock();
+                }
             }
 
             @Override
@@ -636,20 +667,34 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
         FutureCallback<List<TsKvEntry>> callback = new FutureCallback<List<TsKvEntry>>() {
             @Override
             public void onSuccess(List<TsKvEntry> data) {
-                sendWsMsg(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), data));
                 Map<String, Long> subState = new HashMap<>(data.size());
                 data.forEach(v -> subState.put(v.getKey(), v.getTs()));
 
+                Lock subLock = new ReentrantLock();
                 TbTimeseriesSubscription sub = TbTimeseriesSubscription.builder()
                         .serviceId(serviceId)
                         .sessionId(sessionId)
                         .subscriptionId(cmd.getCmdId())
                         .tenantId(sessionRef.getSecurityCtx().getTenantId())
                         .entityId(entityId)
-                        .updateConsumer(DefaultTelemetryWebSocketService.this::sendWsMsg)
+                        .updateConsumer((sessionId, update) -> {
+                            subLock.lock();
+                            try {
+                                sendWsMsg(sessionId, update);
+                            } finally {
+                                subLock.unlock();
+                            }
+                        })
                         .allKeys(true)
                         .keyStates(subState).build();
-                oldSubService.addSubscription(sub);
+
+                subLock.lock();
+                try {
+                    oldSubService.addSubscription(sub);
+                    sendWsMsg(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), data));
+                } finally {
+                    subLock.unlock();
+                }
             }
 
             @Override
@@ -673,21 +718,35 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
         return new FutureCallback<>() {
             @Override
             public void onSuccess(List<TsKvEntry> data) {
-                sendWsMsg(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), data));
                 Map<String, Long> subState = new HashMap<>(keys.size());
                 keys.forEach(key -> subState.put(key, startTs));
                 data.forEach(v -> subState.put(v.getKey(), v.getTs()));
 
+                Lock subLock = new ReentrantLock();
                 TbTimeseriesSubscription sub = TbTimeseriesSubscription.builder()
                         .serviceId(serviceId)
                         .sessionId(sessionId)
                         .subscriptionId(cmd.getCmdId())
                         .tenantId(sessionRef.getSecurityCtx().getTenantId())
                         .entityId(entityId)
-                        .updateConsumer(DefaultTelemetryWebSocketService.this::sendWsMsg)
+                        .updateConsumer((sessionId, update) -> {
+                            subLock.lock();
+                            try {
+                                sendWsMsg(sessionId, update);
+                            } finally {
+                                subLock.unlock();
+                            }
+                        })
                         .allKeys(false)
                         .keyStates(subState).build();
-                oldSubService.addSubscription(sub);
+
+                subLock.lock();
+                try{
+                    oldSubService.addSubscription(sub);
+                    sendWsMsg(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), data));
+                } finally {
+                    subLock.unlock();
+                }
             }
 
             @Override
@@ -793,7 +852,7 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
 
     private void sendWsMsg(TelemetryWebSocketSessionRef sessionRef, int cmdId, Object update) {
         try {
-            String msg = jsonMapper.writeValueAsString(update);
+            String msg = JacksonUtil.OBJECT_MAPPER.writeValueAsString(update);
             executor.submit(() -> {
                 try {
                     msgEndpoint.send(sessionRef, cmdId, msg);

application/src/test/java/org/thingsboard/server/controller/AbstractControllerTest.java

@@ -52,7 +52,7 @@ public abstract class AbstractControllerTest extends AbstractNotifyEntityTest {
     @LocalServerPort
     protected int wsPort;
 
-    private TbTestWebSocketClient wsClient; // lazy
+    private volatile TbTestWebSocketClient wsClient; // lazy
 
     public TbTestWebSocketClient getWsClient() {
         if (wsClient == null) {

application/src/test/java/org/thingsboard/server/controller/BaseWebsocketApiTest.java

@@ -548,7 +548,7 @@ public abstract class BaseWebsocketApiTest extends AbstractControllerTest {
         SingleEntityFilter entityFilter = new SingleEntityFilter();
         entityFilter.setSingleEntity(tenantId);
 
-        assertThatNoException().isThrownBy(() -> {
+        assertThatNoException().as("subscribeForAttributes").isThrownBy(() -> {
             JsonNode update = getWsClient().subscribeForAttributes(tenantId, TbAttributeSubscriptionScope.SERVER_SCOPE.name(), List.of("attr"));
             assertThat(update.get("errorMsg").isNull()).isTrue();
             assertThat(update.get("errorCode").asInt()).isEqualTo(SubscriptionErrorCode.NO_ERROR.getCode());
@@ -560,7 +560,7 @@ public abstract class BaseWebsocketApiTest extends AbstractControllerTest {
                 new BaseAttributeKvEntry(System.currentTimeMillis(), new StringDataEntry("attr", expectedAttrValue))
         ));
         JsonNode update = JacksonUtil.toJsonNode(getWsClient().waitForUpdate());
-        assertThat(update).isNotNull();
+        assertThat(update).as("waitForUpdate").isNotNull();
         assertThat(update.get("data").get("attr").get(0).get(1).asText()).isEqualTo(expectedAttrValue);
     }
 
@@ -569,15 +569,17 @@ public abstract class BaseWebsocketApiTest extends AbstractControllerTest {
         tsService.saveAndNotify(device.getTenantId(), null, device.getId(), tsData, 0, new FutureCallback<Void>() {
             @Override
             public void onSuccess(@Nullable Void result) {
+                log.debug("sendTelemetry callback onSuccess");
                 latch.countDown();
             }
 
             @Override
             public void onFailure(Throwable t) {
+                log.error("Failed to send telemetry", t);
                 latch.countDown();
             }
         });
-        latch.await(3, TimeUnit.SECONDS);
+        assertThat(latch.await(TIMEOUT, TimeUnit.SECONDS)).as("await sendTelemetry callback");
     }
 
     private void sendAttributes(Device device, TbAttributeSubscriptionScope scope, List<AttributeKvEntry> attrData) throws InterruptedException {
@@ -589,14 +591,16 @@ public abstract class BaseWebsocketApiTest extends AbstractControllerTest {
         tsService.saveAndNotify(tenantId, entityId, scope.name(), attrData, new FutureCallback<Void>() {
             @Override
             public void onSuccess(@Nullable Void result) {
+                log.debug("sendAttributes callback onSuccess");
                 latch.countDown();
             }
 
             @Override
             public void onFailure(Throwable t) {
+                log.error("Failed to sendAttributes", t);
                 latch.countDown();
             }
         });
-        latch.await(3, TimeUnit.SECONDS);
+        assertThat(latch.await(TIMEOUT, TimeUnit.SECONDS)).as("await sendAttributes callback").isTrue();
     }
 }

application/src/test/java/org/thingsboard/server/controller/TbTestWebSocketClient.java

@@ -47,6 +47,7 @@ import java.util.concurrent.TimeUnit;
 @Slf4j
 public class TbTestWebSocketClient extends WebSocketClient {
 
+    private static final long TIMEOUT = TimeUnit.SECONDS.toMillis(30);
     private volatile String lastMsg;
     private volatile CountDownLatch reply;
     private volatile CountDownLatch update;
@@ -87,12 +88,14 @@ public class TbTestWebSocketClient extends WebSocketClient {
     }
 
     public void registerWaitForUpdate(int count) {
+        log.debug("registerWaitForUpdate [{}]", count);
         lastMsg = null;
         update = new CountDownLatch(count);
     }
 
     @Override
     public void send(String text) throws NotYetConnectedException {
+        log.debug("send [{}]", text);
         reply = new CountDownLatch(1);
         super.send(text);
     }
@@ -110,21 +113,31 @@ public class TbTestWebSocketClient extends WebSocketClient {
     }
 
     public String waitForUpdate() {
-        return waitForUpdate(TimeUnit.SECONDS.toMillis(3));
+        return waitForUpdate(TIMEOUT);
     }
 
     public String waitForUpdate(long ms) {
+        log.debug("waitForUpdate [{}]", ms);
         try {
-            update.await(ms, TimeUnit.MILLISECONDS);
+            if (!update.await(ms, TimeUnit.MILLISECONDS)) {
+                log.warn("Failed to await update (waiting time [{}]ms elapsed)", ms, new RuntimeException("stacktrace"));
+            }
         } catch (InterruptedException e) {
-            log.warn("Failed to await reply", e);
+            log.warn("Failed to await update", e);
         }
         return lastMsg;
     }
 
     public String waitForReply() {
+        return waitForReply(TIMEOUT);
+    }
+
+    public String waitForReply(long ms) {
+        log.debug("waitForReply [{}]", ms);
         try {
-            reply.await(3, TimeUnit.SECONDS);
+            if (!reply.await(ms, TimeUnit.MILLISECONDS)) {
+                log.warn("Failed to await reply (waiting time [{}]ms elapsed)", ms, new RuntimeException("stacktrace"));
+            }
         } catch (InterruptedException e) {
             log.warn("Failed to await reply", e);
         }