Merge pull request #13511 from thingsboard/task-improvements

Proper rate limit exception for Cassandra queries

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

@@ -36,6 +36,7 @@ import org.thingsboard.common.util.ThingsBoardExecutors;
 import org.thingsboard.server.common.data.AttributeScope;
 import org.thingsboard.server.common.data.StringUtils;
 import org.thingsboard.server.common.data.TenantProfile;
+import org.thingsboard.server.common.data.exception.RateLimitExceededException;
 import org.thingsboard.server.common.data.id.CustomerId;
 import org.thingsboard.server.common.data.id.EntityId;
 import org.thingsboard.server.common.data.id.EntityIdFactory;
@@ -52,7 +53,6 @@ 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;
-import org.thingsboard.server.dao.util.TenantRateLimitException;
 import org.thingsboard.server.exception.UnauthorizedException;
 import org.thingsboard.server.queue.discovery.TbServiceInfoProvider;
 import org.thingsboard.server.queue.util.TbCoreComponent;
@@ -742,7 +742,7 @@ public class DefaultWebSocketService implements WebSocketService {
 
             @Override
             public void onFailure(Throwable e) {
-                if (e instanceof TenantRateLimitException || e.getCause() instanceof TenantRateLimitException) {
+                if (e instanceof RateLimitExceededException || e.getCause() instanceof RateLimitExceededException) {
                     log.trace("[{}] Tenant rate limit detected for subscription: [{}]:{}", sessionRef.getSecurityCtx().getTenantId(), entityId, cmd);
                 } else {
                     log.info(FAILED_TO_FETCH_DATA, e);

dao/src/main/java/org/thingsboard/server/dao/timeseries/BaseTimeseriesService.java

@@ -205,8 +205,8 @@ public class BaseTimeseriesService implements TimeseriesService {
         ListenableFuture<Integer> dpsFuture = saveTs ? Futures.transform(Futures.allAsList(tsFutures), SUM_ALL_INTEGERS, MoreExecutors.directExecutor()) : Futures.immediateFuture(0);
         ListenableFuture<List<Long>> versionsFuture = saveLatest ? Futures.allAsList(latestFutures) : Futures.immediateFuture(null);
         return Futures.whenAllComplete(dpsFuture, versionsFuture).call(() -> {
-            Integer dataPoints = Futures.getUnchecked(dpsFuture);
+            Integer dataPoints = dpsFuture.get();
-            List<Long> versions = Futures.getUnchecked(versionsFuture);
+            List<Long> versions = versionsFuture.get();
             return TimeseriesSaveResult.of(dataPoints, versions);
         }, MoreExecutors.directExecutor());
     }

dao/src/main/java/org/thingsboard/server/dao/util/AbstractBufferedRateExecutor.java

@@ -32,6 +32,7 @@ import lombok.extern.slf4j.Slf4j;
 import org.thingsboard.common.util.ThingsBoardExecutors;
 import org.thingsboard.common.util.ThingsBoardThreadFactory;
 import org.thingsboard.server.cache.limits.RateLimitService;
+import org.thingsboard.server.common.data.exception.RateLimitExceededException;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.limit.LimitedApi;
 import org.thingsboard.server.common.msg.queue.ServiceType;
@@ -124,7 +125,7 @@ public abstract class AbstractBufferedRateExecutor<T extends AsyncTask, F extend
             if (!rateLimitService.checkRateLimit(myLimitedApi, tenantId, tenantId, true)) {
                 stats.incrementRateLimitedTenant(tenantId);
                 stats.getTotalRateLimited().increment();
-                settableFuture.setException(new TenantRateLimitException());
+                settableFuture.setException(new RateLimitExceededException(myLimitedApi));
                 perTenantLimitReached = true;
             }
         } else if (tenantId == null) {

dao/src/main/java/org/thingsboard/server/dao/util/TenantRateLimitException.java

@@ -1,19 +0,0 @@
-/**
- * Copyright © 2016-2025 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.dao.util;
-
-public class TenantRateLimitException extends Exception {
-}