Provide usage stats for sysTenant; refactor

application/src/main/java/org/thingsboard/server/service/apiusage/BaseApiUsageState.java

@@ -1,3 +1,18 @@
+/**
+ * Copyright © 2016-2021 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.service.apiusage;
 
 import lombok.Getter;
@@ -121,23 +136,6 @@ public abstract class BaseApiUsageState {
         return !currentValue.equals(value);
     }
 
-    public Map<ApiFeature, ApiUsageStateValue> checkStateUpdatedDueToThresholds() {
-        return checkStateUpdatedDueToThreshold(new HashSet<>(Arrays.asList(ApiFeature.values())));
-    }
-
-    public Map<ApiFeature, ApiUsageStateValue> checkStateUpdatedDueToThreshold(Set<ApiFeature> features) {
-        Map<ApiFeature, ApiUsageStateValue> result = new HashMap<>();
-        for (ApiFeature feature : features) {
-            Pair<ApiFeature, ApiUsageStateValue> tmp = checkStateUpdatedDueToThreshold(feature);
-            if (tmp != null) {
-                result.put(tmp.getFirst(), tmp.getSecond());
-            }
-        }
-        return result;
-    }
-
-    public abstract Pair<ApiFeature, ApiUsageStateValue> checkStateUpdatedDueToThreshold(ApiFeature feature);
-
     public abstract EntityType getEntityType();
 
     public TenantId getTenantId() {

application/src/main/java/org/thingsboard/server/service/apiusage/CustomerApiUsageState.java

@@ -1,9 +1,21 @@
+/**
+ * Copyright © 2016-2021 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.service.apiusage;
 
-import org.springframework.data.util.Pair;
-import org.thingsboard.server.common.data.ApiFeature;
 import org.thingsboard.server.common.data.ApiUsageState;
-import org.thingsboard.server.common.data.ApiUsageStateValue;
 import org.thingsboard.server.common.data.EntityType;
 
 public class CustomerApiUsageState extends BaseApiUsageState {
@@ -11,12 +23,6 @@ public class CustomerApiUsageState extends BaseApiUsageState {
         super(apiUsageState);
     }
 
-    @Override
-    public Pair<ApiFeature, ApiUsageStateValue> checkStateUpdatedDueToThreshold(ApiFeature feature) {
-        ApiUsageStateValue featureValue = ApiUsageStateValue.ENABLED;
-        return setFeatureValue(feature, featureValue) ? Pair.of(feature, featureValue) : null;
-    }
-
     @Override
     public EntityType getEntityType() {
         return EntityType.CUSTOMER;

application/src/main/java/org/thingsboard/server/service/apiusage/DefaultTbApiUsageStateService.java

@@ -48,6 +48,7 @@ import org.thingsboard.server.common.data.tenant.profile.TenantProfileData;
 import org.thingsboard.server.common.msg.queue.ServiceType;
 import org.thingsboard.server.common.msg.queue.TbCallback;
 import org.thingsboard.server.common.msg.tools.SchedulerUtils;
+import org.thingsboard.server.dao.customer.CustomerService;
 import org.thingsboard.server.dao.tenant.TbTenantProfileCache;
 import org.thingsboard.server.dao.tenant.TenantService;
 import org.thingsboard.server.dao.timeseries.TimeseriesService;
@@ -66,10 +67,10 @@ import javax.annotation.PostConstruct;
 import javax.annotation.PreDestroy;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
@@ -99,6 +100,7 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
     private final TbClusterService clusterService;
     private final PartitionService partitionService;
     private final TenantService tenantService;
+    private final CustomerService customerService;
     private final TimeseriesService tsService;
     private final ApiUsageStateService apiUsageStateService;
     private final SchedulerComponent scheduler;
@@ -127,13 +129,16 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
     public DefaultTbApiUsageStateService(TbClusterService clusterService,
                                          PartitionService partitionService,
                                          TenantService tenantService,
+                                         CustomerService customerService,
                                          TimeseriesService tsService,
                                          ApiUsageStateService apiUsageStateService,
                                          SchedulerComponent scheduler,
-                                         TbTenantProfileCache tenantProfileCache, MailService mailService) {
+                                         TbTenantProfileCache tenantProfileCache,
+                                         MailService mailService) {
         this.clusterService = clusterService;
         this.partitionService = partitionService;
         this.tenantService = tenantService;
+        this.customerService = customerService;
         this.tsService = tsService;
         this.apiUsageStateService = apiUsageStateService;
         this.scheduler = scheduler;
@@ -154,21 +159,20 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
     @Override
     public void process(TbProtoQueueMsg<ToUsageStatsServiceMsg> msg, TbCallback callback) {
         ToUsageStatsServiceMsg statsMsg = msg.getValue();
+
         TenantId tenantId = new TenantId(new UUID(statsMsg.getTenantIdMSB(), statsMsg.getTenantIdLSB()));
-        CustomerId customerId;
+        EntityId initiatorId;
         if (statsMsg.getCustomerIdMSB() != 0 && statsMsg.getCustomerIdLSB() != 0) {
-            customerId = new CustomerId(new UUID(statsMsg.getCustomerIdMSB(), statsMsg.getCustomerIdLSB()));
+            initiatorId = new CustomerId(new UUID(statsMsg.getCustomerIdMSB(), statsMsg.getCustomerIdLSB()));
         } else {
-            customerId = new CustomerId(EntityId.NULL_UUID);
+            initiatorId = tenantId;
         }
 
-        processEntityUsageStats(tenantId, customerId.isNullUid() ? tenantId : customerId, statsMsg.getValuesList());
+        processEntityUsageStats(tenantId, initiatorId, statsMsg.getValuesList());
         callback.onSuccess();
     }
 
     private void processEntityUsageStats(TenantId tenantId, EntityId entityId, List<UsageStatsKVProto> values) {
-        if (tenantProfileCache.get(tenantId) == null) return;
-
         BaseApiUsageState usageState;
         List<TsKvEntry> updatedEntries;
         Map<ApiFeature, ApiUsageStateValue> result;
@@ -192,7 +196,11 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
                 updatedEntries.add(new BasicTsKvEntry(newHourTs, new LongDataEntry(recordKey.getApiCountKey() + HOURLY, newHourlyValue)));
                 apiFeatures.add(recordKey.getApiFeature());
             }
-            result = usageState.checkStateUpdatedDueToThreshold(apiFeatures);
+            if (usageState.getEntityType() == EntityType.TENANT && !usageState.getEntityId().equals(TenantId.SYS_TENANT_ID)) {
+                result = ((TenantApiUsageState) usageState).checkStateUpdatedDueToThreshold(apiFeatures);
+            } else {
+                result = Collections.emptyMap();
+            }
         } finally {
             updateLock.unlock();
         }
@@ -226,7 +234,7 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
                 return state;
             } else {
                 if (partitionService.resolve(ServiceType.TB_CORE, tenantId, tenantId).isMyPartition()) {
-                    return getOrFetchState(tenantId).getApiUsageState();
+                    return getOrFetchState(tenantId, tenantId).getApiUsageState();
                 } else {
                     updateLock.lock();
                     try {
@@ -319,10 +327,10 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
         clusterService.onApiStateChange(state.getApiUsageState(), null);
         long ts = System.currentTimeMillis();
         List<TsKvEntry> stateTelemetry = new ArrayList<>();
-        result.forEach(((apiFeature, aState) -> stateTelemetry.add(new BasicTsKvEntry(ts, new StringDataEntry(apiFeature.getApiStateKey(), aState.name())))));
+        result.forEach((apiFeature, aState) -> stateTelemetry.add(new BasicTsKvEntry(ts, new StringDataEntry(apiFeature.getApiStateKey(), aState.name()))));
         tsWsService.saveAndNotifyInternal(state.getTenantId(), state.getApiUsageState().getId(), stateTelemetry, VOID_CALLBACK);
 
-        if (state.getEntityType() == EntityType.TENANT) {
+        if (state.getEntityType() == EntityType.TENANT && !state.getEntityId().equals(TenantId.SYS_TENANT_ID)) {
             String email = tenantService.findTenantById(state.getTenantId()).getEmail();
             if (StringUtils.isNotEmpty(email)) {
                 result.forEach((apiFeature, stateValue) -> {
@@ -373,12 +381,13 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
             long now = System.currentTimeMillis();
             myUsageStates.values().forEach(state -> {
                 if ((state.getNextCycleTs() < now) && (now - state.getNextCycleTs() < TimeUnit.HOURS.toMillis(1))) {
-                    // FIXME
-                    TenantId tenantId = state.getTenantId();
                     state.setCycles(state.getNextCycleTs(), SchedulerUtils.getStartOfNextNextMonth());
                     saveNewCounts(state, Arrays.asList(ApiUsageRecordKey.values()));
+                    if (state.getEntityType() == EntityType.TENANT && !state.getEntityId().equals(TenantId.SYS_TENANT_ID)) {
+                        TenantId tenantId = state.getTenantId();
                         updateTenantState((TenantApiUsageState) state, tenantProfileCache.get(tenantId));
                     }
+                }
             });
         } finally {
             updateLock.unlock();
@@ -394,29 +403,30 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
     }
 
     private BaseApiUsageState getOrFetchState(TenantId tenantId, EntityId entityId) {
+        if (entityId == null || entityId.isNullUid()) {
+            entityId = tenantId;
+        }
         BaseApiUsageState state = myUsageStates.get(entityId);
         if (state != null) {
             return state;
         }
 
-        ApiUsageState storedState = Optional.ofNullable(apiUsageStateService.findApiUsageStateByEntityId(entityId))
+        ApiUsageState storedState = apiUsageStateService.findApiUsageStateByEntityId(entityId);
-                .orElseGet(() -> {
+        if (storedState == null) {
             try {
-                        return apiUsageStateService.createDefaultApiUsageState(tenantId, entityId);
+                storedState = apiUsageStateService.createDefaultApiUsageState(tenantId, entityId);
             } catch (Exception e) {
-                        return apiUsageStateService.findApiUsageStateByEntityId(entityId);
+                storedState = apiUsageStateService.findApiUsageStateByEntityId(entityId);
             }
-                });
+        }
-
+        if (entityId.getEntityType() == EntityType.TENANT) {
-        switch (entityId.getEntityType()) {
+            if (!entityId.equals(TenantId.SYS_TENANT_ID)) {
-            case TENANT:
+                state = new TenantApiUsageState(tenantProfileCache.get((TenantId) entityId), storedState);
-                TenantProfile tenantProfile = tenantProfileCache.get(tenantId);
+            } else {
-                state = new TenantApiUsageState(tenantProfile, storedState);
+                state = new TenantApiUsageState(storedState);
-                break;
+            }
-            case CUSTOMER:
+        } else {
-            default:
             state = new CustomerApiUsageState(storedState);
-                break;
         }
 
         List<ApiUsageRecordKey> newCounts = new ArrayList<>();
@@ -454,54 +464,6 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
         return state;
     }
 
-    private TenantApiUsageState getOrFetchState(TenantId tenantId) {
-        TenantApiUsageState tenantState = (TenantApiUsageState) myUsageStates.get(tenantId);
-        if (tenantState == null) {
-            ApiUsageState dbStateEntity = apiUsageStateService.findTenantApiUsageState(tenantId);
-            if (dbStateEntity == null) {
-                try {
-                    dbStateEntity = apiUsageStateService.createDefaultApiUsageState(tenantId, null);
-                } catch (Exception e) {
-                    dbStateEntity = apiUsageStateService.findTenantApiUsageState(tenantId);
-                }
-            }
-            TenantProfile tenantProfile = tenantProfileCache.get(tenantId);
-            tenantState = new TenantApiUsageState(tenantProfile, dbStateEntity);
-            List<ApiUsageRecordKey> newCounts = new ArrayList<>();
-            try {
-                List<TsKvEntry> dbValues = tsService.findAllLatest(tenantId, dbStateEntity.getId()).get();
-                for (ApiUsageRecordKey key : ApiUsageRecordKey.values()) {
-                    boolean cycleEntryFound = false;
-                    boolean hourlyEntryFound = false;
-                    for (TsKvEntry tsKvEntry : dbValues) {
-                        if (tsKvEntry.getKey().equals(key.getApiCountKey())) {
-                            cycleEntryFound = true;
-
-                            boolean oldCount = tsKvEntry.getTs() == tenantState.getCurrentCycleTs();
-                            tenantState.put(key, oldCount ? tsKvEntry.getLongValue().get() : 0L);
-
-                            if (!oldCount) {
-                                newCounts.add(key);
-                            }
-                        } else if (tsKvEntry.getKey().equals(key.getApiCountKey() + HOURLY)) {
-                            hourlyEntryFound = true;
-                            tenantState.putHourly(key, tsKvEntry.getTs() == tenantState.getCurrentHourTs() ? tsKvEntry.getLongValue().get() : 0L);
-                        }
-                        if (cycleEntryFound && hourlyEntryFound) {
-                            break;
-                        }
-                    }
-                }
-                log.debug("[{}] Initialized state: {}", tenantId, dbStateEntity);
-                myUsageStates.put(tenantId, tenantState);
-                saveNewCounts(tenantState, newCounts);
-            } catch (InterruptedException | ExecutionException e) {
-                log.warn("[{}] Failed to fetch api usage state from db.", tenantId, e);
-            }
-        }
-        return tenantState;
-    }
-
     private void initStatesFromDataBase() {
         try {
             log.info("Initializing tenant states.");
@@ -516,7 +478,7 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
                             log.debug("[{}] Initializing tenant state.", tenant.getId());
                             futures.add(tmpInitExecutor.submit(() -> {
                                 try {
-                                    updateTenantState(getOrFetchState(tenant.getId()), tenantProfileCache.get(tenant.getTenantProfileId()));
+                                    updateTenantState((TenantApiUsageState) getOrFetchState(tenant.getId(), tenant.getId()), tenantProfileCache.get(tenant.getTenantProfileId()));
                                     log.debug("[{}] Initialized tenant state.", tenant.getId());
                                 } catch (Exception e) {
                                     log.warn("[{}] Failed to initialize tenant API state", tenant.getId(), e);

application/src/main/java/org/thingsboard/server/service/apiusage/TenantApiUsageState.java

@@ -27,6 +27,12 @@ import org.thingsboard.server.common.data.TenantProfile;
 import org.thingsboard.server.common.data.id.TenantProfileId;
 import org.thingsboard.server.common.data.tenant.profile.TenantProfileData;
 
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
 public class TenantApiUsageState extends BaseApiUsageState {
     @Getter
     @Setter
@@ -41,6 +47,10 @@ public class TenantApiUsageState extends BaseApiUsageState {
         this.tenantProfileData = tenantProfile.getProfileData();
     }
 
+    public TenantApiUsageState(ApiUsageState apiUsageState) {
+        super(apiUsageState);
+    }
+
     public long getProfileThreshold(ApiUsageRecordKey key) {
         return tenantProfileData.getConfiguration().getProfileThreshold(key);
     }
@@ -49,8 +59,7 @@ public class TenantApiUsageState extends BaseApiUsageState {
         return tenantProfileData.getConfiguration().getWarnThreshold(key);
     }
 
-    @Override
+    private Pair<ApiFeature, ApiUsageStateValue> checkStateUpdatedDueToThreshold(ApiFeature feature) {
-    public Pair<ApiFeature, ApiUsageStateValue> checkStateUpdatedDueToThreshold(ApiFeature feature) {
         ApiUsageStateValue featureValue = ApiUsageStateValue.ENABLED;
         for (ApiUsageRecordKey recordKey : ApiUsageRecordKey.getKeys(feature)) {
             long value = get(recordKey);
@@ -69,6 +78,22 @@ public class TenantApiUsageState extends BaseApiUsageState {
         return setFeatureValue(feature, featureValue) ? Pair.of(feature, featureValue) : null;
     }
 
+
+    public Map<ApiFeature, ApiUsageStateValue> checkStateUpdatedDueToThresholds() {
+        return checkStateUpdatedDueToThreshold(new HashSet<>(Arrays.asList(ApiFeature.values())));
+    }
+
+    public Map<ApiFeature, ApiUsageStateValue> checkStateUpdatedDueToThreshold(Set<ApiFeature> features) {
+        Map<ApiFeature, ApiUsageStateValue> result = new HashMap<>();
+        for (ApiFeature feature : features) {
+            Pair<ApiFeature, ApiUsageStateValue> tmp = checkStateUpdatedDueToThreshold(feature);
+            if (tmp != null) {
+                result.put(tmp.getFirst(), tmp.getSecond());
+            }
+        }
+        return result;
+    }
+
     @Override
     public EntityType getEntityType() {
         return EntityType.TENANT;

common/queue/src/main/java/org/thingsboard/server/queue/usagestats/DefaultTbApiUsageClient.java

@@ -142,6 +142,7 @@ public class DefaultTbApiUsageClient implements TbApiUsageClient {
             ConcurrentMap<EntityId, AtomicLong> statsForKey = stats.get(key);
 
             statsForKey.computeIfAbsent(tenantId, id -> new AtomicLong()).addAndGet(value);
+            statsForKey.computeIfAbsent(TenantId.SYS_TENANT_ID, id -> new AtomicLong()).addAndGet(value);
 
             if (customerId != null && !customerId.isNullUid()) {
                 statsForKey.computeIfAbsent(customerId, id -> new AtomicLong()).addAndGet(value);

dao/src/main/java/org/thingsboard/server/dao/usagerecord/ApiUsageStateServiceImpl.java

@@ -85,10 +85,6 @@ public class ApiUsageStateServiceImpl extends AbstractEntityService implements A
 
         ApiUsageState saved = apiUsageStateDao.save(apiUsageState.getTenantId(), apiUsageState);
 
-        Tenant tenant = tenantDao.findById(tenantId, tenantId.getId());
-        TenantProfile tenantProfile = tenantProfileDao.findById(tenantId, tenant.getTenantProfileId().getId());
-        TenantProfileConfiguration configuration = tenantProfile.getProfileData().getConfiguration();
-
         List<TsKvEntry> apiUsageStates = new ArrayList<>();
         apiUsageStates.add(new BasicTsKvEntry(saved.getCreatedTime(),
                 new StringDataEntry(ApiFeature.TRANSPORT.getApiStateKey(), ApiUsageStateValue.ENABLED.name())));
@@ -104,12 +100,19 @@ public class ApiUsageStateServiceImpl extends AbstractEntityService implements A
                 new StringDataEntry(ApiFeature.SMS.getApiStateKey(), ApiUsageStateValue.ENABLED.name())));
         tsService.save(tenantId, saved.getId(), apiUsageStates, 0L);
 
-        List<TsKvEntry> profileThresholds = new ArrayList<>();
+        if (entityId.getEntityType() == EntityType.TENANT && !entityId.equals(TenantId.SYS_TENANT_ID)) {
+            tenantId = (TenantId) entityId;
+            Tenant tenant = tenantDao.findById(tenantId, tenantId.getId());
+            TenantProfile tenantProfile = tenantProfileDao.findById(tenantId, tenant.getTenantProfileId().getId());
+            TenantProfileConfiguration configuration = tenantProfile.getProfileData().getConfiguration();
 
+            List<TsKvEntry> profileThresholds = new ArrayList<>();
             for (ApiUsageRecordKey key : ApiUsageRecordKey.values()) {
                 profileThresholds.add(new BasicTsKvEntry(saved.getCreatedTime(), new LongDataEntry(key.getApiLimitKey(), configuration.getProfileThreshold(key))));
             }
             tsService.save(tenantId, saved.getId(), profileThresholds, 0L);
+        }
+
         return saved;
     }
 
@@ -150,8 +153,8 @@ public class ApiUsageStateServiceImpl extends AbstractEntityService implements A
                         throw new DataValidationException("ApiUsageState should be assigned to tenant!");
                     } else {
                         Tenant tenant = tenantDao.findById(requestTenantId, apiUsageState.getTenantId().getId());
-                        if (tenant == null) {
+                        if (tenant == null && !requestTenantId.equals(TenantId.SYS_TENANT_ID)) {
-                            throw new DataValidationException("Asset is referencing to non-existent tenant!");
+                            throw new DataValidationException("ApiUsageState is referencing to non-existent tenant!");
                         }
                     }
                     if (apiUsageState.getEntityId() == null) {