Andrii Shvaika
5 years ago
5 changed files with 369 additions and 179 deletions
@ -0,0 +1,146 @@ |
|||||
|
/** |
||||
|
* Copyright © 2016-2022 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.partition; |
||||
|
|
||||
|
import com.google.common.util.concurrent.ListeningScheduledExecutorService; |
||||
|
import com.google.common.util.concurrent.MoreExecutors; |
||||
|
import lombok.extern.slf4j.Slf4j; |
||||
|
import org.thingsboard.common.util.ThingsBoardThreadFactory; |
||||
|
import org.thingsboard.server.common.data.id.EntityId; |
||||
|
import org.thingsboard.server.common.msg.queue.ServiceType; |
||||
|
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; |
||||
|
import org.thingsboard.server.queue.discovery.TbApplicationEventListener; |
||||
|
import org.thingsboard.server.queue.discovery.event.PartitionChangeEvent; |
||||
|
|
||||
|
import java.util.HashSet; |
||||
|
import java.util.Queue; |
||||
|
import java.util.Set; |
||||
|
import java.util.concurrent.ConcurrentHashMap; |
||||
|
import java.util.concurrent.ConcurrentLinkedQueue; |
||||
|
import java.util.concurrent.ConcurrentMap; |
||||
|
import java.util.concurrent.Executors; |
||||
|
|
||||
|
@Slf4j |
||||
|
public abstract class AbstractPartitionBasedService<T extends EntityId> extends TbApplicationEventListener<PartitionChangeEvent> { |
||||
|
|
||||
|
protected final ConcurrentMap<TopicPartitionInfo, Set<T>> partitionedEntities = new ConcurrentHashMap<>(); |
||||
|
final Queue<Set<TopicPartitionInfo>> subscribeQueue = new ConcurrentLinkedQueue<>(); |
||||
|
|
||||
|
protected ListeningScheduledExecutorService scheduledExecutor; |
||||
|
|
||||
|
abstract protected String getSchedulerExecutorName(); |
||||
|
|
||||
|
abstract protected void onAddedPartitions(Set<TopicPartitionInfo> addedPartitions); |
||||
|
|
||||
|
abstract protected void cleanupEntityOnPartitionRemoval(T entityId); |
||||
|
|
||||
|
public Set<T> getPartitionedEntities(TopicPartitionInfo tpi) { |
||||
|
return partitionedEntities.get(tpi); |
||||
|
} |
||||
|
|
||||
|
protected void init() { |
||||
|
// Should be always single threaded due to absence of locks.
|
||||
|
scheduledExecutor = MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor(ThingsBoardThreadFactory.forName("device-state-scheduled"))); |
||||
|
} |
||||
|
|
||||
|
protected ServiceType getServiceType() { |
||||
|
return ServiceType.TB_CORE; |
||||
|
} |
||||
|
|
||||
|
protected void stop() { |
||||
|
if (scheduledExecutor != null) { |
||||
|
scheduledExecutor.shutdownNow(); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
/** |
||||
|
* DiscoveryService will call this event from the single thread (one-by-one). |
||||
|
* Events order is guaranteed by DiscoveryService. |
||||
|
* The only concurrency is expected from the [main] thread on Application started. |
||||
|
* Async implementation. Locks is not allowed by design. |
||||
|
* Any locks or delays in this module will affect DiscoveryService and entire system |
||||
|
*/ |
||||
|
@Override |
||||
|
protected void onTbApplicationEvent(PartitionChangeEvent partitionChangeEvent) { |
||||
|
if (getServiceType().equals(partitionChangeEvent.getServiceType())) { |
||||
|
log.debug("onTbApplicationEvent, processing event: {}", partitionChangeEvent); |
||||
|
subscribeQueue.add(partitionChangeEvent.getPartitions()); |
||||
|
scheduledExecutor.submit(this::pollInitStateFromDB); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
protected void pollInitStateFromDB() { |
||||
|
final Set<TopicPartitionInfo> partitions = getLatestPartitions(); |
||||
|
if (partitions == null) { |
||||
|
log.debug("Nothing to do. Partitions are empty."); |
||||
|
return; |
||||
|
} |
||||
|
initStateFromDB(partitions); |
||||
|
} |
||||
|
|
||||
|
private void initStateFromDB(Set<TopicPartitionInfo> partitions) { |
||||
|
try { |
||||
|
log.info("CURRENT PARTITIONS: {}", partitionedEntities.keySet()); |
||||
|
log.info("NEW PARTITIONS: {}", partitions); |
||||
|
|
||||
|
Set<TopicPartitionInfo> addedPartitions = new HashSet<>(partitions); |
||||
|
addedPartitions.removeAll(partitionedEntities.keySet()); |
||||
|
|
||||
|
log.info("ADDED PARTITIONS: {}", addedPartitions); |
||||
|
|
||||
|
Set<TopicPartitionInfo> removedPartitions = new HashSet<>(partitionedEntities.keySet()); |
||||
|
removedPartitions.removeAll(partitions); |
||||
|
|
||||
|
log.info("REMOVED PARTITIONS: {}", removedPartitions); |
||||
|
|
||||
|
// We no longer manage current partition of entities;
|
||||
|
removedPartitions.forEach(partition -> { |
||||
|
Set<T> entities = partitionedEntities.remove(partition); |
||||
|
entities.forEach(this::cleanupEntityOnPartitionRemoval); |
||||
|
}); |
||||
|
|
||||
|
onRepartitionEvent(); |
||||
|
|
||||
|
addedPartitions.forEach(tpi -> partitionedEntities.computeIfAbsent(tpi, key -> ConcurrentHashMap.newKeySet())); |
||||
|
|
||||
|
if (!addedPartitions.isEmpty()) { |
||||
|
onAddedPartitions(addedPartitions); |
||||
|
} |
||||
|
|
||||
|
log.info("Managing following partitions:"); |
||||
|
partitionedEntities.forEach((tpi, entities) -> { |
||||
|
log.info("[{}]: {} entities", tpi.getFullTopicName(), entities.size()); |
||||
|
}); |
||||
|
} catch (Throwable t) { |
||||
|
log.warn("Failed to init entities state from DB", t); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
protected void onRepartitionEvent() { |
||||
|
} |
||||
|
|
||||
|
private Set<TopicPartitionInfo> getLatestPartitions() { |
||||
|
log.debug("getLatestPartitionsFromQueue, queue size {}", subscribeQueue.size()); |
||||
|
Set<TopicPartitionInfo> partitions = null; |
||||
|
while (!subscribeQueue.isEmpty()) { |
||||
|
partitions = subscribeQueue.poll(); |
||||
|
log.debug("polled from the queue partitions {}", partitions); |
||||
|
} |
||||
|
log.debug("getLatestPartitionsFromQueue, partitions {}", partitions); |
||||
|
return partitions; |
||||
|
} |
||||
|
|
||||
|
} |
||||
@ -0,0 +1,86 @@ |
|||||
|
/** |
||||
|
* Copyright © 2016-2022 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.junit.Before; |
||||
|
import org.junit.Test; |
||||
|
import org.junit.runner.RunWith; |
||||
|
import org.mockito.Mock; |
||||
|
import org.mockito.Mockito; |
||||
|
import org.mockito.junit.MockitoJUnitRunner; |
||||
|
import org.thingsboard.rule.engine.api.MailService; |
||||
|
import org.thingsboard.server.cluster.TbClusterService; |
||||
|
import org.thingsboard.server.common.data.ApiUsageState; |
||||
|
import org.thingsboard.server.common.data.id.TenantId; |
||||
|
import org.thingsboard.server.common.msg.queue.ServiceType; |
||||
|
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; |
||||
|
import org.thingsboard.server.dao.tenant.TbTenantProfileCache; |
||||
|
import org.thingsboard.server.dao.tenant.TenantService; |
||||
|
import org.thingsboard.server.dao.timeseries.TimeseriesService; |
||||
|
import org.thingsboard.server.dao.usagerecord.ApiUsageStateService; |
||||
|
import org.thingsboard.server.queue.discovery.PartitionService; |
||||
|
import org.thingsboard.server.queue.scheduler.SchedulerComponent; |
||||
|
import org.thingsboard.server.service.executors.DbCallbackExecutorService; |
||||
|
|
||||
|
import java.util.UUID; |
||||
|
|
||||
|
import static org.hamcrest.CoreMatchers.is; |
||||
|
import static org.hamcrest.MatcherAssert.assertThat; |
||||
|
import static org.mockito.BDDMockito.willReturn; |
||||
|
import static org.mockito.Mockito.never; |
||||
|
import static org.mockito.Mockito.spy; |
||||
|
import static org.mockito.Mockito.times; |
||||
|
|
||||
|
@RunWith(MockitoJUnitRunner.class) |
||||
|
public class DefaultTbApiUsageStateServiceTest { |
||||
|
|
||||
|
@Mock |
||||
|
TenantService tenantService; |
||||
|
@Mock |
||||
|
TimeseriesService tsService; |
||||
|
@Mock |
||||
|
TbClusterService clusterService; |
||||
|
@Mock |
||||
|
PartitionService partitionService; |
||||
|
@Mock |
||||
|
TenantApiUsageState tenantUsageStateMock; |
||||
|
@Mock |
||||
|
ApiUsageStateService apiUsageStateService; |
||||
|
@Mock |
||||
|
TbTenantProfileCache tenantProfileCache; |
||||
|
@Mock |
||||
|
MailService mailService; |
||||
|
@Mock |
||||
|
DbCallbackExecutorService dbExecutor; |
||||
|
|
||||
|
TenantId tenantId = TenantId.fromUUID(UUID.fromString("00797a3b-7aeb-4b5b-b57a-c2a810d0f112")); |
||||
|
|
||||
|
DefaultTbApiUsageStateService service; |
||||
|
|
||||
|
@Before |
||||
|
public void setUp() { |
||||
|
service = spy(new DefaultTbApiUsageStateService(clusterService, partitionService, tenantService, tsService, apiUsageStateService, tenantProfileCache, mailService, dbExecutor)); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
public void givenTenantIdFromEntityStatesMap_whenGetApiUsageState() { |
||||
|
service.myUsageStates.put(tenantId, tenantUsageStateMock); |
||||
|
ApiUsageState tenantUsageState = service.getApiUsageState(tenantId); |
||||
|
assertThat(tenantUsageState, is(tenantUsageStateMock.getApiUsageState())); |
||||
|
Mockito.verify(service, never()).getOrFetchState(tenantId, tenantId); |
||||
|
} |
||||
|
|
||||
|
} |
||||
Loading…
Reference in new issue