Rate limit improvements

This commit is contained in:
Andrii Shvaika 2020-10-19 11:50:06 +03:00
parent f181beec61
commit e9bf5bae29
12 changed files with 136 additions and 60 deletions

View File

@ -50,4 +50,8 @@ public class TbRateLimits {
return bucket.tryConsume(1);
}
public boolean tryConsume(long number) {
return bucket.tryConsume(number);
}
}

View File

@ -0,0 +1,26 @@
/**
* Copyright © 2016-2020 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.util;
import org.springframework.boot.autoconfigure.condition.ConditionalOnExpression;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
@Retention(RetentionPolicy.RUNTIME)
@ConditionalOnExpression("('${service.type:null}'=='monolith' && '${transport.api_enabled:true}'=='true') || '${service.type:null}'=='tb-transport'")
public @interface TbTransportComponent {
}

View File

@ -20,6 +20,7 @@ import org.thingsboard.server.common.data.DeviceTransportType;
import org.thingsboard.server.common.data.id.DeviceProfileId;
import org.thingsboard.server.common.transport.auth.GetOrCreateDeviceFromGatewayResponse;
import org.thingsboard.server.common.transport.auth.ValidateDeviceCredentialsResponse;
import org.thingsboard.server.common.transport.limits.TransportRateLimitType;
import org.thingsboard.server.gen.transport.TransportProtos;
import org.thingsboard.server.gen.transport.TransportProtos.ClaimDeviceMsg;
import org.thingsboard.server.gen.transport.TransportProtos.GetAttributeRequestMsg;
@ -69,6 +70,8 @@ public interface TransportService {
boolean checkLimits(SessionInfoProto sessionInfo, Object msg, TransportServiceCallback<Void> callback);
boolean checkLimits(SessionInfoProto sessionInfo, Object msg, TransportServiceCallback<Void> callback, int dataPoints, TransportRateLimitType... limits);
void process(SessionInfoProto sessionInfo, SessionEventMsg msg, TransportServiceCallback<Void> callback);
void process(SessionInfoProto sessionInfo, PostTelemetryMsg msg, TransportServiceCallback<Void> callback);

View File

@ -16,7 +16,6 @@
package org.thingsboard.server.common.transport.limits;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnExpression;
import org.springframework.stereotype.Service;
import org.thingsboard.server.common.data.TenantProfile;
import org.thingsboard.server.common.data.TenantProfileData;
@ -24,12 +23,13 @@ import org.thingsboard.server.common.data.id.DeviceId;
import org.thingsboard.server.common.data.id.TenantId;
import org.thingsboard.server.common.transport.TransportTenantProfileCache;
import org.thingsboard.server.common.transport.profile.TenantProfileUpdateResult;
import org.thingsboard.server.queue.util.TbTransportComponent;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
@Service
@ConditionalOnExpression("('${service.type:null}'=='monolith' && '${transport.api_enabled:true}'=='true') || '${service.type:null}'=='tb-transport'")
@TbTransportComponent
@Slf4j
public class DefaultTransportRateLimitService implements TransportRateLimitService {
@ -45,23 +45,21 @@ public class DefaultTransportRateLimitService implements TransportRateLimitServi
}
@Override
public TransportRateLimit getRateLimit(TenantId tenantId, TransportRateLimitType limitType) {
TransportRateLimit[] limits = perTenantLimits.get(tenantId);
if (limits == null) {
limits = fetchProfileAndInit(tenantId);
perTenantLimits.put(tenantId, limits);
public TransportRateLimitType checkLimits(TenantId tenantId, DeviceId deviceId, int dataPoints, TransportRateLimitType... limits) {
TransportRateLimit[] tenantLimits = getTenantRateLimits(tenantId);
TransportRateLimit[] deviceLimits = getDeviceRateLimits(tenantId, deviceId);
for (TransportRateLimitType limitType : limits) {
TransportRateLimit rateLimit;
if (limitType.isTenantLevel()) {
rateLimit = tenantLimits[limitType.ordinal()];
} else {
rateLimit = deviceLimits[limitType.ordinal()];
}
return limits[limitType.ordinal()];
if (!rateLimit.tryConsume(limitType.isMessageLevel() ? 1L : dataPoints)) {
return limitType;
}
@Override
public TransportRateLimit getRateLimit(TenantId tenantId, DeviceId deviceId, TransportRateLimitType limitType) {
TransportRateLimit[] limits = perDeviceLimits.get(deviceId);
if (limits == null) {
limits = fetchProfileAndInit(tenantId);
perDeviceLimits.put(deviceId, limits);
}
return limits[limitType.ordinal()];
return null;
}
@Override
@ -77,7 +75,17 @@ public class DefaultTransportRateLimitService implements TransportRateLimitServi
mergeLimits(tenantId, fetchProfileAndInit(tenantId));
}
public void mergeLimits(TenantId tenantId, TransportRateLimit[] newRateLimits) {
@Override
public void remove(TenantId tenantId) {
perTenantLimits.remove(tenantId);
}
@Override
public void remove(DeviceId deviceId) {
perDeviceLimits.remove(deviceId);
}
private void mergeLimits(TenantId tenantId, TransportRateLimit[] newRateLimits) {
TransportRateLimit[] oldRateLimits = perTenantLimits.get(tenantId);
if (oldRateLimits == null) {
perTenantLimits.put(tenantId, newRateLimits);
@ -92,16 +100,6 @@ public class DefaultTransportRateLimitService implements TransportRateLimitServi
}
}
@Override
public void remove(TenantId tenantId) {
perTenantLimits.remove(tenantId);
}
@Override
public void remove(DeviceId deviceId) {
perDeviceLimits.remove(deviceId);
}
private TransportRateLimit[] fetchProfileAndInit(TenantId tenantId) {
return perTenantLimits.computeIfAbsent(tenantId, tmp -> createTransportRateLimits(tenantProfileCache.get(tenantId)));
}
@ -114,4 +112,22 @@ public class DefaultTransportRateLimitService implements TransportRateLimitServi
}
return rateLimits;
}
private TransportRateLimit[] getTenantRateLimits(TenantId tenantId) {
TransportRateLimit[] limits = perTenantLimits.get(tenantId);
if (limits == null) {
limits = fetchProfileAndInit(tenantId);
perTenantLimits.put(tenantId, limits);
}
return limits;
}
private TransportRateLimit[] getDeviceRateLimits(TenantId tenantId, DeviceId deviceId) {
TransportRateLimit[] limits = perDeviceLimits.get(deviceId);
if (limits == null) {
limits = fetchProfileAndInit(tenantId);
perDeviceLimits.put(deviceId, limits);
}
return limits;
}
}

View File

@ -22,6 +22,11 @@ public class DummyTransportRateLimit implements TransportRateLimit {
return "";
}
@Override
public boolean tryConsume(long number) {
return true;
}
@Override
public boolean tryConsume() {
return true;

View File

@ -31,4 +31,8 @@ public class SimpleTransportRateLimit implements TransportRateLimit {
return rateLimit.tryConsume();
}
@Override
public boolean tryConsume(long number) {
return number <= 0 || rateLimit.tryConsume(number);
}
}

View File

@ -21,4 +21,6 @@ public interface TransportRateLimit {
boolean tryConsume();
boolean tryConsume(long number);
}

View File

@ -21,9 +21,7 @@ import org.thingsboard.server.common.transport.profile.TenantProfileUpdateResult
public interface TransportRateLimitService {
TransportRateLimit getRateLimit(TenantId tenantId, TransportRateLimitType limit);
TransportRateLimit getRateLimit(TenantId tenantId, DeviceId deviceId, TransportRateLimitType limit);
TransportRateLimitType checkLimits(TenantId tenantId, DeviceId deviceId, int dataPoints, TransportRateLimitType... limits);
void update(TenantProfileUpdateResult update);

View File

@ -19,15 +19,29 @@ import lombok.Getter;
public enum TransportRateLimitType {
TENANT_MAX_MSGS("transport.tenant.max.msg"),
TENANT_MAX_DATA_POINTS("transport.tenant.max.dataPoints"),
DEVICE_MAX_MSGS("transport.device.max.msg"),
DEVICE_MAX_DATA_POINTS("transport.device.max.dataPoints");
TENANT_MAX_MSGS("transport.tenant.msg", true, true),
TENANT_TELEMETRY_MSGS("transport.tenant.telemetry", true, true),
TENANT_MAX_DATA_POINTS("transport.tenant.dataPoints", true, false),
DEVICE_MAX_MSGS("transport.device.msg", false, true),
DEVICE_TELEMETRY_MSGS("transport.device.telemetry", false, true),
DEVICE_MAX_DATA_POINTS("transport.device.dataPoints", false, false);
@Getter
private final String configurationKey;
@Getter
private final boolean tenantLevel;
@Getter
private final boolean deviceLevel;
@Getter
private final boolean messageLevel;
@Getter
private final boolean dataPointLevel;
TransportRateLimitType(String configurationKey) {
TransportRateLimitType(String configurationKey, boolean tenantLevel, boolean messageLevel) {
this.configurationKey = configurationKey;
this.tenantLevel = tenantLevel;
this.deviceLevel = !tenantLevel;
this.messageLevel = messageLevel;
this.dataPointLevel = !messageLevel;
}
}

View File

@ -23,6 +23,7 @@ import org.thingsboard.server.common.data.DeviceProfile;
import org.thingsboard.server.common.data.id.DeviceProfileId;
import org.thingsboard.server.common.transport.TransportDeviceProfileCache;
import org.thingsboard.server.common.transport.util.DataDecodingEncodingService;
import org.thingsboard.server.queue.util.TbTransportComponent;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
@ -30,7 +31,7 @@ import java.util.concurrent.ConcurrentMap;
@Slf4j
@Component
@ConditionalOnExpression("('${service.type:null}'=='monolith' && '${transport.api_enabled:true}'=='true') || '${service.type:null}'=='tb-transport'")
@TbTransportComponent
public class DefaultTransportDeviceProfileCache implements TransportDeviceProfileCache {
private final ConcurrentMap<DeviceProfileId, DeviceProfile> deviceProfiles = new ConcurrentHashMap<>();

View File

@ -79,6 +79,7 @@ import org.thingsboard.server.queue.discovery.PartitionService;
import org.thingsboard.server.queue.discovery.TbServiceInfoProvider;
import org.thingsboard.server.queue.provider.TbQueueProducerProvider;
import org.thingsboard.server.queue.provider.TbTransportQueueFactory;
import org.thingsboard.server.queue.util.TbTransportComponent;
import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
@ -103,7 +104,7 @@ import java.util.concurrent.atomic.AtomicInteger;
*/
@Slf4j
@Service
@ConditionalOnExpression("('${service.type:null}'=='monolith' && '${transport.api_enabled:true}'=='true') || '${service.type:null}'=='tb-transport'")
@TbTransportComponent
public class DefaultTransportService implements TransportService {
@Value("${transport.sessions.inactivity_timeout}")
@ -363,7 +364,11 @@ public class DefaultTransportService implements TransportService {
@Override
public void process(TransportProtos.SessionInfoProto sessionInfo, TransportProtos.PostTelemetryMsg msg, TransportServiceCallback<Void> callback) {
if (checkLimits(sessionInfo, msg, callback)) {
int dataPoints = 0;
for (TransportProtos.TsKvListProto tsKv : msg.getTsKvListList()) {
dataPoints += tsKv.getKvCount();
}
if (checkLimits(sessionInfo, msg, callback, dataPoints, TELEMETRY)) {
reportActivityInternal(sessionInfo);
TenantId tenantId = new TenantId(new UUID(sessionInfo.getTenantIdMSB(), sessionInfo.getTenantIdLSB()));
DeviceId deviceId = new DeviceId(new UUID(sessionInfo.getDeviceIdMSB(), sessionInfo.getDeviceIdLSB()));
@ -384,7 +389,7 @@ public class DefaultTransportService implements TransportService {
@Override
public void process(TransportProtos.SessionInfoProto sessionInfo, TransportProtos.PostAttributeMsg msg, TransportServiceCallback<Void> callback) {
if (checkLimits(sessionInfo, msg, callback)) {
if (checkLimits(sessionInfo, msg, callback, msg.getKvCount(), TELEMETRY)) {
reportActivityInternal(sessionInfo);
TenantId tenantId = new TenantId(new UUID(sessionInfo.getTenantIdMSB(), sessionInfo.getTenantIdLSB()));
DeviceId deviceId = new DeviceId(new UUID(sessionInfo.getDeviceIdMSB(), sessionInfo.getDeviceIdLSB()));
@ -574,37 +579,34 @@ public class DefaultTransportService implements TransportService {
sessions.remove(toSessionId(sessionInfo));
}
private TransportRateLimitType[] DEFAULT = new TransportRateLimitType[]{TransportRateLimitType.TENANT_MAX_MSGS, TransportRateLimitType.DEVICE_MAX_MSGS};
private TransportRateLimitType[] TELEMETRY = TransportRateLimitType.values();
@Override
public boolean checkLimits(TransportProtos.SessionInfoProto sessionInfo, Object msg, TransportServiceCallback<Void> callback) {
return checkLimits(sessionInfo, msg, callback, 0, DEFAULT);
}
@Override
public boolean checkLimits(TransportProtos.SessionInfoProto sessionInfo, Object msg, TransportServiceCallback<Void> callback, int dataPoints, TransportRateLimitType... limits) {
if (log.isTraceEnabled()) {
log.trace("[{}] Processing msg: {}", toSessionId(sessionInfo), msg);
}
TenantId tenantId = new TenantId(new UUID(sessionInfo.getTenantIdMSB(), sessionInfo.getTenantIdLSB()));
TransportRateLimit tenantRateLimit = rateLimitService.getRateLimit(tenantId, TransportRateLimitType.TENANT_MAX_MSGS);
if (!tenantRateLimit.tryConsume()) {
if (callback != null) {
callback.onError(new TbRateLimitsException(EntityType.TENANT));
}
if (log.isTraceEnabled()) {
log.trace("[{}][{}] Tenant level rate limit detected: {}", toSessionId(sessionInfo), tenantId, msg);
}
return false;
}
DeviceId deviceId = new DeviceId(new UUID(sessionInfo.getDeviceIdMSB(), sessionInfo.getDeviceIdLSB()));
TransportRateLimit deviceRateLimit = rateLimitService.getRateLimit(tenantId, deviceId, TransportRateLimitType.DEVICE_MAX_MSGS);
if (!deviceRateLimit.tryConsume()) {
TransportRateLimitType limit = rateLimitService.checkLimits(tenantId, deviceId, 0, limits);
if (limit == null) {
return true;
} else {
if (callback != null) {
callback.onError(new TbRateLimitsException(EntityType.DEVICE));
callback.onError(new TbRateLimitsException(limit.isTenantLevel() ? EntityType.TENANT : EntityType.DEVICE));
}
if (log.isTraceEnabled()) {
log.trace("[{}][{}] Device level rate limit detected: {}", toSessionId(sessionInfo), deviceId, msg);
log.trace("[{}][{}] {} rateLimit detected: {}", toSessionId(sessionInfo), tenantId, limit, msg);
}
return false;
}
return true;
}
protected void processToTransportMsg(TransportProtos.ToTransportMsg toSessionMsg) {

View File

@ -33,6 +33,7 @@ import org.thingsboard.server.common.transport.util.DataDecodingEncodingService;
import org.thingsboard.server.gen.transport.TransportProtos;
import org.thingsboard.server.queue.discovery.TenantRoutingInfo;
import org.thingsboard.server.queue.discovery.TenantRoutingInfoService;
import org.thingsboard.server.queue.util.TbTransportComponent;
import java.util.Collections;
import java.util.Optional;
@ -43,7 +44,7 @@ import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
@Component
@ConditionalOnExpression("('${service.type:null}'=='monolith' && '${transport.api_enabled:true}'=='true') || '${service.type:null}'=='tb-transport'")
@TbTransportComponent
@Slf4j
public class DefaultTransportTenantProfileCache implements TransportTenantProfileCache {