tsai
/
thingsboard
mirror of https://github.com/thingsboard/thingsboard.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

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

[3.5] WebSocket subscription flow fix, WS refactoring and test improvements
pull/8162/head
Andrew Shvayka 4 years ago
committed by GitHub
parent
baddd45d82 71d9f5ef68
commit
f0ca2d48f7
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 125 additions and 44 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 17
      application/src/main/java/org/thingsboard/server/config/WebSocketConfiguration.java
  2. 16
      application/src/main/java/org/thingsboard/server/controller/plugin/TbWebSocketHandler.java
  3. 101
      application/src/main/java/org/thingsboard/server/service/telemetry/DefaultTelemetryWebSocketService.java
  4. 2
      application/src/test/java/org/thingsboard/server/controller/AbstractControllerTest.java
  5. 12
      application/src/test/java/org/thingsboard/server/controller/BaseWebsocketApiTest.java
  6. 21
      application/src/test/java/org/thingsboard/server/controller/TbTestWebSocketClient.java

17
application/src/main/java/org/thingsboard/server/config/WebSocketConfiguration.java
View File

@ -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) {

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

@ -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 static final ConcurrentMap<String, String> externalSessionMap = new ConcurrentHashMap<>();
private final ConcurrentMap<String, SessionMetaData> internalSessionMap = 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 ConcurrentMap<String, TbRateLimits> perSessionUpdateLimits = new ConcurrentHashMap<>();
private final ConcurrentMap<String, TelemetryWebSocketSessionRef> blacklistedSessions = new ConcurrentHashMap<>();
private final ConcurrentMap<String, TbRateLimits> perSessionUpdateLimits = new ConcurrentHashMap<>();
private ConcurrentMap<TenantId, Set<String>> tenantSessionsMap = new ConcurrentHashMap<>();
private ConcurrentMap<CustomerId, Set<String>> customerSessionsMap = new ConcurrentHashMap<>();
private ConcurrentMap<UserId, Set<String>> regularUserSessionsMap = new ConcurrentHashMap<>();
private ConcurrentMap<UserId, Set<String>> publicUserSessionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<TenantId, Set<String>> tenantSessionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<CustomerId, Set<String>> customerSessionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<UserId, Set<String>> regularUserSessionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<UserId, Set<String>> publicUserSessionsMap = new ConcurrentHashMap<>();
@Override
public void handleTextMessage(WebSocketSession session, TextMessage message) {

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

@ -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 ConcurrentMap<CustomerId, Set<String>> customerSubscriptionsMap = new ConcurrentHashMap<>();
private ConcurrentMap<UserId, Set<String>> regularUserSubscriptionsMap = new ConcurrentHashMap<>();
private ConcurrentMap<UserId, Set<String>> publicUserSubscriptionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<TenantId, Set<String>> tenantSubscriptionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<CustomerId, Set<String>> customerSubscriptionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<UserId, Set<String>> regularUserSubscriptionsMap = 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)
.scope(scope).build();
oldSubService.addSubscription(sub);
.updateConsumer((sessionId, update) -> {
subLock.lock();
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);

2
application/src/test/java/org/thingsboard/server/controller/AbstractControllerTest.java
View File

@ -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) {

12
application/src/test/java/org/thingsboard/server/controller/BaseWebsocketApiTest.java
View File

@ -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();
}
}

21
application/src/test/java/org/thingsboard/server/controller/TbTestWebSocketClient.java
View File

@ -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);
}

Write Preview
Loading…
Cancel
Save