Handle tenant and customer deletion for TbApiUsageStateService

@@ -38,6 +38,7 @@ import org.thingsboard.server.common.data.id.EdgeId;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.page.PageData;
 import org.thingsboard.server.common.data.page.PageLink;
+import org.thingsboard.server.common.data.plugin.ComponentLifecycleEvent;
 import org.thingsboard.server.queue.util.TbCoreComponent;
 import org.thingsboard.server.service.security.permission.Operation;
 import org.thingsboard.server.service.security.permission.Resource;
@@ -148,6 +149,7 @@ public class CustomerController extends BaseController {
                     ActionType.DELETED, null, strCustomerId);
 
             sendDeleteNotificationMsg(getTenantId(), customerId, relatedEdgeIds);
+            tbClusterService.onEntityStateChange(getTenantId(), customerId, ComponentLifecycleEvent.DELETED);
         } catch (Exception e) {
 
             logEntityAction(emptyId(EntityType.CUSTOMER),

@@ -111,11 +111,13 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
     @Autowired
     private InternalTelemetryService tsWsService;
 
-    // Tenants that should be processed on this server
+    // Entities that should be processed on this server
     private final Map<EntityId, BaseApiUsageState> myUsageStates = new ConcurrentHashMap<>();
-    // Tenants that should be processed on other servers
+    // Entities that should be processed on other servers
     private final Map<EntityId, ApiUsageState> otherUsageStates = new ConcurrentHashMap<>();
 
+    private final Set<EntityId> deletedEntities = Collections.newSetFromMap(new ConcurrentHashMap<>());
+
     @Value("${usage.stats.report.enabled:true}")
     private boolean enabled;
 
@@ -173,6 +175,8 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
     }
 
     private void processEntityUsageStats(TenantId tenantId, EntityId entityId, List<UsageStatsKVProto> values) {
+        if (deletedEntities.contains(entityId)) return;
+
         BaseApiUsageState usageState;
         List<TsKvEntry> updatedEntries;
         Map<ApiFeature, ApiUsageStateValue> result;
@@ -321,6 +325,18 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
         }
     }
 
+    public void onTenantDelete(TenantId tenantId) {
+        deletedEntities.add(tenantId);
+        myUsageStates.remove(tenantId);
+        otherUsageStates.remove(tenantId);
+    }
+
+    @Override
+    public void onCustomerDelete(CustomerId customerId) {
+        deletedEntities.add(customerId);
+        myUsageStates.remove(customerId);
+    }
+
     private void persistAndNotify(BaseApiUsageState state, Map<ApiFeature, ApiUsageStateValue> result) {
         log.info("[{}] Detected update of the API state for {}: {}", state.getEntityId(), state.getEntityType(), result);
         apiUsageStateService.update(state.getApiUsageState());

@@ -17,6 +17,7 @@ package org.thingsboard.server.service.apiusage;
 
 import org.springframework.context.ApplicationListener;
 import org.thingsboard.server.common.data.ApiUsageState;
+import org.thingsboard.server.common.data.id.CustomerId;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.id.TenantProfileId;
 import org.thingsboard.server.common.msg.queue.TbCallback;
@@ -34,5 +35,9 @@ public interface TbApiUsageStateService extends ApplicationListener<PartitionCha
 
     void onTenantUpdate(TenantId tenantId);
 
+    void onTenantDelete(TenantId tenantId);
+
+    void onCustomerDelete(CustomerId customerId);
+
     void onApiUsageStateUpdate(TenantId tenantId);
 }

@@ -18,14 +18,15 @@ package org.thingsboard.server.service.queue.processing;
 import com.google.protobuf.ByteString;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.boot.context.event.ApplicationReadyEvent;
-import org.springframework.context.ApplicationListener;
 import org.springframework.context.event.EventListener;
 import org.springframework.core.annotation.Order;
 import org.thingsboard.common.util.ThingsBoardThreadFactory;
 import org.thingsboard.server.actors.ActorSystemContext;
 import org.thingsboard.server.common.data.EntityType;
+import org.thingsboard.server.common.data.id.CustomerId;
 import org.thingsboard.server.common.data.id.DeviceId;
 import org.thingsboard.server.common.data.id.DeviceProfileId;
+import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.id.TenantProfileId;
 import org.thingsboard.server.common.data.plugin.ComponentLifecycleEvent;
 import org.thingsboard.server.common.msg.TbActorMsg;
@@ -72,7 +73,8 @@ public abstract class AbstractConsumerService<N extends com.google.protobuf.Gene
     protected final TbQueueConsumer<TbProtoQueueMsg<N>> nfConsumer;
 
     public AbstractConsumerService(ActorSystemContext actorContext, DataDecodingEncodingService encodingService,
-                                   TbTenantProfileCache tenantProfileCache, TbDeviceProfileCache deviceProfileCache, TbApiUsageStateService apiUsageStateService, TbQueueConsumer<TbProtoQueueMsg<N>> nfConsumer) {
+                                   TbTenantProfileCache tenantProfileCache, TbDeviceProfileCache deviceProfileCache,
+                                   TbApiUsageStateService apiUsageStateService, TbQueueConsumer<TbProtoQueueMsg<N>> nfConsumer) {
         this.actorContext = actorContext;
         this.encodingService = encodingService;
         this.tenantProfileCache = tenantProfileCache;
@@ -166,6 +168,8 @@ public abstract class AbstractConsumerService<N extends com.google.protobuf.Gene
                     tenantProfileCache.evict(componentLifecycleMsg.getTenantId());
                     if (componentLifecycleMsg.getEvent().equals(ComponentLifecycleEvent.UPDATED)) {
                         apiUsageStateService.onTenantUpdate(componentLifecycleMsg.getTenantId());
+                    } else if (componentLifecycleMsg.getEvent().equals(ComponentLifecycleEvent.DELETED)) {
+                        apiUsageStateService.onTenantDelete((TenantId) componentLifecycleMsg.getEntityId());
                     }
                 } else if (EntityType.DEVICE_PROFILE.equals(componentLifecycleMsg.getEntityId().getEntityType())) {
                     deviceProfileCache.evict(componentLifecycleMsg.getTenantId(), new DeviceProfileId(componentLifecycleMsg.getEntityId().getId()));
@@ -173,6 +177,10 @@ public abstract class AbstractConsumerService<N extends com.google.protobuf.Gene
                     deviceProfileCache.evict(componentLifecycleMsg.getTenantId(), new DeviceId(componentLifecycleMsg.getEntityId().getId()));
                 } else if (EntityType.API_USAGE_STATE.equals(componentLifecycleMsg.getEntityId().getEntityType())) {
                     apiUsageStateService.onApiUsageStateUpdate(componentLifecycleMsg.getTenantId());
+                } else if (EntityType.CUSTOMER.equals(componentLifecycleMsg.getEntityId().getEntityType())) {
+                    if (componentLifecycleMsg.getEvent() == ComponentLifecycleEvent.DELETED) {
+                        apiUsageStateService.onCustomerDelete((CustomerId) componentLifecycleMsg.getEntityId());
+                    }
                 }
             }
             log.trace("[{}] Forwarding message to App Actor {}", id, actorMsg);

@@ -32,5 +32,7 @@ public interface ApiUsageStateService {
 
     void deleteApiUsageStateByTenantId(TenantId tenantId);
 
+    void deleteApiUsageStateByEntityId(EntityId entityId);
+
     ApiUsageState findApiUsageStateById(TenantId tenantId, ApiUsageStateId id);
 }

@@ -43,6 +43,7 @@ import org.thingsboard.server.dao.service.PaginatedRemover;
 import org.thingsboard.server.dao.service.Validator;
 import org.thingsboard.server.dao.tenant.TbTenantProfileCache;
 import org.thingsboard.server.dao.tenant.TenantDao;
+import org.thingsboard.server.dao.usagerecord.ApiUsageStateService;
 import org.thingsboard.server.dao.user.UserService;
 
 import java.io.IOException;
@@ -80,6 +81,9 @@ public class CustomerServiceImpl extends AbstractEntityService implements Custom
     private DashboardService dashboardService;
 
     @Autowired
+    private ApiUsageStateService apiUsageStateService;
+
+    @Autowired
     @Lazy
     private TbTenantProfileCache tenantProfileCache;
 
@@ -128,6 +132,7 @@ public class CustomerServiceImpl extends AbstractEntityService implements Custom
         edgeService.unassignCustomerEdges(customer.getTenantId(), customerId);
         userService.deleteCustomerUsers(customer.getTenantId(), customerId);
         deleteEntityRelations(tenantId, customerId);
+        apiUsageStateService.deleteApiUsageStateByEntityId(customerId);
         customerDao.removeById(tenantId, customerId.getId());
     }
 

@@ -39,4 +39,8 @@ public interface ApiUsageStateRepository extends CrudRepository<ApiUsageStateEnt
     @Modifying
     @Query("DELETE FROM ApiUsageStateEntity ur WHERE ur.tenantId = :tenantId")
     void deleteApiUsageStateByTenantId(@Param("tenantId") UUID tenantId);
+
+    @Transactional
+    @Modifying
+    void deleteByEntityIdAndEntityType(UUID entityId, String entityType);
 }

@@ -63,4 +63,9 @@ public class JpaApiUsageStateDao extends JpaAbstractDao<ApiUsageStateEntity, Api
     public void deleteApiUsageStateByTenantId(TenantId tenantId) {
         apiUsageStateRepository.deleteApiUsageStateByTenantId(tenantId.getId());
     }
+
+    @Override
+    public void deleteApiUsageStateByEntityId(EntityId entityId) {
+        apiUsageStateRepository.deleteByEntityIdAndEntityType(entityId.getId(), entityId.getEntityType().name());
+    }
 }

@@ -48,4 +48,6 @@ public interface ApiUsageStateDao extends Dao<ApiUsageState> {
      * @param tenantId the tenantId
      */
     void deleteApiUsageStateByTenantId(TenantId tenantId);
+
+    void deleteApiUsageStateByEntityId(EntityId entityId);
 }

@@ -70,6 +70,13 @@ public class ApiUsageStateServiceImpl extends AbstractEntityService implements A
     }
 
     @Override
+    public void deleteApiUsageStateByEntityId(EntityId entityId) {
+        log.trace("Executing deleteApiUsageStateByEntityId [{}]", entityId);
+        validateId(entityId.getId(), "Invalid entity id");
+        apiUsageStateDao.deleteApiUsageStateByEntityId(entityId);
+    }
+
+    @Override
     public ApiUsageState createDefaultApiUsageState(TenantId tenantId, EntityId entityId) {
         entityId = Objects.requireNonNullElse(entityId, tenantId);
         log.trace("Executing createDefaultUsageRecord [{}]", entityId);