dendi-wang/thingsboard
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rate limit for WS subscriptions

Browse Source
This commit is contained in:
ViacheslavKlimov 2024-08-05 17:22:58 +03:00
parent fc36fcaacf
commit 5759610340
13 changed files with 123 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
application/src/main/java/org/thingsboard/server/service/subscription/DefaultTbEntityDataSubscriptionService.java
View File

@ -19,6 +19,8 @@ import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.MoreExecutors;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ArrayUtils;
@ -42,6 +44,7 @@ import org.thingsboard.server.common.data.query.EntityDataQuery;
import org.thingsboard.server.common.data.query.EntityKey;
import org.thingsboard.server.common.data.query.EntityKeyType;
import org.thingsboard.server.common.data.query.TsValue;
import org.thingsboard.server.common.msg.tools.TbRateLimitsException;
import org.thingsboard.server.dao.alarm.AlarmService;
import org.thingsboard.server.dao.attributes.AttributesService;
import org.thingsboard.server.dao.entity.EntityService;
@ -66,8 +69,6 @@ import org.thingsboard.server.service.ws.telemetry.cmd.v2.LatestValueCmd;
import org.thingsboard.server.service.ws.telemetry.cmd.v2.TimeSeriesCmd;
import org.thingsboard.server.service.ws.telemetry.cmd.v2.UnsubscribeCmd;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
@ -355,6 +356,9 @@ public class DefaultTbEntityDataSubscriptionService implements TbEntityDataSubsc
private void handleWsCmdRuntimeException(String sessionId, RuntimeException e, EntityDataCmd cmd) {
log.debug("[{}] Failed to process ws cmd: {}", sessionId, cmd, e);
if (e instanceof TbRateLimitsException) {
return;
}
wsService.close(sessionId, CloseStatus.SERVICE_RESTARTED);
}

41
application/src/main/java/org/thingsboard/server/service/subscription/DefaultTbLocalSubscriptionService.java
View File

@ -15,16 +15,20 @@
*/
package org.thingsboard.server.service.subscription;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Lazy;
import org.springframework.context.event.EventListener;
import org.springframework.stereotype.Service;
import org.thingsboard.common.util.DeduplicationUtil;
import org.thingsboard.common.util.DonAsynchron;
import org.thingsboard.common.util.ThingsBoardExecutors;
import org.thingsboard.common.util.ThingsBoardThreadFactory;
import org.thingsboard.server.cache.limits.RateLimitService;
import org.thingsboard.server.cluster.TbClusterService;
import org.thingsboard.server.common.data.AttributeScope;
import org.thingsboard.server.common.data.DataConstants;
import org.thingsboard.server.common.data.EntityType;
import org.thingsboard.server.common.data.alarm.AlarmInfo;
import org.thingsboard.server.common.data.id.EntityId;
@ -34,10 +38,12 @@ import org.thingsboard.server.common.data.kv.BaseReadTsKvQuery;
import org.thingsboard.server.common.data.kv.BasicTsKvEntry;
import org.thingsboard.server.common.data.kv.ReadTsKvQuery;
import org.thingsboard.server.common.data.kv.TsKvEntry;
import org.thingsboard.server.common.data.limit.LimitedApi;
import org.thingsboard.server.common.data.plugin.ComponentLifecycleEvent;
import org.thingsboard.server.common.msg.queue.ServiceType;
import org.thingsboard.server.common.msg.queue.TbCallback;
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
import org.thingsboard.server.common.msg.tools.TbRateLimitsException;
import org.thingsboard.server.dao.attributes.AttributesService;
import org.thingsboard.server.dao.timeseries.TimeseriesService;
import org.thingsboard.server.gen.transport.TransportProtos;
@ -46,13 +52,12 @@ import org.thingsboard.server.queue.discovery.TbServiceInfoProvider;
import org.thingsboard.server.queue.discovery.event.ClusterTopologyChangeEvent;
import org.thingsboard.server.queue.util.TbCoreComponent;
import org.thingsboard.server.service.ws.WebSocketService;
import org.thingsboard.server.service.ws.WebSocketSessionRef;
import org.thingsboard.server.service.ws.notification.sub.NotificationRequestUpdate;
import org.thingsboard.server.service.ws.notification.sub.NotificationsSubscriptionUpdate;
import org.thingsboard.server.service.ws.telemetry.sub.AlarmSubscriptionUpdate;
import org.thingsboard.server.service.ws.telemetry.sub.TelemetrySubscriptionUpdate;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
@ -88,13 +93,20 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
private final TbClusterService clusterService;
private final SubscriptionManagerService subscriptionManagerService;
private final WebSocketService webSocketService;
private final RateLimitService rateLimitService;
private ExecutorService tsCallBackExecutor;
private ScheduledExecutorService staleSessionCleanupExecutor;
@Value("${server.ws.rate_limits.subscriptions_per_tenant:2000:60}")
private String subscriptionsPerTenantRateLimit;
@Value("${server.ws.rate_limits.subscriptions_per_user:500:60}")
private String subscriptionsPerUserRateLimit;
public DefaultTbLocalSubscriptionService(AttributesService attrService, TimeseriesService tsService, TbServiceInfoProvider serviceInfoProvider,
PartitionService partitionService, TbClusterService clusterService,
@Lazy SubscriptionManagerService subscriptionManagerService, @Lazy WebSocketService webSocketService) {
@Lazy SubscriptionManagerService subscriptionManagerService, @Lazy WebSocketService webSocketService,
RateLimitService rateLimitService) {
this.attrService = attrService;
this.tsService = tsService;
this.serviceInfoProvider = serviceInfoProvider;
@ -102,6 +114,7 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
this.clusterService = clusterService;
this.subscriptionManagerService = subscriptionManagerService;
this.webSocketService = webSocketService;
this.rateLimitService = rateLimitService;
}
private String serviceId;
@ -164,9 +177,18 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
}
@Override
public void addSubscription(TbSubscription<?> subscription) {
public void addSubscription(TbSubscription<?> subscription, WebSocketSessionRef sessionRef) {
TenantId tenantId = subscription.getTenantId();
EntityId entityId = subscription.getEntityId();
if (!rateLimitService.checkRateLimit(LimitedApi.WS_SUBSCRIPTIONS, (Object) tenantId, subscriptionsPerTenantRateLimit)) {
handleRateLimitError(subscription, sessionRef, "Exceeded rate limit for WS subscriptions per tenant");
return;
}
if (sessionRef.getSecurityCtx() != null && !rateLimitService.checkRateLimit(LimitedApi.WS_SUBSCRIPTIONS, sessionRef.getSecurityCtx().getId(), subscriptionsPerUserRateLimit)) {
handleRateLimitError(subscription, sessionRef, "Exceeded rate limit for WS subscriptions per user");
return;
}
log.debug("[{}][{}] Register subscription: {}", tenantId, entityId, subscription);
SubscriptionModificationResult result;
subsLock.lock();
@ -563,4 +585,13 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
subscriptionsBySessionId.keySet().forEach(webSocketService::cleanupIfStale);
}
private void handleRateLimitError(TbSubscription<?> subscription, WebSocketSessionRef sessionRef, String message) {
String deduplicationKey = sessionRef.getSessionId() + message;
if (!DeduplicationUtil.alreadyProcessed(deduplicationKey, TimeUnit.SECONDS.toMillis(15))) {
log.info("{} {}", sessionRef, message);
webSocketService.sendError(sessionRef, subscription.getSubscriptionId(), SubscriptionErrorCode.BAD_REQUEST, message);
}
throw new TbRateLimitsException(message);
}
}

5
application/src/main/java/org/thingsboard/server/service/subscription/TbAbstractDataSubCtx.java
View File

@ -132,7 +132,7 @@ public abstract class TbAbstractDataSubCtx<T extends AbstractDataQuery<? extends
int subIdx = sessionRef.getSessionSubIdSeq().incrementAndGet();
subToEntityIdMap.put(subIdx, entityData.getEntityId());
localSubscriptionService.addSubscription(
createTsSub(entityData, subIdx, false, startTs, endTs, keyStates, resultToLatestValues));
createTsSub(entityData, subIdx, false, startTs, endTs, keyStates, resultToLatestValues), sessionRef);
});
}
@ -140,7 +140,7 @@ public abstract class TbAbstractDataSubCtx<T extends AbstractDataQuery<? extends
Map<EntityKeyType, List<EntityKey>> keysByType = getEntityKeyByTypeMap(keys);
for (EntityData entityData : data.getData()) {
List<TbSubscription> entitySubscriptions = addSubscriptions(entityData, keysByType, latestValues, startTs, endTs);
entitySubscriptions.forEach(localSubscriptionService::addSubscription);
entitySubscriptions.forEach(subscription -> localSubscriptionService.addSubscription(subscription, sessionRef));
}
}
@ -254,4 +254,5 @@ public abstract class TbAbstractDataSubCtx<T extends AbstractDataQuery<? extends
abstract void sendWsMsg(String sessionId, TelemetrySubscriptionUpdate subscriptionUpdate, EntityKeyType keyType, boolean resultToLatestValues);
protected abstract Aggregation getCurrentAggregation();
}

2
application/src/main/java/org/thingsboard/server/service/subscription/TbAbstractSubCtx.java
View File

@ -152,7 +152,7 @@ public abstract class TbAbstractSubCtx<T extends EntityCountQuery> {
.scope(TbAttributeSubscriptionScope.SERVER_SCOPE)
.build();
subToDynamicValueKeySet.add(subIdx);
localSubscriptionService.addSubscription(sub);
localSubscriptionService.addSubscription(sub, sessionRef);
}
} catch (InterruptedException | ExecutionException e) {
log.info("[{}][{}][{}] Failed to resolve dynamic values: {}", tenantId, customerId, userId, dynamicValues.keySet());

5
application/src/main/java/org/thingsboard/server/service/subscription/TbAlarmDataSubCtx.java
View File

@ -177,7 +177,7 @@ public class TbAlarmDataSubCtx extends TbAbstractDataSubCtx<AlarmDataQuery> {
.updateProcessor((sub, update) -> sendWsMsg(sub.getSessionId(), update))
.ts(startTs)
.build();
localSubscriptionService.addSubscription(subscription);
localSubscriptionService.addSubscription(subscription, sessionRef);
}
@Override
@ -342,7 +342,7 @@ public class TbAlarmDataSubCtx extends TbAbstractDataSubCtx<AlarmDataQuery> {
newSubsList.forEach(entity -> createAlarmSubscriptionForEntity(query.getPageLink(), startTs, entity));
}
subIdsToCancel.forEach(subId -> localSubscriptionService.cancelSubscription(getSessionId(), subId));
subsToAdd.forEach(localSubscriptionService::addSubscription);
subsToAdd.forEach(subscription -> localSubscriptionService.addSubscription(subscription, sessionRef));
}
private void resetInvocationCounter() {
@ -361,4 +361,5 @@ public class TbAlarmDataSubCtx extends TbAbstractDataSubCtx<AlarmDataQuery> {
EntityDataPageLink edpl = new EntityDataPageLink(maxEntitiesPerAlarmSubscription, 0, null, entitiesSortOrder);
return new EntityDataQuery(query.getEntityFilter(), edpl, query.getEntityFields(), query.getLatestValues(), query.getKeyFilters());
}
}

3
application/src/main/java/org/thingsboard/server/service/subscription/TbEntityDataSubCtx.java
View File

@ -226,7 +226,7 @@ public class TbEntityDataSubCtx extends TbAbstractDataSubCtx<EntityDataQuery> {
}
}
subIdsToCancel.forEach(subId -> localSubscriptionService.cancelSubscription(getSessionId(), subId));
subsToAdd.forEach(localSubscriptionService::addSubscription);
subsToAdd.forEach(subscription -> localSubscriptionService.addSubscription(subscription, sessionRef));
sendWsMsg(new EntityDataUpdate(cmdId, data, null, maxEntitiesPerDataSubscription));
}
@ -239,4 +239,5 @@ public class TbEntityDataSubCtx extends TbAbstractDataSubCtx<EntityDataQuery> {
protected EntityDataQuery buildEntityDataQuery() {
return query;
}
}

3
application/src/main/java/org/thingsboard/server/service/subscription/TbLocalSubscriptionService.java
View File

@ -22,6 +22,7 @@ import org.thingsboard.server.common.data.kv.TsKvEntry;
import org.thingsboard.server.common.msg.queue.TbCallback;
import org.thingsboard.server.gen.transport.TransportProtos;
import org.thingsboard.server.queue.discovery.event.ClusterTopologyChangeEvent;
import org.thingsboard.server.service.ws.WebSocketSessionRef;
import org.thingsboard.server.service.ws.notification.sub.NotificationRequestUpdate;
import org.thingsboard.server.service.ws.notification.sub.NotificationsSubscriptionUpdate;
@ -29,7 +30,7 @@ import java.util.List;
public interface TbLocalSubscriptionService {
void addSubscription(TbSubscription<?> subscription);
void addSubscription(TbSubscription<?> subscription, WebSocketSessionRef sessionRef);
void onSubEventCallback(TransportProtos.TbEntitySubEventCallbackProto subEventCallback, TbCallback callback);

21
application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java
View File

@ -21,6 +21,9 @@ import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.MoreExecutors;
import jakarta.annotation.Nullable;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
@ -46,6 +49,7 @@ import org.thingsboard.server.common.data.kv.BasicTsKvEntry;
import org.thingsboard.server.common.data.kv.ReadTsKvQuery;
import org.thingsboard.server.common.data.kv.TsKvEntry;
import org.thingsboard.server.common.data.tenant.profile.DefaultTenantProfileConfiguration;
import org.thingsboard.server.common.msg.tools.TbRateLimitsException;
import org.thingsboard.server.dao.attributes.AttributesService;
import org.thingsboard.server.dao.tenant.TbTenantProfileCache;
import org.thingsboard.server.dao.timeseries.TimeseriesService;
@ -80,10 +84,6 @@ import org.thingsboard.server.service.ws.telemetry.cmd.v2.EntityDataUpdate;
import org.thingsboard.server.service.ws.telemetry.cmd.v2.UnsubscribeCmd;
import org.thingsboard.server.service.ws.telemetry.sub.TelemetrySubscriptionUpdate;
import jakarta.annotation.Nullable;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
@ -224,9 +224,10 @@ public class DefaultWebSocketService implements WebSocketService {
try {
Optional.ofNullable(cmdsHandlers.get(cmd.getType()))
.ifPresent(cmdHandler -> cmdHandler.handle(sessionRef, cmd));
} catch (TbRateLimitsException e) {
log.debug("{} Failed to handle WS cmd: {}", sessionRef, cmd, e);
} catch (Exception e) {
log.error("[sessionId: {}, tenantId: {}, userId: {}] Failed to handle WS cmd: {}", sessionId,
sessionRef.getSecurityCtx().getTenantId(), sessionRef.getSecurityCtx().getId(), cmd, e);
log.error("{} Failed to handle WS cmd: {}", sessionRef, cmd, e);
}
}
}
@ -468,7 +469,7 @@ public class DefaultWebSocketService implements WebSocketService {
subLock.lock();
try {
oldSubService.addSubscription(sub);
oldSubService.addSubscription(sub, sessionRef);
sendUpdate(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), attributesData));
} finally {
subLock.unlock();
@ -581,7 +582,7 @@ public class DefaultWebSocketService implements WebSocketService {
subLock.lock();
try {
oldSubService.addSubscription(sub);
oldSubService.addSubscription(sub, sessionRef);
sendUpdate(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), attributesData));
} finally {
subLock.unlock();
@ -678,7 +679,7 @@ public class DefaultWebSocketService implements WebSocketService {
subLock.lock();
try {
oldSubService.addSubscription(sub);
oldSubService.addSubscription(sub, sessionRef);
sendUpdate(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), data));
} finally {
subLock.unlock();
@ -733,7 +734,7 @@ public class DefaultWebSocketService implements WebSocketService {
subLock.lock();
try {
oldSubService.addSubscription(sub);
oldSubService.addSubscription(sub, sessionRef);
sendUpdate(sessionRef, new TelemetrySubscriptionUpdate(cmd.getCmdId(), data));
} finally {
subLock.unlock();

4
application/src/main/java/org/thingsboard/server/service/ws/notification/DefaultNotificationCommandsHandler.java
View File

@ -79,7 +79,7 @@ public class DefaultNotificationCommandsHandler implements NotificationCommandsH
.updateProcessor(this::handleNotificationsSubscriptionUpdate)
.limit(cmd.getLimit())
.build();
localSubscriptionService.addSubscription(subscription);
localSubscriptionService.addSubscription(subscription, sessionRef);
fetchUnreadNotifications(subscription);
sendUpdate(sessionRef.getSessionId(), subscription.createFullUpdate());
@ -97,7 +97,7 @@ public class DefaultNotificationCommandsHandler implements NotificationCommandsH
.entityId(securityCtx.getId())
.updateProcessor(this::handleNotificationsCountSubscriptionUpdate)
.build();
localSubscriptionService.addSubscription(subscription);
localSubscriptionService.addSubscription(subscription, sessionRef);
fetchUnreadNotificationsCount(subscription);
sendUpdate(sessionRef.getSessionId(), subscription.createUpdate());

5
application/src/main/resources/thingsboard.yml
View File

@ -78,6 +78,11 @@ server:
max_queue_messages_per_session: "${TB_SERVER_WS_DEFAULT_QUEUE_MESSAGES_PER_SESSION:1000}"
# Maximum time between WS session opening and sending auth command
auth_timeout_ms: "${TB_SERVER_WS_AUTH_TIMEOUT_MS:10000}"
rate_limits:
# Per-tenant rate limit for WS subscriptions
subscriptions_per_tenant: "${TB_SERVER_WS_SUBSCRIPTIONS_PER_TENANT_RATE_LIMIT:2000:60}"
# Per-user rate limit for WS subscriptions
subscriptions_per_user: "${TB_SERVER_WS_SUBSCRIPTIONS_PER_USER_RATE_LIMIT:500:60}"
rest:
server_side_rpc:
# Minimum value of the server-side RPC timeout. May override value provided in the REST API call.

3
common/data/src/main/java/org/thingsboard/server/common/data/limit/LimitedApi.java
View File

@ -42,7 +42,8 @@ public enum LimitedApi {
TRANSPORT_MESSAGES_PER_DEVICE("transport messages per device", false),
TRANSPORT_MESSAGES_PER_GATEWAY("transport messages per gateway", false),
TRANSPORT_MESSAGES_PER_GATEWAY_DEVICE("transport messages per gateway device", false),
EMAILS("emails sending", true);
EMAILS("emails sending", true),
WS_SUBSCRIPTIONS("WS subscriptions", false);
private Function<DefaultTenantProfileConfiguration, String> configExtractor;
@Getter

6
common/message/src/main/java/org/thingsboard/server/common/msg/tools/TbRateLimitsException.java
View File

@ -30,4 +30,10 @@ public class TbRateLimitsException extends AbstractRateLimitException {
super(entityType.name() + " rate limits reached!");
this.entityType = entityType;
}
public TbRateLimitsException(String message) {
super(message);
this.entityType = null;
}
}

44
common/util/src/main/java/org/thingsboard/common/util/DeduplicationUtil.java Normal file
View File

@ -0,0 +1,44 @@
/**
* Copyright © 2016-2024 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.common.util;
import org.springframework.util.ConcurrentReferenceHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicBoolean;
import static org.springframework.util.ConcurrentReferenceHashMap.ReferenceType.SOFT;
public class DeduplicationUtil {
private static final ConcurrentMap<Object, Long> cache = new ConcurrentReferenceHashMap<>(16, SOFT);
public static boolean alreadyProcessed(Object deduplicationKey, long deduplicationDuration) {
AtomicBoolean alreadyProcessed = new AtomicBoolean(false);
cache.compute(deduplicationKey, (key, lastProcessedTs) -> {
if (lastProcessedTs != null) {
long passed = System.currentTimeMillis() - lastProcessedTs;
if (passed <= deduplicationDuration) {
alreadyProcessed.set(true);
return lastProcessedTs;
}
}
return System.currentTimeMillis();
});
return alreadyProcessed.get();
}
}