简柏林 / thingskit · Commits

Commit 9ec4b776728a8da85bc552a5eefb448133ae83bc

Authored by Andrii Shvaika
1 parent 2a65ae7a

Improvements to ApiUsageState persistence

Inline Side-by-side
Showing 7 changed files with 202 additions and 34 deletions
... ... @@ -16,9 +16,11 @@
16 16 package org.thingsboard.server.service.apiusage;
17 17
18 18 import lombok.extern.slf4j.Slf4j;
  19 +import org.springframework.beans.factory.annotation.Value;
19 20 import org.springframework.stereotype.Service;
20 21 import org.thingsboard.server.common.data.ApiUsageRecordKey;
21 22 import org.thingsboard.server.common.data.ApiUsageState;
  23 +import org.thingsboard.server.common.data.TenantProfile;
22 24 import org.thingsboard.server.common.data.id.TenantId;
23 25 import org.thingsboard.server.common.data.kv.BasicTsKvEntry;
24 26 import org.thingsboard.server.common.data.kv.LongDataEntry;
... ... @@ -30,69 +32,144 @@ import org.thingsboard.server.dao.usagerecord.ApiUsageStateService;
30 32 import org.thingsboard.server.gen.transport.TransportProtos.ToUsageStatsServiceMsg;
31 33 import org.thingsboard.server.gen.transport.TransportProtos.UsageStatsKVProto;
32 34 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
  35 +import org.thingsboard.server.queue.scheduler.SchedulerComponent;
33 36 import org.thingsboard.server.queue.util.TbCoreComponent;
  37 +import org.thingsboard.server.service.profile.TbTenantProfileCache;
34 38
35   -import java.time.LocalDate;
36   -import java.time.ZoneId;
  39 +import javax.annotation.PostConstruct;
37 40 import java.util.ArrayList;
38 41 import java.util.List;
39 42 import java.util.Map;
40 43 import java.util.UUID;
41 44 import java.util.concurrent.ConcurrentHashMap;
42 45 import java.util.concurrent.ExecutionException;
  46 +import java.util.concurrent.TimeUnit;
  47 +import java.util.concurrent.locks.Lock;
  48 +import java.util.concurrent.locks.ReentrantLock;
43 49
44 50 @Slf4j
45 51 @TbCoreComponent
46 52 @Service
47 53 public class DefaultTbApiUsageStateService implements TbApiUsageStateService {
48 54
  55 + public static final String HOURLY = "HOURLY_";
49 56 private final ApiUsageStateService apiUsageStateService;
50 57 private final TimeseriesService tsService;
51   - private final ZoneId zoneId;
  58 + private final SchedulerComponent scheduler;
  59 + private final TbTenantProfileCache tenantProfileCache;
52 60 private final Map<TenantId, TenantApiUsageState> tenantStates = new ConcurrentHashMap<>();
53 61
54   - public DefaultTbApiUsageStateService(ApiUsageStateService apiUsageStateService, TimeseriesService tsService) {
  62 + @Value("${usage.stats.report.enabled:true}")
  63 + private boolean enabled;
  64 +
  65 + @Value("${usage.stats.check.cycle:60000}")
  66 + private long nextCycleCheckInterval;
  67 +
  68 + private final Lock updateLock = new ReentrantLock();
  69 +
  70 + public DefaultTbApiUsageStateService(ApiUsageStateService apiUsageStateService, TimeseriesService tsService, SchedulerComponent scheduler, TbTenantProfileCache tenantProfileCache) {
55 71 this.apiUsageStateService = apiUsageStateService;
56 72 this.tsService = tsService;
57   - this.zoneId = SchedulerUtils.getZoneId("UTC");
  73 + this.scheduler = scheduler;
  74 + this.tenantProfileCache = tenantProfileCache;
  75 + }
  76 +
  77 + @PostConstruct
  78 + public void init() {
  79 + if (enabled) {
  80 + scheduler.scheduleAtFixedRate(this::checkStartOfNextCycle, nextCycleCheckInterval, nextCycleCheckInterval, TimeUnit.MILLISECONDS);
  81 + }
58 82 }
59 83
60 84 @Override
61 85 public void process(TbProtoQueueMsg<ToUsageStatsServiceMsg> msg, TbCallback callback) {
62 86 ToUsageStatsServiceMsg statsMsg = msg.getValue();
63 87 TenantId tenantId = new TenantId(new UUID(statsMsg.getTenantIdMSB(), statsMsg.getTenantIdLSB()));
64   - TenantApiUsageState tenantState = getOrFetchState(tenantId);
65   - long ts = tenantState.getCurrentMonthTs();
66   - List<TsKvEntry> updatedEntries = new ArrayList<>(ApiUsageRecordKey.values().length);
67   - for (UsageStatsKVProto kvProto : statsMsg.getValuesList()) {
68   - ApiUsageRecordKey recordKey = ApiUsageRecordKey.valueOf(kvProto.getKey());
69   - long newValue = tenantState.add(recordKey, kvProto.getValue());
70   - updatedEntries.add(new BasicTsKvEntry(ts, new LongDataEntry(recordKey.name(), newValue)));
  88 + TenantApiUsageState tenantState;
  89 + List<TsKvEntry> updatedEntries;
  90 + updateLock.lock();
  91 + try {
  92 + tenantState = getOrFetchState(tenantId);
  93 + long ts = tenantState.getCurrentCycleTs();
  94 + long hourTs = tenantState.getCurrentHourTs();
  95 + long newHourTs = SchedulerUtils.getStartOfCurrentHour();
  96 + if (newHourTs != hourTs) {
  97 + tenantState.setHour(newHourTs);
  98 + }
  99 + updatedEntries = new ArrayList<>(ApiUsageRecordKey.values().length);
  100 + for (UsageStatsKVProto kvProto : statsMsg.getValuesList()) {
  101 + ApiUsageRecordKey recordKey = ApiUsageRecordKey.valueOf(kvProto.getKey());
  102 + long newValue = tenantState.add(recordKey, kvProto.getValue());
  103 + updatedEntries.add(new BasicTsKvEntry(ts, new LongDataEntry(recordKey.name(), newValue)));
  104 + newValue = tenantState.addToHourly(recordKey, kvProto.getValue());
  105 + updatedEntries.add(new BasicTsKvEntry(hourTs, new LongDataEntry(HOURLY + recordKey.name(), newValue)));
  106 + }
  107 + } finally {
  108 + updateLock.unlock();
71 109 }
72 110 tsService.save(tenantId, tenantState.getEntityId(), updatedEntries, 0L);
73 111 }
74 112
  113 + @Override
  114 + public TenantApiUsageState getApiUsageState(TenantId tenantId) {
  115 + return null;
  116 + }
  117 +
  118 + @Override
  119 + public void onAddedToAllowList(TenantId tenantId) {
  120 +
  121 + }
  122 +
  123 + @Override
  124 + public void onAddedToDenyList(TenantId tenantId) {
  125 +
  126 + }
  127 +
  128 + private void checkStartOfNextCycle() {
  129 + updateLock.lock();
  130 + try {
  131 + long now = System.currentTimeMillis();
  132 + tenantStates.values().forEach(state -> {
  133 + if ((state.getNextCycleTs() > now) && (state.getNextCycleTs() - now < TimeUnit.HOURS.toMillis(1))) {
  134 + state.setCycles(state.getNextCycleTs(), SchedulerUtils.getStartOfNextNextMonth());
  135 + }
  136 + });
  137 + } finally {
  138 + updateLock.unlock();
  139 + }
  140 + }
  141 +
75 142 private TenantApiUsageState getOrFetchState(TenantId tenantId) {
76 143 TenantApiUsageState tenantState = tenantStates.get(tenantId);
77 144 if (tenantState == null) {
78   - long currentMonthTs = LocalDate.now().withDayOfMonth(1).atStartOfDay(zoneId).toInstant().toEpochMilli();
79 145 ApiUsageState dbStateEntity = apiUsageStateService.findTenantApiUsageState(tenantId);
80   - tenantState = new TenantApiUsageState(currentMonthTs, dbStateEntity.getEntityId());
  146 + if (dbStateEntity == null) {
  147 + try {
  148 + dbStateEntity = apiUsageStateService.createDefaultApiUsageState(tenantId);
  149 + } catch (Exception e) {
  150 + dbStateEntity = apiUsageStateService.findTenantApiUsageState(tenantId);
  151 + }
  152 + }
  153 + tenantState = new TenantApiUsageState(dbStateEntity.getEntityId());
81 154 try {
82 155 List<TsKvEntry> dbValues = tsService.findAllLatest(tenantId, dbStateEntity.getEntityId()).get();
83 156 for (ApiUsageRecordKey key : ApiUsageRecordKey.values()) {
84   - TsKvEntry keyEntry = null;
  157 + boolean cycleEntryFound = false;
  158 + boolean hourlyEntryFound = false;
85 159 for (TsKvEntry tsKvEntry : dbValues) {
86 160 if (tsKvEntry.getKey().equals(key.name())) {
87   - keyEntry = tsKvEntry;
  161 + cycleEntryFound = true;
  162 + tenantState.put(key, tsKvEntry.getLongValue().get());
  163 + } else if (tsKvEntry.getKey().equals(HOURLY + key.name())) {
  164 + hourlyEntryFound = true;
  165 + if (tsKvEntry.getTs() == tenantState.getCurrentHourTs()) {
  166 + tenantState.putHourly(key, tsKvEntry.getLongValue().get());
  167 + }
  168 + }
  169 + if (cycleEntryFound && hourlyEntryFound) {
88 170 break;
89 171 }
90 172 }
91   - if (keyEntry != null) {
92   - tenantState.put(key, keyEntry.getLongValue().get());
93   - } else {
94   - tenantState.put(key, 0L);
95   - }
96 173 }
97 174 tenantStates.put(tenantId, tenantState);
98 175 } catch (InterruptedException | ExecutionException e) {
... ...
... ... @@ -15,11 +15,19 @@
15 15 */
16 16 package org.thingsboard.server.service.apiusage;
17 17
  18 +import org.thingsboard.server.common.data.id.TenantId;
18 19 import org.thingsboard.server.common.msg.queue.TbCallback;
19   -import org.thingsboard.server.gen.transport.TransportProtos;
  20 +import org.thingsboard.server.gen.transport.TransportProtos.ToUsageStatsServiceMsg;
20 21 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
21 22
22 23 public interface TbApiUsageStateService {
23 24
24   - void process(TbProtoQueueMsg<TransportProtos.ToUsageStatsServiceMsg> msg, TbCallback callback);
  25 + void process(TbProtoQueueMsg<ToUsageStatsServiceMsg> msg, TbCallback callback);
  26 +
  27 + TenantApiUsageState getApiUsageState(TenantId tenantId);
  28 +
  29 + void onAddedToAllowList(TenantId tenantId);
  30 +
  31 + void onAddedToDenyList(TenantId tenantId);
  32 +
25 33 }
... ...
... ... @@ -18,30 +18,65 @@ package org.thingsboard.server.service.apiusage;
18 18 import lombok.Getter;
19 19 import org.thingsboard.server.common.data.ApiUsageRecordKey;
20 20 import org.thingsboard.server.common.data.id.EntityId;
  21 +import org.thingsboard.server.common.msg.tools.SchedulerUtils;
21 22
22 23 import java.util.Map;
23 24 import java.util.concurrent.ConcurrentHashMap;
24 25
25 26 public class TenantApiUsageState {
26 27
27   - private final Map<ApiUsageRecordKey, Long> values = new ConcurrentHashMap<>();
  28 + private final Map<ApiUsageRecordKey, Long> currentCycleValues = new ConcurrentHashMap<>();
  29 + private final Map<ApiUsageRecordKey, Long> currentHourValues = new ConcurrentHashMap<>();
  30 +
28 31 @Getter
29 32 private final EntityId entityId;
30 33 @Getter
31   - private volatile long currentMonthTs;
  34 + private volatile long currentCycleTs;
  35 + @Getter
  36 + private volatile long nextCycleTs;
  37 + @Getter
  38 + private volatile long currentHourTs;
32 39
33   - public TenantApiUsageState(long currentMonthTs, EntityId entityId) {
  40 + public TenantApiUsageState(EntityId entityId) {
34 41 this.entityId = entityId;
35   - this.currentMonthTs = currentMonthTs;
  42 + this.currentCycleTs = SchedulerUtils.getStartOfCurrentMonth();
  43 + this.nextCycleTs = SchedulerUtils.getStartOfNextMonth();
  44 + this.currentHourTs = SchedulerUtils.getStartOfCurrentHour();
36 45 }
37 46
38 47 public void put(ApiUsageRecordKey key, Long value) {
39   - values.put(key, value);
  48 + currentCycleValues.put(key, value);
  49 + }
  50 +
  51 + public void putHourly(ApiUsageRecordKey key, Long value) {
  52 + currentHourValues.put(key, value);
40 53 }
41 54
42 55 public long add(ApiUsageRecordKey key, long value) {
43   - long result = values.getOrDefault(key, 0L) + value;
44   - values.put(key, result);
  56 + long result = currentCycleValues.getOrDefault(key, 0L) + value;
  57 + currentCycleValues.put(key, result);
45 58 return result;
46 59 }
  60 +
  61 + public long addToHourly(ApiUsageRecordKey key, long value) {
  62 + long result = currentHourValues.getOrDefault(key, 0L) + value;
  63 + currentHourValues.put(key, result);
  64 + return result;
  65 + }
  66 +
  67 + public void setHour(long currentHourTs) {
  68 + this.currentHourTs = currentHourTs;
  69 + for (ApiUsageRecordKey key : ApiUsageRecordKey.values()) {
  70 + currentHourValues.put(key, 0L);
  71 + }
  72 + }
  73 +
  74 + public void setCycles(long currentCycleTs, long nextCycleTs) {
  75 + this.currentCycleTs = currentCycleTs;
  76 + this.nextCycleTs = nextCycleTs;
  77 + for (ApiUsageRecordKey key : ApiUsageRecordKey.values()) {
  78 + currentCycleValues.put(key, 0L);
  79 + }
  80 + }
  81 +
47 82 }
... ...
... ... @@ -24,5 +24,5 @@ public interface ApiUsageStateService {
24 24
25 25 void deleteApiUsageStateByTenantId(TenantId tenantId);
26 26
27   - void createDefaultApiUsageState(TenantId id);
  27 + ApiUsageState createDefaultApiUsageState(TenantId id);
28 28 }
... ...
... ... @@ -21,5 +21,5 @@ import java.io.Serializable;
21 21 * @author Andrew Shvayka
22 22 */
23 23 public enum ComponentLifecycleEvent implements Serializable {
24   - CREATED, STARTED, ACTIVATED, SUSPENDED, UPDATED, STOPPED, DELETED
  24 + CREATED, STARTED, ACTIVATED, SUSPENDED, UPDATED, STOPPED, DELETED, ADDED_TO_ALLOW_LIST, ADDED_TO_DENY_LIST
25 25 }
\ No newline at end of file
... ...
... ... @@ -15,16 +15,64 @@
15 15 */
16 16 package org.thingsboard.server.common.msg.tools;
17 17
  18 +import java.time.LocalDate;
  19 +import java.time.LocalDateTime;
  20 +import java.time.LocalTime;
18 21 import java.time.ZoneId;
  22 +import java.time.ZoneOffset;
  23 +import java.time.temporal.ChronoUnit;
  24 +import java.time.temporal.TemporalAdjuster;
  25 +import java.time.temporal.TemporalAdjusters;
  26 +import java.time.temporal.TemporalUnit;
19 27 import java.util.concurrent.ConcurrentHashMap;
20 28 import java.util.concurrent.ConcurrentMap;
21 29
  30 +import static java.time.temporal.ChronoField.DAY_OF_MONTH;
  31 +import static java.time.temporal.ChronoUnit.MONTHS;
  32 +
22 33 public class SchedulerUtils {
23 34
  35 + private final static ZoneId UTC = ZoneId.of("UTC");
24 36 private static final ConcurrentMap<String, ZoneId> tzMap = new ConcurrentHashMap<>();
25 37
26 38 public static ZoneId getZoneId(String tz) {
27 39 return tzMap.computeIfAbsent(tz == null || tz.isEmpty() ? "UTC" : tz, ZoneId::of);
28 40 }
29 41
  42 + public static long getStartOfCurrentHour() {
  43 + return getStartOfCurrentHour(UTC);
  44 + }
  45 +
  46 + public static long getStartOfCurrentHour(ZoneId zoneId) {
  47 + return LocalDateTime.now(ZoneOffset.UTC).atZone(zoneId).truncatedTo(ChronoUnit.HOURS).toInstant().toEpochMilli();
  48 + }
  49 +
  50 + public static long getStartOfCurrentMonth() {
  51 + return getStartOfCurrentMonth(UTC);
  52 + }
  53 +
  54 + public static long getStartOfCurrentMonth(ZoneId zoneId) {
  55 + return LocalDate.now().withDayOfMonth(1).atStartOfDay(zoneId).toInstant().toEpochMilli();
  56 + }
  57 +
  58 + public static long getStartOfNextMonth() {
  59 + return getStartOfNextMonth(UTC);
  60 + }
  61 +
  62 + public static long getStartOfNextMonth(ZoneId zoneId) {
  63 + return LocalDate.now().with(TemporalAdjusters.firstDayOfNextMonth()).atStartOfDay(zoneId).toInstant().toEpochMilli();
  64 + }
  65 +
  66 + public static long getStartOfNextNextMonth() {
  67 + return getStartOfNextNextMonth(UTC);
  68 + }
  69 +
  70 + public static long getStartOfNextNextMonth(ZoneId zoneId) {
  71 + return LocalDate.now().with(firstDayOfNextNextMonth()).atStartOfDay(zoneId).toInstant().toEpochMilli();
  72 + }
  73 +
  74 + public static TemporalAdjuster firstDayOfNextNextMonth() {
  75 + return (temporal) -> temporal.with(DAY_OF_MONTH, 1).plus(2, MONTHS);
  76 + }
  77 +
30 78 }
... ...
... ... @@ -49,14 +49,14 @@ public class ApiApiUsageStateServiceImpl extends AbstractEntityService implement
49 49 }
50 50
51 51 @Override
52   - public void createDefaultApiUsageState(TenantId tenantId) {
  52 + public ApiUsageState createDefaultApiUsageState(TenantId tenantId) {
53 53 log.trace("Executing createDefaultUsageRecord [{}]", tenantId);
54 54 validateId(tenantId, INCORRECT_TENANT_ID + tenantId);
55 55 ApiUsageState apiUsageState = new ApiUsageState();
56 56 apiUsageState.setTenantId(tenantId);
57 57 apiUsageState.setEntityId(tenantId);
58 58 apiUsageStateValidator.validate(apiUsageState, ApiUsageState::getTenantId);
59   - apiUsageStateDao.save(apiUsageState.getTenantId(), apiUsageState);
  59 + return apiUsageStateDao.save(apiUsageState.getTenantId(), apiUsageState);
60 60 }
61 61
62 62 @Override
... ...