Andrii Shvaika
3 years ago
48 changed files with 2080 additions and 666 deletions
@ -0,0 +1,65 @@ |
|||
/** |
|||
* Copyright © 2016-2023 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.queue; |
|||
|
|||
import org.thingsboard.server.common.data.EntityType; |
|||
import org.thingsboard.server.common.data.id.EntityIdFactory; |
|||
import org.thingsboard.server.common.data.id.TenantId; |
|||
import org.thingsboard.server.common.data.plugin.ComponentLifecycleEvent; |
|||
import org.thingsboard.server.common.msg.plugin.ComponentLifecycleMsg; |
|||
import org.thingsboard.server.gen.transport.TransportProtos; |
|||
|
|||
import java.util.Arrays; |
|||
import java.util.UUID; |
|||
|
|||
public class ProtoUtils { |
|||
|
|||
private static final EntityType[] entityTypeByProtoNumber; |
|||
|
|||
static { |
|||
int arraySize = Arrays.stream(EntityType.values()).mapToInt(EntityType::getProtoNumber).max().orElse(0); |
|||
entityTypeByProtoNumber = new EntityType[arraySize + 1]; |
|||
Arrays.stream(EntityType.values()).forEach(entityType -> entityTypeByProtoNumber[entityType.getProtoNumber()] = entityType); |
|||
} |
|||
|
|||
public static TransportProtos.ComponentLifecycleMsgProto toProto(ComponentLifecycleMsg msg) { |
|||
return TransportProtos.ComponentLifecycleMsgProto.newBuilder() |
|||
.setTenantIdMSB(msg.getTenantId().getId().getMostSignificantBits()) |
|||
.setTenantIdLSB(msg.getTenantId().getId().getLeastSignificantBits()) |
|||
.setEntityType(toProto(msg.getEntityId().getEntityType())) |
|||
.setEntityIdMSB(msg.getEntityId().getId().getMostSignificantBits()) |
|||
.setEntityIdLSB(msg.getEntityId().getId().getLeastSignificantBits()) |
|||
.setEvent(TransportProtos.ComponentLifecycleEvent.forNumber(msg.getEvent().ordinal())) |
|||
.build(); |
|||
} |
|||
|
|||
public static TransportProtos.EntityTypeProto toProto(EntityType entityType) { |
|||
return TransportProtos.EntityTypeProto.forNumber(entityType.getProtoNumber()); |
|||
} |
|||
|
|||
public static ComponentLifecycleMsg fromProto(TransportProtos.ComponentLifecycleMsgProto proto) { |
|||
return new ComponentLifecycleMsg( |
|||
TenantId.fromUUID(new UUID(proto.getTenantIdMSB(), proto.getTenantIdLSB())), |
|||
EntityIdFactory.getByTypeAndUuid(fromProto(proto.getEntityType()), new UUID(proto.getEntityIdMSB(), proto.getEntityIdLSB())), |
|||
ComponentLifecycleEvent.values()[proto.getEventValue()] |
|||
); |
|||
} |
|||
|
|||
public static EntityType fromProto(TransportProtos.EntityTypeProto entityType) { |
|||
return entityTypeByProtoNumber[entityType.getNumber()]; |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,24 @@ |
|||
/** |
|||
* Copyright © 2016-2023 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.queue.ruleengine; |
|||
|
|||
import java.io.Serializable; |
|||
|
|||
public enum QueueEvent implements Serializable { |
|||
|
|||
PARTITION_CHANGE, CONFIG_UPDATE, DELETE |
|||
|
|||
} |
|||
@ -0,0 +1,47 @@ |
|||
/** |
|||
* Copyright © 2016-2023 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.queue.ruleengine; |
|||
|
|||
import lombok.Getter; |
|||
import lombok.ToString; |
|||
import org.thingsboard.server.common.data.queue.Queue; |
|||
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; |
|||
|
|||
import java.util.Set; |
|||
|
|||
@Getter |
|||
@ToString |
|||
public class TbQueueConsumerManagerTask { |
|||
|
|||
private final QueueEvent event; |
|||
private Queue queue; |
|||
private Set<TopicPartitionInfo> partitions; |
|||
|
|||
public TbQueueConsumerManagerTask(QueueEvent event) { |
|||
this.event = event; |
|||
} |
|||
|
|||
public TbQueueConsumerManagerTask(QueueEvent event, Queue queue) { |
|||
this.event = event; |
|||
this.queue = queue; |
|||
} |
|||
|
|||
public TbQueueConsumerManagerTask(QueueEvent event, Set<TopicPartitionInfo> partitions) { |
|||
this.event = event; |
|||
this.partitions = partitions; |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,70 @@ |
|||
/** |
|||
* Copyright © 2016-2023 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.queue.ruleengine; |
|||
|
|||
import lombok.Getter; |
|||
import lombok.RequiredArgsConstructor; |
|||
import lombok.Setter; |
|||
import lombok.extern.slf4j.Slf4j; |
|||
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; |
|||
import org.thingsboard.server.gen.transport.TransportProtos; |
|||
import org.thingsboard.server.queue.TbQueueConsumer; |
|||
import org.thingsboard.server.queue.common.TbProtoQueueMsg; |
|||
|
|||
import java.util.Set; |
|||
import java.util.concurrent.Future; |
|||
import java.util.concurrent.TimeUnit; |
|||
|
|||
@RequiredArgsConstructor |
|||
@Slf4j |
|||
public class TbQueueConsumerTask { |
|||
|
|||
@Getter |
|||
private final Object key; |
|||
@Getter |
|||
private final TbQueueConsumer<TbProtoQueueMsg<TransportProtos.ToRuleEngineMsg>> consumer; |
|||
|
|||
@Setter |
|||
private Future<?> task; |
|||
|
|||
public void subscribe(Set<TopicPartitionInfo> partitions) { |
|||
log.trace("[{}] Subscribing to partitions: {}", key, partitions); |
|||
consumer.subscribe(partitions); |
|||
} |
|||
|
|||
public void initiateStop() { |
|||
log.debug("[{}] Initiating stop", key); |
|||
consumer.stop(); |
|||
} |
|||
|
|||
public void awaitCompletion() { |
|||
log.trace("[{}] Awaiting finish", key); |
|||
if (isRunning()) { |
|||
try { |
|||
task.get(30, TimeUnit.SECONDS); |
|||
log.trace("[{}] Awaited finish", key); |
|||
} catch (Exception e) { |
|||
log.warn("[{}] Failed to await for consumer to stop", key, e); |
|||
} |
|||
task = null; |
|||
} |
|||
} |
|||
|
|||
public boolean isRunning() { |
|||
return task != null; |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,89 @@ |
|||
/** |
|||
* Copyright © 2016-2023 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.queue.ruleengine; |
|||
|
|||
import lombok.Data; |
|||
import lombok.extern.slf4j.Slf4j; |
|||
import org.springframework.beans.factory.annotation.Value; |
|||
import org.springframework.stereotype.Component; |
|||
import org.thingsboard.common.util.ThingsBoardExecutors; |
|||
import org.thingsboard.common.util.ThingsBoardThreadFactory; |
|||
import org.thingsboard.server.actors.ActorSystemContext; |
|||
import org.thingsboard.server.common.stats.StatsFactory; |
|||
import org.thingsboard.server.queue.TbQueueAdmin; |
|||
import org.thingsboard.server.queue.discovery.PartitionService; |
|||
import org.thingsboard.server.queue.discovery.TbServiceInfoProvider; |
|||
import org.thingsboard.server.queue.provider.TbQueueProducerProvider; |
|||
import org.thingsboard.server.queue.provider.TbRuleEngineQueueFactory; |
|||
import org.thingsboard.server.queue.util.TbRuleEngineComponent; |
|||
import org.thingsboard.server.service.queue.processing.TbRuleEngineProcessingStrategyFactory; |
|||
import org.thingsboard.server.service.queue.processing.TbRuleEngineSubmitStrategyFactory; |
|||
import org.thingsboard.server.service.stats.RuleEngineStatisticsService; |
|||
|
|||
import javax.annotation.PostConstruct; |
|||
import java.util.concurrent.ExecutorService; |
|||
import java.util.concurrent.Executors; |
|||
import java.util.concurrent.ScheduledExecutorService; |
|||
|
|||
@Component |
|||
@TbRuleEngineComponent |
|||
@Slf4j |
|||
@Data |
|||
public class TbRuleEngineConsumerContext { |
|||
|
|||
@Value("${queue.rule-engine.poll-interval}") |
|||
private long pollDuration; |
|||
@Value("${queue.rule-engine.pack-processing-timeout}") |
|||
private long packProcessingTimeout; |
|||
@Value("${queue.rule-engine.stats.enabled:true}") |
|||
private boolean statsEnabled; |
|||
@Value("${queue.rule-engine.prometheus-stats.enabled:false}") |
|||
private boolean prometheusStatsEnabled; |
|||
@Value("${queue.rule-engine.topic-deletion-delay:15}") |
|||
private int topicDeletionDelayInSec; |
|||
@Value("${queue.rule-engine.management-thread-pool-size:12}") |
|||
private int mgmtThreadPoolSize; |
|||
|
|||
private final ActorSystemContext actorContext; |
|||
private final StatsFactory statsFactory; |
|||
private final TbRuleEngineSubmitStrategyFactory submitStrategyFactory; |
|||
private final TbRuleEngineProcessingStrategyFactory processingStrategyFactory; |
|||
private final TbRuleEngineQueueFactory queueFactory; |
|||
private final RuleEngineStatisticsService statisticsService; |
|||
private final TbServiceInfoProvider serviceInfoProvider; |
|||
private final PartitionService partitionService; |
|||
private final TbQueueProducerProvider producerProvider; |
|||
private final TbQueueAdmin queueAdmin; |
|||
|
|||
private ExecutorService consumersExecutor; |
|||
private ExecutorService mgmtExecutor; |
|||
private ScheduledExecutorService scheduler; |
|||
|
|||
private volatile boolean isReady = false; |
|||
|
|||
@PostConstruct |
|||
void init() { |
|||
this.consumersExecutor = Executors.newCachedThreadPool(ThingsBoardThreadFactory.forName("tb-rule-engine-consumer")); |
|||
this.mgmtExecutor = ThingsBoardExecutors.newWorkStealingPool(mgmtThreadPoolSize, "tb-rule-engine-mgmt"); |
|||
this.scheduler = Executors.newSingleThreadScheduledExecutor(ThingsBoardThreadFactory.forName("tb-rule-engine-consumer-scheduler")); |
|||
} |
|||
|
|||
public void stop() { |
|||
scheduler.shutdownNow(); |
|||
consumersExecutor.shutdownNow(); |
|||
mgmtExecutor.shutdownNow(); |
|||
} |
|||
} |
|||
@ -0,0 +1,486 @@ |
|||
/** |
|||
* Copyright © 2016-2023 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.queue.ruleengine; |
|||
|
|||
import com.google.protobuf.ProtocolStringList; |
|||
import lombok.Getter; |
|||
import lombok.extern.slf4j.Slf4j; |
|||
import org.thingsboard.common.util.ThingsBoardThreadFactory; |
|||
import org.thingsboard.server.common.data.id.EntityId; |
|||
import org.thingsboard.server.common.data.id.EntityIdFactory; |
|||
import org.thingsboard.server.common.data.id.TenantId; |
|||
import org.thingsboard.server.common.data.queue.Queue; |
|||
import org.thingsboard.server.common.msg.TbMsg; |
|||
import org.thingsboard.server.common.msg.gen.MsgProtos; |
|||
import org.thingsboard.server.common.msg.queue.QueueToRuleEngineMsg; |
|||
import org.thingsboard.server.common.msg.queue.RuleEngineException; |
|||
import org.thingsboard.server.common.msg.queue.RuleNodeInfo; |
|||
import org.thingsboard.server.common.msg.queue.ServiceType; |
|||
import org.thingsboard.server.common.msg.queue.TbMsgCallback; |
|||
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; |
|||
import org.thingsboard.server.gen.transport.TransportProtos.ToRuleEngineMsg; |
|||
import org.thingsboard.server.queue.TbQueueConsumer; |
|||
import org.thingsboard.server.queue.common.TbProtoQueueMsg; |
|||
import org.thingsboard.server.queue.discovery.QueueKey; |
|||
import org.thingsboard.server.service.queue.TbMsgPackCallback; |
|||
import org.thingsboard.server.service.queue.TbMsgPackProcessingContext; |
|||
import org.thingsboard.server.service.queue.TbRuleEngineConsumerStats; |
|||
import org.thingsboard.server.service.queue.processing.TbRuleEngineProcessingDecision; |
|||
import org.thingsboard.server.service.queue.processing.TbRuleEngineProcessingResult; |
|||
import org.thingsboard.server.service.queue.processing.TbRuleEngineProcessingStrategy; |
|||
import org.thingsboard.server.service.queue.processing.TbRuleEngineSubmitStrategy; |
|||
|
|||
import java.util.Collection; |
|||
import java.util.Collections; |
|||
import java.util.HashMap; |
|||
import java.util.HashSet; |
|||
import java.util.List; |
|||
import java.util.Map; |
|||
import java.util.Set; |
|||
import java.util.UUID; |
|||
import java.util.concurrent.ConcurrentLinkedQueue; |
|||
import java.util.concurrent.Future; |
|||
import java.util.concurrent.TimeUnit; |
|||
import java.util.concurrent.locks.ReentrantLock; |
|||
import java.util.stream.Collectors; |
|||
|
|||
@Slf4j |
|||
public class TbRuleEngineQueueConsumerManager { |
|||
|
|||
public static final String SUCCESSFUL_STATUS = "successful"; |
|||
public static final String FAILED_STATUS = "failed"; |
|||
|
|||
private final TbRuleEngineConsumerContext ctx; |
|||
private final QueueKey queueKey; |
|||
private final TbRuleEngineConsumerStats stats; |
|||
private final ReentrantLock lock = new ReentrantLock(); //NonfairSync
|
|||
|
|||
@Getter |
|||
private volatile Queue queue; |
|||
@Getter |
|||
private volatile Set<TopicPartitionInfo> partitions; |
|||
private volatile ConsumerWrapper consumerWrapper; |
|||
|
|||
private volatile boolean stopped; |
|||
|
|||
private final java.util.Queue<TbQueueConsumerManagerTask> tasks = new ConcurrentLinkedQueue<>(); |
|||
|
|||
public TbRuleEngineQueueConsumerManager(TbRuleEngineConsumerContext ctx, QueueKey queueKey) { |
|||
this.ctx = ctx; |
|||
this.queueKey = queueKey; |
|||
this.stats = new TbRuleEngineConsumerStats(queueKey, ctx.getStatsFactory()); |
|||
} |
|||
|
|||
public void init(Queue queue) { |
|||
this.queue = queue; |
|||
if (queue.isConsumerPerPartition()) { |
|||
this.consumerWrapper = new ConsumerPerPartitionWrapper(); |
|||
} else { |
|||
this.consumerWrapper = new SingleConsumerWrapper(); |
|||
} |
|||
log.debug("[{}] Initialized consumer for queue: {}", queueKey, queue); |
|||
} |
|||
|
|||
public void update(Queue queue) { |
|||
addTask(new TbQueueConsumerManagerTask(QueueEvent.CONFIG_UPDATE, queue)); |
|||
} |
|||
|
|||
public void update(Set<TopicPartitionInfo> partitions) { |
|||
addTask(new TbQueueConsumerManagerTask(QueueEvent.PARTITION_CHANGE, partitions)); |
|||
} |
|||
|
|||
public void delete() { |
|||
addTask(new TbQueueConsumerManagerTask(QueueEvent.DELETE)); |
|||
} |
|||
|
|||
private void addTask(TbQueueConsumerManagerTask todo) { |
|||
if (stopped) { |
|||
return; |
|||
} |
|||
tasks.add(todo); |
|||
log.trace("[{}] Added task: {}", queueKey, todo); |
|||
tryProcessTasks(); |
|||
} |
|||
|
|||
private void tryProcessTasks() { |
|||
if (!ctx.isReady()) { |
|||
log.debug("[{}] TbRuleEngineConsumerContext is not ready yet, will process tasks later", queueKey); |
|||
ctx.getScheduler().schedule(this::tryProcessTasks, 1, TimeUnit.SECONDS); |
|||
return; |
|||
} |
|||
ctx.getMgmtExecutor().submit(() -> { |
|||
if (lock.tryLock()) { |
|||
try { |
|||
Queue newConfiguration = null; |
|||
Set<TopicPartitionInfo> newPartitions = null; |
|||
while (!stopped) { |
|||
TbQueueConsumerManagerTask task = tasks.poll(); |
|||
if (task == null) { |
|||
break; |
|||
} |
|||
log.trace("[{}] Processing task: {}", queueKey, task); |
|||
|
|||
if (task.getEvent() == QueueEvent.PARTITION_CHANGE) { |
|||
newPartitions = task.getPartitions(); |
|||
} else if (task.getEvent() == QueueEvent.CONFIG_UPDATE) { |
|||
newConfiguration = task.getQueue(); |
|||
} else if (task.getEvent() == QueueEvent.DELETE) { |
|||
doDelete(); |
|||
return; |
|||
} |
|||
} |
|||
if (stopped) { |
|||
return; |
|||
} |
|||
if (newConfiguration != null) { |
|||
doUpdate(newConfiguration); |
|||
} |
|||
if (newPartitions != null) { |
|||
doUpdate(newPartitions); |
|||
} |
|||
} catch (Exception e) { |
|||
log.error("[{}] Failed to process tasks", queueKey, e); |
|||
} finally { |
|||
lock.unlock(); |
|||
} |
|||
} else { |
|||
log.trace("[{}] Failed to acquire lock", queueKey); |
|||
ctx.getScheduler().schedule(this::tryProcessTasks, 1, TimeUnit.SECONDS); |
|||
} |
|||
}); |
|||
} |
|||
|
|||
private void doUpdate(Queue newQueue) { |
|||
log.info("[{}] Processing queue update: {}", queueKey, newQueue); |
|||
var oldQueue = this.queue; |
|||
this.queue = newQueue; |
|||
if (log.isTraceEnabled()) { |
|||
log.trace("[{}] Old queue configuration: {}", queueKey, oldQueue); |
|||
log.trace("[{}] New queue configuration: {}", queueKey, newQueue); |
|||
} |
|||
|
|||
if (oldQueue == null) { |
|||
init(queue); |
|||
} else if (newQueue.isConsumerPerPartition() != oldQueue.isConsumerPerPartition()) { |
|||
consumerWrapper.getConsumers().forEach(TbQueueConsumerTask::initiateStop); |
|||
consumerWrapper.getConsumers().forEach(TbQueueConsumerTask::awaitCompletion); |
|||
|
|||
init(queue); |
|||
if (partitions != null) { |
|||
doUpdate(partitions); // even if partitions number was changed, there can be no partition change event
|
|||
} |
|||
} else { |
|||
// do nothing, because partitions change (if they changed) will be handled on PartitionChangeEvent,
|
|||
// and changes to pollInterval/packProcessingTimeout/submitStrategy/processingStrategy will be picked up by consumer on the fly,
|
|||
// and queue topic and name are immutable
|
|||
} |
|||
} |
|||
|
|||
private void doUpdate(Set<TopicPartitionInfo> partitions) { |
|||
this.partitions = partitions; |
|||
consumerWrapper.updatePartitions(partitions); |
|||
} |
|||
|
|||
public void stop() { |
|||
log.debug("[{}] Stopping consumers", queueKey); |
|||
consumerWrapper.getConsumers().forEach(TbQueueConsumerTask::initiateStop); |
|||
stopped = true; |
|||
} |
|||
|
|||
public void awaitStop() { |
|||
consumerWrapper.getConsumers().forEach(TbQueueConsumerTask::awaitCompletion); |
|||
log.debug("[{}] Unsubscribed and stopped consumers", queueKey); |
|||
} |
|||
|
|||
private void doDelete() { |
|||
stopped = true; |
|||
log.info("[{}] Handling queue deletion", queueKey); |
|||
consumerWrapper.getConsumers().forEach(TbQueueConsumerTask::awaitCompletion); |
|||
|
|||
List<TbQueueConsumer<TbProtoQueueMsg<ToRuleEngineMsg>>> queueConsumers = consumerWrapper.getConsumers().stream() |
|||
.map(TbQueueConsumerTask::getConsumer).collect(Collectors.toList()); |
|||
ctx.getConsumersExecutor().submit(() -> { |
|||
drainQueue(queueConsumers); |
|||
|
|||
queueConsumers.forEach(consumer -> { |
|||
for (String topic : consumer.getFullTopicNames()) { |
|||
try { |
|||
ctx.getQueueAdmin().deleteTopic(topic); |
|||
log.info("Deleted topic {}", topic); |
|||
} catch (Exception e) { |
|||
log.error("Failed to delete topic {}", topic, e); |
|||
} |
|||
} |
|||
try { |
|||
consumer.unsubscribe(); |
|||
} catch (Exception e) { |
|||
log.error("[{}] Failed to unsubscribe consumer", queueKey, e); |
|||
} |
|||
}); |
|||
}); |
|||
} |
|||
|
|||
private void launchConsumer(TbQueueConsumerTask consumerTask) { |
|||
log.info("[{}] Launching consumer", consumerTask.getKey()); |
|||
Future<?> consumerLoop = ctx.getConsumersExecutor().submit(() -> { |
|||
ThingsBoardThreadFactory.updateCurrentThreadName(consumerTask.getKey().toString()); |
|||
try { |
|||
consumerLoop(consumerTask.getConsumer()); |
|||
} catch (Throwable e) { |
|||
log.error("Failure in consumer loop", e); |
|||
} |
|||
}); |
|||
consumerTask.setTask(consumerLoop); |
|||
} |
|||
|
|||
private void consumerLoop(TbQueueConsumer<TbProtoQueueMsg<ToRuleEngineMsg>> consumer) { |
|||
while (!stopped && !consumer.isStopped()) { |
|||
try { |
|||
List<TbProtoQueueMsg<ToRuleEngineMsg>> msgs = consumer.poll(queue.getPollInterval()); |
|||
if (msgs.isEmpty()) { |
|||
continue; |
|||
} |
|||
processMsgs(msgs, consumer, queue); |
|||
} catch (Exception e) { |
|||
if (!consumer.isStopped()) { |
|||
log.warn("Failed to process messages from queue", e); |
|||
try { |
|||
Thread.sleep(ctx.getPollDuration()); |
|||
} catch (InterruptedException e2) { |
|||
log.trace("Failed to wait until the server has capacity to handle new requests", e2); |
|||
} |
|||
} |
|||
} |
|||
} |
|||
if (consumer.isStopped()) { |
|||
consumer.unsubscribe(); |
|||
} |
|||
log.info("Rule Engine consumer stopped"); |
|||
} |
|||
|
|||
private void processMsgs(List<TbProtoQueueMsg<ToRuleEngineMsg>> msgs, |
|||
TbQueueConsumer<TbProtoQueueMsg<ToRuleEngineMsg>> consumer, |
|||
Queue queue) throws InterruptedException { |
|||
TbRuleEngineSubmitStrategy submitStrategy = getSubmitStrategy(queue); |
|||
TbRuleEngineProcessingStrategy ackStrategy = getProcessingStrategy(queue); |
|||
submitStrategy.init(msgs); |
|||
while (!stopped && !consumer.isStopped()) { |
|||
TbMsgPackProcessingContext packCtx = new TbMsgPackProcessingContext(queue.getName(), submitStrategy, ackStrategy.isSkipTimeoutMsgs()); |
|||
submitStrategy.submitAttempt((id, msg) -> submitMessage(packCtx, id, msg)); |
|||
|
|||
final boolean timeout = !packCtx.await(queue.getPackProcessingTimeout(), TimeUnit.MILLISECONDS); |
|||
|
|||
TbRuleEngineProcessingResult result = new TbRuleEngineProcessingResult(queue.getName(), timeout, packCtx); |
|||
if (timeout) { |
|||
printFirstOrAll(packCtx, packCtx.getPendingMap(), "Timeout"); |
|||
} |
|||
if (!packCtx.getFailedMap().isEmpty()) { |
|||
printFirstOrAll(packCtx, packCtx.getFailedMap(), "Failed"); |
|||
} |
|||
packCtx.printProfilerStats(); |
|||
|
|||
TbRuleEngineProcessingDecision decision = ackStrategy.analyze(result); |
|||
if (ctx.isStatsEnabled()) { |
|||
stats.log(result, decision.isCommit()); |
|||
} |
|||
|
|||
packCtx.cleanup(); |
|||
|
|||
if (decision.isCommit()) { |
|||
submitStrategy.stop(); |
|||
consumer.commit(); |
|||
break; |
|||
} else { |
|||
submitStrategy.update(decision.getReprocessMap()); |
|||
} |
|||
} |
|||
} |
|||
|
|||
private TbRuleEngineSubmitStrategy getSubmitStrategy(Queue queue) { |
|||
return ctx.getSubmitStrategyFactory().newInstance(queue.getName(), queue.getSubmitStrategy()); |
|||
} |
|||
|
|||
private TbRuleEngineProcessingStrategy getProcessingStrategy(Queue queue) { |
|||
return ctx.getProcessingStrategyFactory().newInstance(queue.getName(), queue.getProcessingStrategy()); |
|||
} |
|||
|
|||
private void submitMessage(TbMsgPackProcessingContext packCtx, UUID id, TbProtoQueueMsg<ToRuleEngineMsg> msg) { |
|||
log.trace("[{}] Creating callback for topic {} message: {}", id, queue.getName(), msg.getValue()); |
|||
ToRuleEngineMsg toRuleEngineMsg = msg.getValue(); |
|||
TenantId tenantId = TenantId.fromUUID(new UUID(toRuleEngineMsg.getTenantIdMSB(), toRuleEngineMsg.getTenantIdLSB())); |
|||
TbMsgCallback callback = ctx.isPrometheusStatsEnabled() ? |
|||
new TbMsgPackCallback(id, tenantId, packCtx, stats.getTimer(tenantId, SUCCESSFUL_STATUS), stats.getTimer(tenantId, FAILED_STATUS)) : |
|||
new TbMsgPackCallback(id, tenantId, packCtx); |
|||
try { |
|||
if (!toRuleEngineMsg.getTbMsg().isEmpty()) { |
|||
forwardToRuleEngineActor(queue.getName(), tenantId, toRuleEngineMsg, callback); |
|||
} else { |
|||
callback.onSuccess(); |
|||
} |
|||
} catch (Exception e) { |
|||
callback.onFailure(new RuleEngineException(e.getMessage(), e)); |
|||
} |
|||
} |
|||
|
|||
private void forwardToRuleEngineActor(String queueName, TenantId tenantId, ToRuleEngineMsg toRuleEngineMsg, TbMsgCallback callback) { |
|||
TbMsg tbMsg = TbMsg.fromBytes(queueName, toRuleEngineMsg.getTbMsg().toByteArray(), callback); |
|||
QueueToRuleEngineMsg msg; |
|||
ProtocolStringList relationTypesList = toRuleEngineMsg.getRelationTypesList(); |
|||
Set<String> relationTypes; |
|||
if (relationTypesList.size() == 1) { |
|||
relationTypes = Collections.singleton(relationTypesList.get(0)); |
|||
} else { |
|||
relationTypes = new HashSet<>(relationTypesList); |
|||
} |
|||
msg = new QueueToRuleEngineMsg(tenantId, tbMsg, relationTypes, toRuleEngineMsg.getFailureMessage()); |
|||
ctx.getActorContext().tell(msg); |
|||
} |
|||
|
|||
private void printFirstOrAll(TbMsgPackProcessingContext ctx, Map<UUID, TbProtoQueueMsg<ToRuleEngineMsg>> map, String prefix) { |
|||
boolean printAll = log.isTraceEnabled(); |
|||
log.info("[{}] {} to process [{}] messages", queueKey, prefix, map.size()); |
|||
for (Map.Entry<UUID, TbProtoQueueMsg<ToRuleEngineMsg>> pending : map.entrySet()) { |
|||
ToRuleEngineMsg tmp = pending.getValue().getValue(); |
|||
TbMsg tmpMsg = TbMsg.fromBytes(queue.getName(), tmp.getTbMsg().toByteArray(), TbMsgCallback.EMPTY); |
|||
RuleNodeInfo ruleNodeInfo = ctx.getLastVisitedRuleNode(pending.getKey()); |
|||
if (printAll) { |
|||
log.trace("[{}][{}] {} to process message: {}, Last Rule Node: {}", queueKey, TenantId.fromUUID(new UUID(tmp.getTenantIdMSB(), tmp.getTenantIdLSB())), prefix, tmpMsg, ruleNodeInfo); |
|||
} else { |
|||
log.info("[{}] {} to process message: {}, Last Rule Node: {}", TenantId.fromUUID(new UUID(tmp.getTenantIdMSB(), tmp.getTenantIdLSB())), prefix, tmpMsg, ruleNodeInfo); |
|||
break; |
|||
} |
|||
} |
|||
} |
|||
|
|||
public void printStats(long ts) { |
|||
stats.printStats(); |
|||
ctx.getStatisticsService().reportQueueStats(ts, stats); |
|||
stats.reset(); |
|||
} |
|||
|
|||
private void drainQueue(List<TbQueueConsumer<TbProtoQueueMsg<ToRuleEngineMsg>>> consumers) { |
|||
long finishTs = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(ctx.getTopicDeletionDelayInSec()); |
|||
try { |
|||
int n = 0; |
|||
while (System.currentTimeMillis() <= finishTs) { |
|||
for (TbQueueConsumer<TbProtoQueueMsg<ToRuleEngineMsg>> consumer : consumers) { |
|||
List<TbProtoQueueMsg<ToRuleEngineMsg>> msgs = consumer.poll(queue.getPollInterval()); |
|||
if (msgs.isEmpty()) { |
|||
continue; |
|||
} |
|||
for (TbProtoQueueMsg<ToRuleEngineMsg> msg : msgs) { |
|||
try { |
|||
MsgProtos.TbMsgProto tbMsgProto = MsgProtos.TbMsgProto.parseFrom(msg.getValue().getTbMsg().toByteArray()); |
|||
EntityId originator = EntityIdFactory.getByTypeAndUuid(tbMsgProto.getEntityType(), new UUID(tbMsgProto.getEntityIdMSB(), tbMsgProto.getEntityIdLSB())); |
|||
|
|||
TopicPartitionInfo tpi = ctx.getPartitionService().resolve(ServiceType.TB_RULE_ENGINE, queue.getName(), TenantId.SYS_TENANT_ID, originator); |
|||
ctx.getProducerProvider().getRuleEngineMsgProducer().send(tpi, msg, null); |
|||
n++; |
|||
} catch (Throwable e) { |
|||
log.warn("Failed to move message to system {}: {}", consumer.getTopic(), msg, e); |
|||
} |
|||
} |
|||
consumer.commit(); |
|||
} |
|||
} |
|||
if (n > 0) { |
|||
log.info("Moved {} messages from {} to system {}", n, queueKey, queue.getName()); |
|||
} |
|||
} catch (Exception e) { |
|||
log.error("[{}] Failed to drain queue", queueKey, e); |
|||
} |
|||
} |
|||
|
|||
private static String partitionsToString(Collection<TopicPartitionInfo> partitions) { |
|||
return partitions.stream().map(TopicPartitionInfo::getFullTopicName).collect(Collectors.joining(", ", "[", "]")); |
|||
} |
|||
|
|||
interface ConsumerWrapper { |
|||
|
|||
void updatePartitions(Set<TopicPartitionInfo> partitions); |
|||
|
|||
Collection<TbQueueConsumerTask> getConsumers(); |
|||
|
|||
} |
|||
|
|||
class ConsumerPerPartitionWrapper implements ConsumerWrapper { |
|||
private final Map<TopicPartitionInfo, TbQueueConsumerTask> consumers = new HashMap<>(); |
|||
|
|||
@Override |
|||
public void updatePartitions(Set<TopicPartitionInfo> partitions) { |
|||
Set<TopicPartitionInfo> addedPartitions = new HashSet<>(partitions); |
|||
addedPartitions.removeAll(consumers.keySet()); |
|||
|
|||
Set<TopicPartitionInfo> removedPartitions = new HashSet<>(consumers.keySet()); |
|||
removedPartitions.removeAll(partitions); |
|||
log.info("[{}] Added partitions: {}, removed partitions: {}", queueKey, partitionsToString(addedPartitions), partitionsToString(removedPartitions)); |
|||
|
|||
removedPartitions.forEach((tpi) -> { |
|||
consumers.get(tpi).initiateStop(); |
|||
}); |
|||
removedPartitions.forEach((tpi) -> { |
|||
consumers.remove(tpi).awaitCompletion(); |
|||
}); |
|||
|
|||
addedPartitions.forEach((tpi) -> { |
|||
String key = queueKey + "-" + tpi.getPartition().orElse(-999999); |
|||
TbQueueConsumerTask consumer = new TbQueueConsumerTask(key, ctx.getQueueFactory().createToRuleEngineMsgConsumer(queue)); |
|||
consumers.put(tpi, consumer); |
|||
consumer.subscribe(Set.of(tpi)); |
|||
launchConsumer(consumer); |
|||
}); |
|||
} |
|||
|
|||
@Override |
|||
public Collection<TbQueueConsumerTask> getConsumers() { |
|||
return consumers.values(); |
|||
} |
|||
} |
|||
|
|||
class SingleConsumerWrapper implements ConsumerWrapper { |
|||
private TbQueueConsumerTask consumer; |
|||
|
|||
@Override |
|||
public void updatePartitions(Set<TopicPartitionInfo> partitions) { |
|||
log.info("[{}] New partitions: {}", queueKey, partitionsToString(partitions)); |
|||
if (partitions.isEmpty()) { |
|||
if (consumer != null && consumer.isRunning()) { |
|||
consumer.initiateStop(); |
|||
consumer.awaitCompletion(); |
|||
} |
|||
consumer = null; |
|||
return; |
|||
} |
|||
|
|||
if (consumer == null) { |
|||
consumer = new TbQueueConsumerTask(queueKey, ctx.getQueueFactory().createToRuleEngineMsgConsumer(queue)); |
|||
} |
|||
consumer.subscribe(partitions); |
|||
if (!consumer.isRunning()) { |
|||
launchConsumer(consumer); |
|||
} |
|||
} |
|||
|
|||
@Override |
|||
public Collection<TbQueueConsumerTask> getConsumers() { |
|||
if (consumer == null) { |
|||
return Collections.emptyList(); |
|||
} |
|||
return List.of(consumer); |
|||
} |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,49 @@ |
|||
/** |
|||
* Copyright © 2016-2023 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.queue; |
|||
|
|||
import org.junit.jupiter.api.Test; |
|||
import org.thingsboard.server.common.data.EntityType; |
|||
import org.thingsboard.server.common.data.id.EntityId; |
|||
import org.thingsboard.server.common.data.id.RuleChainId; |
|||
import org.thingsboard.server.common.data.id.TenantId; |
|||
import org.thingsboard.server.common.data.plugin.ComponentLifecycleEvent; |
|||
import org.thingsboard.server.common.msg.plugin.ComponentLifecycleMsg; |
|||
import org.thingsboard.server.gen.transport.TransportProtos; |
|||
import java.util.UUID; |
|||
|
|||
import static org.assertj.core.api.Assertions.assertThat; |
|||
|
|||
class ProtoUtilsTest { |
|||
|
|||
TenantId tenantId = TenantId.fromUUID(UUID.fromString("35e10f77-16e7-424d-ae46-ee780f87ac4f")); |
|||
EntityId entityId = new RuleChainId(UUID.fromString("c640b635-4f0f-41e6-b10b-25a86003094e")); |
|||
@Test |
|||
void protoComponentLifecycleSerialization() { |
|||
ComponentLifecycleMsg msg = new ComponentLifecycleMsg(tenantId, entityId, ComponentLifecycleEvent.UPDATED); |
|||
assertThat(ProtoUtils.fromProto(ProtoUtils.toProto(msg))).as("deserialized").isEqualTo(msg); |
|||
msg = new ComponentLifecycleMsg(tenantId, entityId, ComponentLifecycleEvent.STARTED); |
|||
assertThat(ProtoUtils.fromProto(ProtoUtils.toProto(msg))).as("deserialized").isEqualTo(msg); |
|||
} |
|||
|
|||
@Test |
|||
void protoEntityTypeSerialization() { |
|||
for(EntityType entityType : EntityType.values()){ |
|||
assertThat(ProtoUtils.fromProto(ProtoUtils.toProto(entityType))).as(entityType.getNormalName()).isEqualTo(entityType); |
|||
} |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,773 @@ |
|||
/** |
|||
* Copyright © 2016-2023 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.queue.ruleengine; |
|||
|
|||
import lombok.Getter; |
|||
import lombok.SneakyThrows; |
|||
import lombok.extern.slf4j.Slf4j; |
|||
import org.junit.After; |
|||
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.testcontainers.shaded.org.apache.commons.lang3.RandomUtils; |
|||
import org.thingsboard.common.util.JacksonUtil; |
|||
import org.thingsboard.server.actors.ActorSystemContext; |
|||
import org.thingsboard.server.common.data.id.DeviceId; |
|||
import org.thingsboard.server.common.data.id.QueueId; |
|||
import org.thingsboard.server.common.data.id.TenantId; |
|||
import org.thingsboard.server.common.data.msg.TbMsgType; |
|||
import org.thingsboard.server.common.data.queue.ProcessingStrategy; |
|||
import org.thingsboard.server.common.data.queue.ProcessingStrategyType; |
|||
import org.thingsboard.server.common.data.queue.Queue; |
|||
import org.thingsboard.server.common.data.queue.SubmitStrategy; |
|||
import org.thingsboard.server.common.data.queue.SubmitStrategyType; |
|||
import org.thingsboard.server.common.msg.TbMsg; |
|||
import org.thingsboard.server.common.msg.TbMsgMetaData; |
|||
import org.thingsboard.server.common.msg.queue.QueueToRuleEngineMsg; |
|||
import org.thingsboard.server.common.msg.queue.ServiceType; |
|||
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; |
|||
import org.thingsboard.server.common.stats.StatsFactory; |
|||
import org.thingsboard.server.gen.transport.TransportProtos.ToRuleEngineMsg; |
|||
import org.thingsboard.server.queue.TbQueueAdmin; |
|||
import org.thingsboard.server.queue.TbQueueProducer; |
|||
import org.thingsboard.server.queue.common.AbstractTbQueueConsumerTemplate; |
|||
import org.thingsboard.server.queue.common.TbProtoQueueMsg; |
|||
import org.thingsboard.server.queue.discovery.PartitionService; |
|||
import org.thingsboard.server.queue.discovery.QueueKey; |
|||
import org.thingsboard.server.queue.discovery.TbServiceInfoProvider; |
|||
import org.thingsboard.server.queue.provider.TbQueueProducerProvider; |
|||
import org.thingsboard.server.queue.provider.TbRuleEngineQueueFactory; |
|||
import org.thingsboard.server.service.queue.processing.TbRuleEngineProcessingStrategyFactory; |
|||
import org.thingsboard.server.service.queue.processing.TbRuleEngineSubmitStrategyFactory; |
|||
import org.thingsboard.server.service.stats.RuleEngineStatisticsService; |
|||
|
|||
import java.io.IOException; |
|||
import java.util.Arrays; |
|||
import java.util.Collections; |
|||
import java.util.List; |
|||
import java.util.Objects; |
|||
import java.util.Set; |
|||
import java.util.UUID; |
|||
import java.util.concurrent.ConcurrentHashMap; |
|||
import java.util.concurrent.TimeUnit; |
|||
import java.util.concurrent.atomic.AtomicInteger; |
|||
import java.util.function.Supplier; |
|||
import java.util.stream.Collectors; |
|||
import java.util.stream.IntStream; |
|||
|
|||
import static org.assertj.core.api.Assertions.assertThat; |
|||
import static org.awaitility.Awaitility.await; |
|||
import static org.mockito.ArgumentMatchers.any; |
|||
import static org.mockito.ArgumentMatchers.argThat; |
|||
import static org.mockito.ArgumentMatchers.eq; |
|||
import static org.mockito.Mockito.after; |
|||
import static org.mockito.Mockito.atLeast; |
|||
import static org.mockito.Mockito.atLeastOnce; |
|||
import static org.mockito.Mockito.clearInvocations; |
|||
import static org.mockito.Mockito.doAnswer; |
|||
import static org.mockito.Mockito.mock; |
|||
import static org.mockito.Mockito.never; |
|||
import static org.mockito.Mockito.spy; |
|||
import static org.mockito.Mockito.times; |
|||
import static org.mockito.Mockito.verify; |
|||
import static org.mockito.Mockito.verifyNoInteractions; |
|||
import static org.mockito.Mockito.verifyNoMoreInteractions; |
|||
import static org.mockito.Mockito.when; |
|||
|
|||
@Slf4j |
|||
@RunWith(MockitoJUnitRunner.class) |
|||
public class TbRuleEngineQueueConsumerManagerTest { |
|||
|
|||
@Mock |
|||
private ActorSystemContext actorContext; |
|||
@Mock |
|||
private StatsFactory statsFactory; |
|||
@Mock |
|||
private TbRuleEngineQueueFactory queueFactory; |
|||
@Mock |
|||
private RuleEngineStatisticsService statisticsService; |
|||
@Mock |
|||
private TbServiceInfoProvider serviceInfoProvider; |
|||
@Mock |
|||
private PartitionService partitionService; |
|||
@Mock |
|||
private TbQueueProducerProvider producerProvider; |
|||
private TbQueueProducer<TbProtoQueueMsg<ToRuleEngineMsg>> ruleEngineMsgProducer; |
|||
@Mock |
|||
private TbQueueAdmin queueAdmin; |
|||
private TbRuleEngineConsumerContext ruleEngineConsumerContext; |
|||
|
|||
private TbRuleEngineQueueConsumerManager consumerManager; |
|||
private Queue queue; |
|||
|
|||
private Set<TestConsumer> consumers; |
|||
private boolean generateQueueMsgs; |
|||
private AtomicInteger totalConsumedMsgs; |
|||
private AtomicInteger totalProcessedMsgs; |
|||
|
|||
@Before |
|||
public void beforeEach() { |
|||
ruleEngineConsumerContext = new TbRuleEngineConsumerContext( |
|||
actorContext, statsFactory, spy(new TbRuleEngineSubmitStrategyFactory()), |
|||
spy(new TbRuleEngineProcessingStrategyFactory()), queueFactory, statisticsService, |
|||
serviceInfoProvider, partitionService, producerProvider, queueAdmin |
|||
); |
|||
consumers = ConcurrentHashMap.newKeySet(); |
|||
generateQueueMsgs = true; |
|||
totalConsumedMsgs = new AtomicInteger(); |
|||
totalProcessedMsgs = new AtomicInteger(); |
|||
doAnswer(inv -> { |
|||
QueueToRuleEngineMsg msg = inv.getArgument(0); |
|||
msg.getMsg().getCallback().onSuccess(); |
|||
totalProcessedMsgs.incrementAndGet(); |
|||
log.trace("totalProcessedMsgs = {}", totalProcessedMsgs); |
|||
return null; |
|||
}).when(actorContext).tell(any()); |
|||
ruleEngineMsgProducer = mock(TbQueueProducer.class); |
|||
when(producerProvider.getRuleEngineMsgProducer()).thenReturn(ruleEngineMsgProducer); |
|||
ruleEngineConsumerContext.setMgmtThreadPoolSize(2); |
|||
ruleEngineConsumerContext.setTopicDeletionDelayInSec(5); |
|||
ruleEngineConsumerContext.init(); |
|||
ruleEngineConsumerContext.setReady(false); |
|||
|
|||
queue = new Queue(); |
|||
queue.setName("Test"); |
|||
queue.setTenantId(TenantId.SYS_TENANT_ID); |
|||
queue.setId(new QueueId(UUID.randomUUID())); |
|||
queue.setTopic("tb_test"); |
|||
queue.setPartitions(10); |
|||
queue.setConsumerPerPartition(true); |
|||
queue.setPollInterval(250); |
|||
queue.setPackProcessingTimeout(2000); |
|||
SubmitStrategy submitStrategy = new SubmitStrategy(); |
|||
submitStrategy.setType(SubmitStrategyType.BURST); |
|||
submitStrategy.setBatchSize(200); |
|||
queue.setSubmitStrategy(submitStrategy); |
|||
ProcessingStrategy processingStrategy = new ProcessingStrategy(); |
|||
processingStrategy.setType(ProcessingStrategyType.SKIP_ALL_FAILURES_AND_TIMED_OUT); |
|||
processingStrategy.setRetries(0); |
|||
queue.setProcessingStrategy(processingStrategy); |
|||
|
|||
doAnswer(i -> { |
|||
TestConsumer consumer = spy(new TestConsumer(queue.getTopic())); |
|||
if (generateQueueMsgs) { |
|||
consumer.setUpTestMsg(); |
|||
} |
|||
consumers.add(consumer); |
|||
return consumer; |
|||
}).when(queueFactory).createToRuleEngineMsgConsumer(any()); |
|||
|
|||
QueueKey queueKey = new QueueKey(ServiceType.TB_RULE_ENGINE, queue); |
|||
consumerManager = new TbRuleEngineQueueConsumerManager(ruleEngineConsumerContext, queueKey); |
|||
} |
|||
|
|||
@After |
|||
public void afterEach() { |
|||
consumerManager.stop(); |
|||
consumerManager.awaitStop(); |
|||
ruleEngineConsumerContext.stop(); |
|||
|
|||
if (generateQueueMsgs) { |
|||
await().atMost(10, TimeUnit.SECONDS) |
|||
.untilAsserted(() -> { |
|||
log.debug("totalConsumedMsgs = {}, totalProcessedMsgs = {}", totalConsumedMsgs.get(), totalProcessedMsgs.get()); |
|||
assertThat(totalProcessedMsgs.get()).isEqualTo(totalConsumedMsgs.get()); |
|||
}); |
|||
} |
|||
} |
|||
|
|||
@Test |
|||
public void testInit_consumerPerPartition() { |
|||
queue.setConsumerPerPartition(true); |
|||
consumerManager.init(queue); |
|||
|
|||
Set<TopicPartitionInfo> partitions = createTpis(2, 3, 4); |
|||
consumerManager.update(partitions); |
|||
partitions = createTpis(3, 4, 5); |
|||
consumerManager.update(partitions); |
|||
partitions = createTpis(1, 2, 3); |
|||
consumerManager.update(partitions); |
|||
// simulated multiple partition change events before consumer is ready; only latest partitions should be processed
|
|||
verifyNoInteractions(queueFactory); |
|||
|
|||
ruleEngineConsumerContext.setReady(true); |
|||
await().atMost(2, TimeUnit.SECONDS) |
|||
.until(() -> consumers.size() == 3); |
|||
for (TopicPartitionInfo partition : partitions) { |
|||
TestConsumer consumer = getConsumer(partition); |
|||
verifySubscribedAndLaunched(consumer, Set.of(partition)); |
|||
} |
|||
} |
|||
|
|||
@Test |
|||
public void testInit_singleConsumer() { |
|||
queue.setConsumerPerPartition(false); |
|||
consumerManager.init(queue); |
|||
|
|||
Set<TopicPartitionInfo> partitions = createTpis(2, 3, 4); |
|||
consumerManager.update(partitions); |
|||
partitions = createTpis(3, 4, 5); |
|||
consumerManager.update(partitions); |
|||
partitions = createTpis(1, 2, 3); |
|||
consumerManager.update(partitions); |
|||
|
|||
verifyNoInteractions(queueFactory); |
|||
|
|||
ruleEngineConsumerContext.setReady(true); |
|||
await().atMost(2, TimeUnit.SECONDS) |
|||
.until(() -> consumers.size() == 1); |
|||
TestConsumer consumer = getConsumer(); |
|||
verifySubscribedAndLaunched(consumer, partitions); |
|||
} |
|||
|
|||
@Test |
|||
public void testPartitionsUpdate_singleConsumer() { |
|||
queue.setConsumerPerPartition(false); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
|
|||
Set<TopicPartitionInfo> partitions = Collections.emptySet(); |
|||
consumerManager.update(partitions); |
|||
verify(queueFactory, after(1000).never()).createToRuleEngineMsgConsumer(any()); |
|||
|
|||
partitions = createTpis(1); |
|||
consumerManager.update(partitions); |
|||
TestConsumer consumer = getConsumer(); |
|||
verifySubscribedAndLaunched(consumer, partitions); |
|||
|
|||
partitions = createTpis(1, 2, 3); |
|||
consumerManager.update(partitions); |
|||
verifySubscribedAndLaunched(consumer, partitions); |
|||
|
|||
partitions = createTpis(4, 5, 6); |
|||
consumerManager.update(partitions); |
|||
verifySubscribedAndLaunched(consumer, partitions); |
|||
|
|||
partitions = Collections.emptySet(); |
|||
consumerManager.update(partitions); |
|||
verifyUnsubscribedAndStopped(consumer); |
|||
|
|||
partitions = createTpis(1, 2, 3); |
|||
consumerManager.update(partitions); |
|||
consumer = getConsumer(); |
|||
verifySubscribedAndLaunched(consumer, partitions); |
|||
} |
|||
|
|||
@Test |
|||
public void testPartitionsUpdate_consumerPerPartition() { |
|||
queue.setConsumerPerPartition(true); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
|
|||
consumerManager.update(Collections.emptySet()); |
|||
verify(queueFactory, after(1000).never()).createToRuleEngineMsgConsumer(any()); |
|||
|
|||
consumerManager.update(createTpis(1)); |
|||
TestConsumer consumer1 = getConsumer(1); |
|||
verifySubscribedAndLaunched(consumer1, 1); |
|||
|
|||
consumerManager.update(createTpis(1, 2, 3)); |
|||
TestConsumer consumer2 = getConsumer(2); |
|||
TestConsumer consumer3 = getConsumer(3); |
|||
verifySubscribedAndLaunched(consumer2, 2); |
|||
verifySubscribedAndLaunched(consumer3, 3); |
|||
verifyNotTouched(consumer1); |
|||
|
|||
consumerManager.update(createTpis(3, 4, 5)); |
|||
TestConsumer consumer4 = getConsumer(4); |
|||
TestConsumer consumer5 = getConsumer(5); |
|||
verifySubscribedAndLaunched(consumer4, 4); |
|||
verifySubscribedAndLaunched(consumer5, 5); |
|||
verifyUnsubscribedAndStopped(consumer1); |
|||
verifyUnsubscribedAndStopped(consumer2); |
|||
verifyNotTouched(consumer3); |
|||
|
|||
consumerManager.update(Collections.emptySet()); |
|||
verifyUnsubscribedAndStopped(consumer3); |
|||
verifyUnsubscribedAndStopped(consumer4); |
|||
verifyUnsubscribedAndStopped(consumer5); |
|||
|
|||
consumerManager.update(createTpis(1, 2, 3)); |
|||
consumer1 = getConsumer(1); |
|||
consumer2 = getConsumer(2); |
|||
consumer3 = getConsumer(3); |
|||
verifySubscribedAndLaunched(consumer1, 1); |
|||
verifySubscribedAndLaunched(consumer2, 2); |
|||
verifySubscribedAndLaunched(consumer3, 3); |
|||
} |
|||
|
|||
@Test |
|||
public void testConfigUpdate_singleConsumer() { |
|||
queue.setConsumerPerPartition(false); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
Set<TopicPartitionInfo> partitions = createTpis(1, 2, 3); |
|||
consumerManager.update(partitions); |
|||
TestConsumer consumer = getConsumer(); |
|||
verifySubscribedAndLaunched(consumer, partitions); |
|||
|
|||
Queue newConfig = JacksonUtil.clone(queue); |
|||
newConfig.setPollInterval(queue.getPollInterval() / 2); |
|||
newConfig.setPartitions(queue.getPartitions() / 2); |
|||
newConfig.setPackProcessingTimeout(queue.getPackProcessingTimeout() * 2); |
|||
newConfig.getSubmitStrategy().setType(SubmitStrategyType.SEQUENTIAL_BY_ORIGINATOR); |
|||
newConfig.getProcessingStrategy().setType(ProcessingStrategyType.RETRY_ALL); |
|||
consumerManager.update(newConfig); |
|||
|
|||
await().atMost(2, TimeUnit.SECONDS) |
|||
.untilAsserted(() -> { |
|||
verify(consumer, atLeastOnce()).poll(eq((long) newConfig.getPollInterval())); |
|||
verify(ruleEngineConsumerContext.getSubmitStrategyFactory(), atLeastOnce()).newInstance(any(), eq(newConfig.getSubmitStrategy())); |
|||
verify(ruleEngineConsumerContext.getProcessingStrategyFactory(), atLeastOnce()).newInstance(any(), eq(newConfig.getProcessingStrategy())); |
|||
}); |
|||
} |
|||
|
|||
@Test |
|||
public void testConfigUpdate_consumerPerPartition() { |
|||
queue.setConsumerPerPartition(true); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
Set<TopicPartitionInfo> partitions = createTpis(1, 2, 3); |
|||
consumerManager.update(partitions); |
|||
TestConsumer consumer1 = getConsumer(1); |
|||
TestConsumer consumer2 = getConsumer(2); |
|||
TestConsumer consumer3 = getConsumer(3); |
|||
verifySubscribedAndLaunched(consumer1, 1); |
|||
verifySubscribedAndLaunched(consumer2, 2); |
|||
verifySubscribedAndLaunched(consumer3, 3); |
|||
|
|||
Queue newConfig = JacksonUtil.clone(queue); |
|||
newConfig.setPollInterval(queue.getPollInterval() / 2); |
|||
newConfig.setPartitions(queue.getPartitions() / 2); |
|||
newConfig.setPackProcessingTimeout(queue.getPackProcessingTimeout() * 2); |
|||
newConfig.getSubmitStrategy().setType(SubmitStrategyType.SEQUENTIAL_BY_ORIGINATOR); |
|||
newConfig.getProcessingStrategy().setType(ProcessingStrategyType.RETRY_ALL); |
|||
consumerManager.update(newConfig); |
|||
|
|||
await().atMost(2, TimeUnit.SECONDS) |
|||
.untilAsserted(() -> { |
|||
verify(consumer1, atLeastOnce()).poll(eq((long) newConfig.getPollInterval())); |
|||
verify(consumer2, atLeastOnce()).poll(eq((long) newConfig.getPollInterval())); |
|||
verify(consumer3, atLeastOnce()).poll(eq((long) newConfig.getPollInterval())); |
|||
}); |
|||
verifyNotTouched(consumer1); |
|||
verifyNotTouched(consumer2); |
|||
verifyNotTouched(consumer3); |
|||
} |
|||
|
|||
@Test |
|||
public void testConfigUpdate_fromSingleToConsumerPerPartition() { |
|||
queue.setConsumerPerPartition(false); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
Set<TopicPartitionInfo> partitions = createTpis(1, 2, 3); |
|||
consumerManager.update(partitions); |
|||
TestConsumer consumer = getConsumer(); |
|||
verifySubscribedAndLaunched(consumer, partitions); |
|||
|
|||
Queue newConfig = JacksonUtil.clone(queue); |
|||
newConfig.setConsumerPerPartition(true); |
|||
consumerManager.update(newConfig); |
|||
|
|||
verifyUnsubscribedAndStopped(consumer); |
|||
verifySubscribedAndLaunched(getConsumer(1), 1); |
|||
verifySubscribedAndLaunched(getConsumer(2), 2); |
|||
verifySubscribedAndLaunched(getConsumer(3), 3); |
|||
} |
|||
|
|||
@Test |
|||
public void testConfigUpdate_fromConsumerPerPartitionToSingle() { |
|||
queue.setConsumerPerPartition(true); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
Set<TopicPartitionInfo> partitions = createTpis(1, 2, 3); |
|||
consumerManager.update(partitions); |
|||
TestConsumer consumer1 = getConsumer(1); |
|||
TestConsumer consumer2 = getConsumer(2); |
|||
TestConsumer consumer3 = getConsumer(3); |
|||
verifySubscribedAndLaunched(consumer1, 1); |
|||
verifySubscribedAndLaunched(consumer2, 2); |
|||
verifySubscribedAndLaunched(consumer3, 3); |
|||
|
|||
Queue newConfig = JacksonUtil.clone(queue); |
|||
newConfig.setConsumerPerPartition(false); |
|||
consumerManager.update(newConfig); |
|||
|
|||
verifyUnsubscribedAndStopped(consumer1); |
|||
verifyUnsubscribedAndStopped(consumer2); |
|||
verifyUnsubscribedAndStopped(consumer3); |
|||
verifySubscribedAndLaunched(getConsumer(), partitions); |
|||
} |
|||
|
|||
@Test |
|||
public void testStop() { |
|||
queue.setConsumerPerPartition(true); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
consumerManager.update(createTpis(1)); |
|||
TestConsumer consumer = getConsumer(1); |
|||
verifySubscribedAndLaunched(consumer, 1); |
|||
verify(queueFactory, times(1)).createToRuleEngineMsgConsumer(any()); |
|||
|
|||
consumerManager.stop(); |
|||
consumerManager.update(createTpis(1, 2, 3, 4)); // to check that no new tasks after stop are processed
|
|||
consumerManager.update(createTpis(5, 6, 7)); |
|||
|
|||
verifyUnsubscribedAndStopped(consumer); |
|||
verifyNoMoreInteractions(queueFactory); |
|||
} |
|||
|
|||
@Test |
|||
public void testDelete_consumerPerPartition() { |
|||
queue.setConsumerPerPartition(true); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
Set<TopicPartitionInfo> partitions = createTpis(1, 2); |
|||
consumerManager.update(partitions); |
|||
TestConsumer consumer1 = getConsumer(1); |
|||
TestConsumer consumer2 = getConsumer(2); |
|||
verifySubscribedAndLaunched(consumer1, 1); |
|||
verifySubscribedAndLaunched(consumer2, 2); |
|||
verifyMsgProcessed(consumer1.testMsg); |
|||
verifyMsgProcessed(consumer2.testMsg); |
|||
|
|||
consumerManager.delete(); |
|||
|
|||
await().atMost(2, TimeUnit.SECONDS) |
|||
.untilAsserted(() -> { |
|||
verify(ruleEngineMsgProducer).send(any(), any(), any()); |
|||
}); |
|||
clearInvocations(actorContext); |
|||
verify(consumer1, never()).unsubscribe(); |
|||
verify(consumer2, never()).unsubscribe(); |
|||
int msgCount = totalConsumedMsgs.get(); |
|||
|
|||
await().atLeast(4, TimeUnit.SECONDS) // based on topicDeletionDelayInSec
|
|||
.atMost(7, TimeUnit.SECONDS) |
|||
.untilAsserted(() -> { |
|||
partitions.stream() |
|||
.map(TopicPartitionInfo::getFullTopicName) |
|||
.forEach(topic -> { |
|||
verify(queueAdmin).deleteTopic(eq(topic)); |
|||
}); |
|||
}); |
|||
verify(consumer1).unsubscribe(); |
|||
verify(consumer2).unsubscribe(); |
|||
|
|||
int totalMovedMsgs = totalConsumedMsgs.get() - msgCount; |
|||
assertThat(totalMovedMsgs).isNotZero(); |
|||
verify(ruleEngineMsgProducer, atLeast(totalMovedMsgs)).send(any(), any(), any()); |
|||
verify(actorContext, never()).tell(any()); |
|||
generateQueueMsgs = false; |
|||
} |
|||
|
|||
@Test |
|||
public void testDelete_singleConsumer() { |
|||
queue.setConsumerPerPartition(false); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
Set<TopicPartitionInfo> partitions = createTpis(1, 2); |
|||
consumerManager.update(partitions); |
|||
TestConsumer consumer = getConsumer(); |
|||
verifySubscribedAndLaunched(consumer, partitions); |
|||
verifyMsgProcessed(consumer.testMsg); |
|||
|
|||
consumerManager.delete(); |
|||
|
|||
await().atMost(2, TimeUnit.SECONDS) |
|||
.untilAsserted(() -> { |
|||
verify(ruleEngineMsgProducer).send(any(), any(), any()); |
|||
}); |
|||
clearInvocations(actorContext); |
|||
verify(consumer, never()).unsubscribe(); |
|||
int msgCount = totalConsumedMsgs.get(); |
|||
|
|||
await().atLeast(4, TimeUnit.SECONDS) |
|||
.atMost(7, TimeUnit.SECONDS) |
|||
.untilAsserted(() -> { |
|||
partitions.stream() |
|||
.map(TopicPartitionInfo::getFullTopicName) |
|||
.forEach(topic -> { |
|||
verify(queueAdmin).deleteTopic(eq(topic)); |
|||
}); |
|||
}); |
|||
verify(consumer).unsubscribe(); |
|||
|
|||
int movedMsgs = totalConsumedMsgs.get() - msgCount; |
|||
assertThat(movedMsgs).isNotZero(); |
|||
verify(ruleEngineMsgProducer, atLeast(movedMsgs)).send(any(), any(), any()); |
|||
verify(actorContext, never()).tell(any()); |
|||
generateQueueMsgs = false; |
|||
} |
|||
|
|||
@Test |
|||
public void testManyDifferentUpdates() throws Exception { |
|||
queue.setConsumerPerPartition(RandomUtils.nextBoolean()); |
|||
consumerManager.init(queue); |
|||
ruleEngineConsumerContext.setReady(true); |
|||
|
|||
Supplier<Queue> queueConfigUpdater = () -> { |
|||
Queue oldConfig = consumerManager.getQueue(); |
|||
Queue newConfig = JacksonUtil.clone(oldConfig); |
|||
newConfig.setConsumerPerPartition(RandomUtils.nextBoolean()); |
|||
newConfig.setPollInterval(RandomUtils.nextInt(100, 501)); |
|||
newConfig.setPartitions(RandomUtils.nextInt(1, 10)); |
|||
newConfig.setPackProcessingTimeout(RandomUtils.nextLong(100, 5001)); |
|||
newConfig.getSubmitStrategy().setType(SubmitStrategyType.values()[RandomUtils.nextInt(0, SubmitStrategyType.values().length)]); |
|||
newConfig.getProcessingStrategy().setType(ProcessingStrategyType.values()[RandomUtils.nextInt(0, ProcessingStrategyType.values().length)]); |
|||
log.info("Generated new config: consumerPerPartition={}, pollInterval={}, processingStrategy={}", |
|||
newConfig.isConsumerPerPartition(), newConfig.getPollInterval(), newConfig.getProcessingStrategy().getType()); |
|||
return newConfig; |
|||
}; |
|||
Supplier<Set<TopicPartitionInfo>> partitionsUpdater = () -> { |
|||
int partitionsCount = RandomUtils.nextInt(0, 20); |
|||
int[] partitions = IntStream.generate(() -> RandomUtils.nextInt(0, 20)) |
|||
.distinct().limit(partitionsCount) |
|||
.sorted().toArray(); |
|||
log.info("Generated new partitions: {}", Arrays.toString(partitions)); |
|||
return createTpis(partitions); |
|||
}; |
|||
|
|||
int iterations = 100; |
|||
Queue latestConfig = queue; |
|||
Set<TopicPartitionInfo> latestPartitions = Collections.emptySet(); |
|||
for (int i = 1; i <= iterations; i++) { |
|||
boolean updateQueueConfig = RandomUtils.nextBoolean(); |
|||
boolean updatePartitions = !updateQueueConfig; |
|||
if (updateQueueConfig) { |
|||
latestConfig = queueConfigUpdater.get(); |
|||
consumerManager.update(latestConfig); |
|||
} |
|||
if (updatePartitions) { |
|||
latestPartitions = partitionsUpdater.get(); |
|||
consumerManager.update(latestPartitions); |
|||
} |
|||
Thread.sleep(RandomUtils.nextLong(0, 200)); |
|||
} |
|||
if (latestPartitions.isEmpty()) { |
|||
do { |
|||
latestPartitions = partitionsUpdater.get(); |
|||
} while (latestPartitions.isEmpty()); |
|||
consumerManager.update(latestPartitions); |
|||
} |
|||
|
|||
Queue expectedConfig = latestConfig; |
|||
Set<TopicPartitionInfo> expectedPartitions = latestPartitions; |
|||
await().atMost(5, TimeUnit.SECONDS) |
|||
.untilAsserted(() -> { |
|||
assertThat(consumerManager.getQueue()).isEqualTo(expectedConfig); |
|||
assertThat(consumerManager.getPartitions()).isEqualTo(expectedPartitions); |
|||
}); |
|||
|
|||
if (expectedConfig.isConsumerPerPartition()) { |
|||
await().atMost(5, TimeUnit.SECONDS).until(() -> { |
|||
for (TopicPartitionInfo partition : expectedPartitions) { |
|||
if (consumers.stream().noneMatch(consumer -> consumer.subscribed && |
|||
consumer.pollingStarted && Set.of(partition).equals(consumer.getPartitions()))) { |
|||
return false; |
|||
} |
|||
} |
|||
return consumers.size() == expectedPartitions.size(); |
|||
}); |
|||
} else { |
|||
await().atMost(5, TimeUnit.SECONDS).until(() -> { |
|||
return consumers.size() == 1 && consumers.stream() |
|||
.anyMatch(consumer -> consumer.subscribed && consumer.pollingStarted && |
|||
expectedPartitions.equals(consumer.getPartitions())); |
|||
}); |
|||
} |
|||
Mockito.reset(ruleEngineConsumerContext.getSubmitStrategyFactory()); |
|||
Mockito.reset(ruleEngineConsumerContext.getProcessingStrategyFactory()); |
|||
consumers.forEach(Mockito::clearInvocations); |
|||
|
|||
await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> { |
|||
for (TestConsumer consumer : consumers) { |
|||
verify(consumer, atLeastOnce().description("consumer " + consumer.topics)).poll(expectedConfig.getPollInterval()); |
|||
} |
|||
verify(ruleEngineConsumerContext.getSubmitStrategyFactory(), atLeastOnce()).newInstance(any(), eq(expectedConfig.getSubmitStrategy())); |
|||
verify(ruleEngineConsumerContext.getProcessingStrategyFactory(), atLeastOnce()).newInstance(any(), eq(expectedConfig.getProcessingStrategy())); |
|||
}); |
|||
} |
|||
|
|||
private void verifySubscribedAndLaunched(TestConsumer consumer, Set<TopicPartitionInfo> expectedPartitions) { |
|||
await().atMost(2, TimeUnit.SECONDS) |
|||
.until(() -> consumer.subscribed && consumer.getPartitions().equals(expectedPartitions) && consumer.pollingStarted); |
|||
verify(consumer, times(1)).subscribe(any()); |
|||
verify(consumer).subscribe(eq(expectedPartitions)); |
|||
verify(consumer).doSubscribe(argThat(topics -> topics.containsAll(expectedPartitions.stream() |
|||
.map(TopicPartitionInfo::getFullTopicName).collect(Collectors.toList())))); |
|||
verify(consumer, atLeastOnce()).poll(eq((long) queue.getPollInterval())); |
|||
verify(consumer, atLeastOnce()).doPoll(eq((long) queue.getPollInterval())); |
|||
verify(consumer, never()).unsubscribe(); |
|||
Mockito.reset(consumer); |
|||
} |
|||
|
|||
private void verifySubscribedAndLaunched(TestConsumer consumer, int... expectedPartitions) { |
|||
verifySubscribedAndLaunched(consumer, createTpis(expectedPartitions)); |
|||
} |
|||
|
|||
private void verifyUnsubscribedAndStopped(TestConsumer consumer) { |
|||
await().atMost(2, TimeUnit.SECONDS) |
|||
.until(() -> !consumer.subscribed && !consumer.topics.isEmpty()); |
|||
verify(consumer, never()).subscribe(any()); |
|||
verify(consumer, never()).doSubscribe(any()); |
|||
assertThat(consumers).doesNotContain(consumer); |
|||
Mockito.reset(consumer); |
|||
} |
|||
|
|||
private void verifyNotTouched(TestConsumer consumer) { |
|||
verify(consumer, never()).subscribe(any()); |
|||
verify(consumer, never()).subscribe(); |
|||
verify(consumer, never()).doSubscribe(any()); |
|||
verify(consumer, never()).unsubscribe(); |
|||
verify(consumer, never()).doUnsubscribe(); |
|||
} |
|||
|
|||
private void verifyMsgProcessed(TbMsg tbMsg) { |
|||
await().atMost(2, TimeUnit.SECONDS).untilAsserted(() -> { |
|||
verify(actorContext, atLeastOnce()).tell(argThat(msg -> { |
|||
return ((QueueToRuleEngineMsg) msg).getMsg().getId().equals(tbMsg.getId()); |
|||
})); |
|||
}); |
|||
} |
|||
|
|||
// for consumer-per-partition
|
|||
private TestConsumer getConsumer(TopicPartitionInfo tpi) { |
|||
return await().atMost(5, TimeUnit.SECONDS) |
|||
.until(() -> consumers.stream() |
|||
.filter(consumer -> consumer.getPartitions() != null && |
|||
consumer.getPartitions().size() == 1 && |
|||
consumer.getPartitions().contains(tpi)) |
|||
.findFirst().orElse(null), Objects::nonNull); |
|||
} |
|||
|
|||
private TestConsumer getConsumer(int partition) { |
|||
return await().atMost(5, TimeUnit.SECONDS) |
|||
.until(() -> consumers.stream() |
|||
.filter(consumer -> consumer.getPartitions() != null && |
|||
consumer.getPartitions().size() == 1 && |
|||
consumer.getPartitions().stream() |
|||
.anyMatch(tpi -> tpi.getPartition().get().equals(partition))) |
|||
.findFirst().orElse(null), Objects::nonNull); |
|||
} |
|||
|
|||
// for single consumer
|
|||
private TestConsumer getConsumer() { |
|||
return await().atMost(5, TimeUnit.SECONDS) |
|||
.until(() -> consumers.size() == 1 ? consumers.iterator().next() : null, Objects::nonNull); |
|||
} |
|||
|
|||
private Set<TopicPartitionInfo> createTpis(int... partitions) { |
|||
return Arrays.stream(partitions) |
|||
.mapToObj(n -> TopicPartitionInfo.builder() |
|||
.tenantId(queue.getTenantId()) |
|||
.topic(queue.getTopic()) |
|||
.partition(n) |
|||
.myPartition(true) |
|||
.build()) |
|||
.collect(Collectors.toSet()); |
|||
} |
|||
|
|||
|
|||
class TestConsumer extends AbstractTbQueueConsumerTemplate<TbMsg, TbProtoQueueMsg<ToRuleEngineMsg>> { |
|||
|
|||
@Getter |
|||
private List<String> topics; |
|||
|
|||
private boolean subscribed; |
|||
private boolean pollingStarted; |
|||
|
|||
private TbMsg testMsg; |
|||
|
|||
public TestConsumer(String topic) { |
|||
super(topic); |
|||
} |
|||
|
|||
@SneakyThrows |
|||
@Override |
|||
protected List<TbMsg> doPoll(long durationInMillis) { |
|||
log.debug("doPoll({} ms)", durationInMillis); |
|||
if (!subscribed) { |
|||
throw new IllegalStateException("Cannot poll because not subscribed"); |
|||
} |
|||
pollingStarted = true; |
|||
if (testMsg != null && RandomUtils.nextBoolean()) { |
|||
Thread.sleep(100); |
|||
return List.of(testMsg); |
|||
} |
|||
return Collections.emptyList(); |
|||
} |
|||
|
|||
@Override |
|||
protected TbProtoQueueMsg<ToRuleEngineMsg> decode(TbMsg tbMsg) throws IOException { |
|||
log.debug("decode()"); |
|||
UUID tenantId = UUID.randomUUID(); |
|||
return new TbProtoQueueMsg<>(UUID.randomUUID(), ToRuleEngineMsg.newBuilder() |
|||
.setTenantIdMSB(tenantId.getMostSignificantBits()) |
|||
.setTenantIdLSB(tenantId.getLeastSignificantBits()) |
|||
.addRelationTypes("Success") |
|||
.setTbMsg(TbMsg.toByteString(tbMsg)) |
|||
.build()); |
|||
} |
|||
|
|||
@Override |
|||
protected void doSubscribe(List<String> topicNames) { |
|||
log.debug("doSubscribe({})", topicNames); |
|||
this.topics = topicNames; |
|||
subscribed = true; |
|||
} |
|||
|
|||
@Override |
|||
protected void doCommit() { |
|||
if (!subscribed) { |
|||
throw new IllegalStateException("Cannot commit because not subscribed"); |
|||
} |
|||
log.debug("doCommit() totalConsumedMsgs = {}", totalConsumedMsgs.incrementAndGet()); |
|||
} |
|||
|
|||
@Override |
|||
public void unsubscribe() { |
|||
super.unsubscribe(); |
|||
consumers.remove(this); |
|||
} |
|||
|
|||
@Override |
|||
protected void doUnsubscribe() { |
|||
log.debug("doUnsubscribe()"); |
|||
if (!subscribed) { |
|||
throw new IllegalStateException("Already unsubscribed!"); |
|||
} |
|||
subscribed = false; |
|||
} |
|||
|
|||
@Override |
|||
protected boolean isLongPollingSupported() { |
|||
return false; |
|||
} |
|||
|
|||
public Set<TopicPartitionInfo> getPartitions() { |
|||
return partitions; |
|||
} |
|||
|
|||
public void setUpTestMsg() { |
|||
testMsg = TbMsg.newMsg(TbMsgType.POST_TELEMETRY_REQUEST, new DeviceId(UUID.randomUUID()), new TbMsgMetaData(), "{}"); |
|||
} |
|||
} |
|||
|
|||
} |
|||
Loading…
Reference in new issue