Commit 8d5c38b743a91c2ad3739c25c47f93ece5480f08

Authored by Andrii Shvaika
1 parent c7f282d3

Queue refactoring

Showing 14 changed files with 232 additions and 350 deletions
... ... @@ -23,7 +23,6 @@ import org.thingsboard.common.util.ThingsBoardThreadFactory;
23 23 import org.thingsboard.server.actors.ActorSystemContext;
24 24 import org.thingsboard.server.common.msg.queue.ServiceType;
25 25 import org.thingsboard.server.common.msg.queue.TbCallback;
26   -import org.thingsboard.server.gen.transport.TransportProtos;
27 26 import org.thingsboard.server.queue.TbQueueConsumer;
28 27 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
29 28 import org.thingsboard.server.queue.discovery.PartitionChangeEvent;
... ...
... ... @@ -23,7 +23,6 @@ import java.util.LinkedHashMap;
23 23 import java.util.Map;
24 24 import java.util.UUID;
25 25 import java.util.concurrent.ConcurrentMap;
26   -import java.util.concurrent.ExecutorService;
27 26 import java.util.concurrent.atomic.AtomicInteger;
28 27 import java.util.function.BiConsumer;
29 28
... ... @@ -77,8 +76,8 @@ public class BatchTbRuleEngineSubmitStrategy extends AbstractTbRuleEngineSubmitS
77 76 }
78 77 }
79 78 int submitSize = pendingPack.size();
80   - if (log.isInfoEnabled() && submitSize > 0) {
81   - log.info("[{}] submitting [{}] messages to rule engine", queueName, submitSize);
  79 + if (log.isDebugEnabled() && submitSize > 0) {
  80 + log.debug("[{}] submitting [{}] messages to rule engine", queueName, submitSize);
82 81 }
83 82 pendingPack.forEach(msgConsumer);
84 83 }
... ...
... ... @@ -19,14 +19,8 @@ import lombok.extern.slf4j.Slf4j;
19 19 import org.thingsboard.server.gen.transport.TransportProtos;
20 20 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
21 21
22   -import java.util.ArrayList;
23   -import java.util.List;
24 22 import java.util.UUID;
25   -import java.util.concurrent.ConcurrentMap;
26   -import java.util.concurrent.ExecutorService;
27 23 import java.util.function.BiConsumer;
28   -import java.util.function.Function;
29   -import java.util.stream.Collectors;
30 24
31 25 @Slf4j
32 26 public class BurstTbRuleEngineSubmitStrategy extends AbstractTbRuleEngineSubmitStrategy {
... ... @@ -37,8 +31,8 @@ public class BurstTbRuleEngineSubmitStrategy extends AbstractTbRuleEngineSubmitS
37 31
38 32 @Override
39 33 public void submitAttempt(BiConsumer<UUID, TbProtoQueueMsg<TransportProtos.ToRuleEngineMsg>> msgConsumer) {
40   - if (log.isInfoEnabled()) {
41   - log.info("[{}] submitting [{}] messages to rule engine", queueName, orderedMsgList.size());
  34 + if (log.isDebugEnabled()) {
  35 + log.debug("[{}] submitting [{}] messages to rule engine", queueName, orderedMsgList.size());
42 36 }
43 37 orderedMsgList.forEach(pair -> msgConsumer.accept(pair.uuid, pair.msg));
44 38 }
... ...
... ... @@ -15,26 +15,18 @@
15 15 */
16 16 package org.thingsboard.server.service.queue.processing;
17 17
18   -import com.google.protobuf.InvalidProtocolBufferException;
19 18 import lombok.extern.slf4j.Slf4j;
20 19 import org.thingsboard.server.common.data.id.EntityId;
21   -import org.thingsboard.server.common.data.id.EntityIdFactory;
22   -import org.thingsboard.server.common.msg.TbMsg;
23   -import org.thingsboard.server.common.msg.gen.MsgProtos;
24   -import org.thingsboard.server.common.msg.queue.TbMsgCallback;
25 20 import org.thingsboard.server.gen.transport.TransportProtos;
26 21 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
27 22
28   -import java.util.ArrayList;
29 23 import java.util.LinkedList;
30 24 import java.util.List;
31 25 import java.util.Queue;
32 26 import java.util.UUID;
33 27 import java.util.concurrent.ConcurrentHashMap;
34 28 import java.util.concurrent.ConcurrentMap;
35   -import java.util.concurrent.atomic.AtomicInteger;
36 29 import java.util.function.BiConsumer;
37   -import java.util.stream.Collectors;
38 30
39 31 @Slf4j
40 32 public abstract class SequentialByEntityIdTbRuleEngineSubmitStrategy extends AbstractTbRuleEngineSubmitStrategy {
... ...
... ... @@ -30,6 +30,5 @@ public class SequentialByTenantIdTbRuleEngineSubmitStrategy extends SequentialBy
30 30 @Override
31 31 protected EntityId getEntityId(TransportProtos.ToRuleEngineMsg msg) {
32 32 return new TenantId(new UUID(msg.getTenantIdMSB(), msg.getTenantIdLSB()));
33   -
34 33 }
35 34 }
... ...
... ... @@ -19,8 +19,6 @@ import lombok.extern.slf4j.Slf4j;
19 19 import org.thingsboard.server.gen.transport.TransportProtos;
20 20 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
21 21
22   -import java.util.LinkedHashMap;
23   -import java.util.Map;
24 22 import java.util.UUID;
25 23 import java.util.concurrent.ConcurrentMap;
26 24 import java.util.concurrent.atomic.AtomicInteger;
... ... @@ -63,8 +61,8 @@ public class SequentialTbRuleEngineSubmitStrategy extends AbstractTbRuleEngineSu
63 61 if (idx < listSize) {
64 62 IdMsgPair pair = orderedMsgList.get(idx);
65 63 expectedMsgId = pair.uuid;
66   - if (log.isInfoEnabled()) {
67   - log.info("[{}] submitting [{}] message to rule engine", queueName, pair.msg);
  64 + if (log.isDebugEnabled()) {
  65 + log.debug("[{}] submitting [{}] message to rule engine", queueName, pair.msg);
68 66 }
69 67 msgConsumer.accept(pair.uuid, pair.msg);
70 68 }
... ...
... ... @@ -82,10 +82,10 @@ public class TbRuleEngineProcessingStrategyFactory {
82 82 retryCount++;
83 83 double failedCount = result.getFailedMap().size() + result.getPendingMap().size();
84 84 if (maxRetries > 0 && retryCount > maxRetries) {
85   - log.info("[{}] Skip reprocess of the rule engine pack due to max retries", queueName);
  85 + log.debug("[{}] Skip reprocess of the rule engine pack due to max retries", queueName);
86 86 return new TbRuleEngineProcessingDecision(true, null);
87 87 } else if (maxAllowedFailurePercentage > 0 && (failedCount / initialTotalCount) > maxAllowedFailurePercentage) {
88   - log.info("[{}] Skip reprocess of the rule engine pack due to max allowed failure percentage", queueName);
  88 + log.debug("[{}] Skip reprocess of the rule engine pack due to max allowed failure percentage", queueName);
89 89 return new TbRuleEngineProcessingDecision(true, null);
90 90 } else {
91 91 ConcurrentMap<UUID, TbProtoQueueMsg<TransportProtos.ToRuleEngineMsg>> toReprocess = new ConcurrentHashMap<>(initialTotalCount);
... ... @@ -98,7 +98,7 @@ public class TbRuleEngineProcessingStrategyFactory {
98 98 if (retrySuccessful) {
99 99 result.getSuccessMap().forEach(toReprocess::put);
100 100 }
101   - log.info("[{}] Going to reprocess {} messages", queueName, toReprocess.size());
  101 + log.debug("[{}] Going to reprocess {} messages", queueName, toReprocess.size());
102 102 if (log.isTraceEnabled()) {
103 103 toReprocess.forEach((id, msg) -> log.trace("Going to reprocess [{}]: {}", id, TbMsg.fromBytes(msg.getValue().getTbMsg().toByteArray(), TbMsgCallback.EMPTY)));
104 104 }
... ... @@ -126,7 +126,7 @@ public class TbRuleEngineProcessingStrategyFactory {
126 126 @Override
127 127 public TbRuleEngineProcessingDecision analyze(TbRuleEngineProcessingResult result) {
128 128 if (!result.isSuccess()) {
129   - log.info("[{}] Reprocessing skipped for {} failed and {} timeout messages", queueName, result.getFailedMap().size(), result.getPendingMap().size());
  129 + log.debug("[{}] Reprocessing skipped for {} failed and {} timeout messages", queueName, result.getFailedMap().size(), result.getPendingMap().size());
130 130 }
131 131 if (log.isTraceEnabled()) {
132 132 result.getFailedMap().forEach((id, msg) -> log.trace("Failed messages [{}]: {}", id, TbMsg.fromBytes(msg.getValue().getTbMsg().toByteArray(), TbMsgCallback.EMPTY)));
... ...
... ... @@ -31,9 +31,9 @@ import org.apache.qpid.proton.amqp.transport.SenderSettleMode;
31 31 import org.springframework.util.CollectionUtils;
32 32 import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
33 33 import org.thingsboard.server.queue.TbQueueAdmin;
34   -import org.thingsboard.server.queue.TbQueueConsumer;
35 34 import org.thingsboard.server.queue.TbQueueMsg;
36 35 import org.thingsboard.server.queue.TbQueueMsgDecoder;
  36 +import org.thingsboard.server.queue.common.AbstractTbQueueConsumerTemplate;
37 37 import org.thingsboard.server.queue.common.DefaultTbQueueMsg;
38 38
39 39 import java.time.Duration;
... ... @@ -50,102 +50,72 @@ import java.util.stream.Collectors;
50 50 import java.util.stream.Stream;
51 51
52 52 @Slf4j
53   -public class TbServiceBusConsumerTemplate<T extends TbQueueMsg> implements TbQueueConsumer<T> {
  53 +public class TbServiceBusConsumerTemplate<T extends TbQueueMsg> extends AbstractTbQueueConsumerTemplate<MessageWithDeliveryTag, T> {
54 54 private final TbQueueAdmin admin;
55   - private final String topic;
56 55 private final TbQueueMsgDecoder<T> decoder;
57 56 private final TbServiceBusSettings serviceBusSettings;
58 57
59 58 private final Gson gson = new Gson();
60 59
61 60 private Set<CoreMessageReceiver> receivers;
62   - private volatile Set<TopicPartitionInfo> partitions;
63   - private volatile boolean subscribed;
64   - private volatile boolean stopped = false;
65 61 private Map<CoreMessageReceiver, Collection<MessageWithDeliveryTag>> pendingMessages = new ConcurrentHashMap<>();
66 62 private volatile int messagesPerQueue;
67 63
68 64 public TbServiceBusConsumerTemplate(TbQueueAdmin admin, TbServiceBusSettings serviceBusSettings, String topic, TbQueueMsgDecoder<T> decoder) {
  65 + super(topic);
69 66 this.admin = admin;
70 67 this.decoder = decoder;
71   - this.topic = topic;
72 68 this.serviceBusSettings = serviceBusSettings;
73 69 }
74 70
75 71 @Override
76   - public String getTopic() {
77   - return topic;
  72 + protected List<MessageWithDeliveryTag> doPoll(long durationInMillis) {
  73 + List<CompletableFuture<Collection<MessageWithDeliveryTag>>> messageFutures =
  74 + receivers.stream()
  75 + .map(receiver -> receiver
  76 + .receiveAsync(messagesPerQueue, Duration.ofMillis(durationInMillis))
  77 + .whenComplete((messages, err) -> {
  78 + if (!CollectionUtils.isEmpty(messages)) {
  79 + pendingMessages.put(receiver, messages);
  80 + } else if (err != null) {
  81 + log.error("Failed to receive messages.", err);
  82 + }
  83 + }))
  84 + .collect(Collectors.toList());
  85 + try {
  86 + return fromList(messageFutures)
  87 + .get()
  88 + .stream()
  89 + .flatMap(messages -> CollectionUtils.isEmpty(messages) ? Stream.empty() : messages.stream())
  90 + .collect(Collectors.toList());
  91 + } catch (InterruptedException | ExecutionException e) {
  92 + if (stopped) {
  93 + log.info("[{}] Service Bus consumer is stopped.", getTopic());
  94 + } else {
  95 + log.error("Failed to receive messages", e);
  96 + }
  97 + return Collections.emptyList();
  98 + }
78 99 }
79 100
80 101 @Override
81   - public void subscribe() {
82   - partitions = Collections.singleton(new TopicPartitionInfo(topic, null, null, true));
83   - subscribed = false;
  102 + protected void doSubscribe(List<String> topicNames) {
  103 + createReceivers();
  104 + messagesPerQueue = receivers.size() / partitions.size();
84 105 }
85 106
86 107 @Override
87   - public void subscribe(Set<TopicPartitionInfo> partitions) {
88   - this.partitions = partitions;
89   - subscribed = false;
  108 + protected void doCommit() {
  109 + pendingMessages.forEach((receiver, msgs) ->
  110 + msgs.forEach(msg -> receiver.completeMessageAsync(msg.getDeliveryTag(), TransactionContext.NULL_TXN)));
  111 + pendingMessages.clear();
90 112 }
91 113
92 114 @Override
93   - public void unsubscribe() {
94   - stopped = true;
  115 + protected void doUnsubscribe() {
95 116 receivers.forEach(CoreMessageReceiver::closeAsync);
96 117 }
97 118
98   - @Override
99   - public List<T> poll(long durationInMillis) {
100   - if (!subscribed && partitions == null) {
101   - try {
102   - Thread.sleep(durationInMillis);
103   - } catch (InterruptedException e) {
104   - log.debug("Failed to await subscription", e);
105   - }
106   - } else {
107   - if (!subscribed) {
108   - createReceivers();
109   - messagesPerQueue = receivers.size() / partitions.size();
110   - subscribed = true;
111   - }
112   -
113   - List<CompletableFuture<Collection<MessageWithDeliveryTag>>> messageFutures =
114   - receivers.stream()
115   - .map(receiver -> receiver
116   - .receiveAsync(messagesPerQueue, Duration.ofMillis(durationInMillis))
117   - .whenComplete((messages, err) -> {
118   - if (!CollectionUtils.isEmpty(messages)) {
119   - pendingMessages.put(receiver, messages);
120   - } else if (err != null) {
121   - log.error("Failed to receive messages.", err);
122   - }
123   - }))
124   - .collect(Collectors.toList());
125   - try {
126   - return fromList(messageFutures)
127   - .get()
128   - .stream()
129   - .flatMap(messages -> CollectionUtils.isEmpty(messages) ? Stream.empty() : messages.stream())
130   - .map(message -> {
131   - try {
132   - return decode(message);
133   - } catch (InvalidProtocolBufferException e) {
134   - log.error("Failed to parse message.", e);
135   - throw new RuntimeException("Failed to parse message.", e);
136   - }
137   - }).collect(Collectors.toList());
138   - } catch (InterruptedException | ExecutionException e) {
139   - if (stopped) {
140   - log.info("[{}] Service Bus consumer is stopped.", topic);
141   - } else {
142   - log.error("Failed to receive messages", e);
143   - }
144   - }
145   - }
146   - return Collections.emptyList();
147   - }
148   -
149 119 private void createReceivers() {
150 120 List<CompletableFuture<CoreMessageReceiver>> receiverFutures = partitions.stream()
151 121 .map(TopicPartitionInfo::getFullTopicName)
... ... @@ -167,7 +137,7 @@ public class TbServiceBusConsumerTemplate<T extends TbQueueMsg> implements TbQue
167 137 receivers = new HashSet<>(fromList(receiverFutures).get());
168 138 } catch (InterruptedException | ExecutionException e) {
169 139 if (stopped) {
170   - log.info("[{}] Service Bus consumer is stopped.", topic);
  140 + log.info("[{}] Service Bus consumer is stopped.", getTopic());
171 141 } else {
172 142 log.error("Failed to create receivers", e);
173 143 }
... ... @@ -196,13 +166,7 @@ public class TbServiceBusConsumerTemplate<T extends TbQueueMsg> implements TbQue
196 166 }
197 167
198 168 @Override
199   - public void commit() {
200   - pendingMessages.forEach((receiver, msgs) ->
201   - msgs.forEach(msg -> receiver.completeMessageAsync(msg.getDeliveryTag(), TransactionContext.NULL_TXN)));
202   - pendingMessages.clear();
203   - }
204   -
205   - private T decode(MessageWithDeliveryTag data) throws InvalidProtocolBufferException {
  169 + protected T decode(MessageWithDeliveryTag data) throws InvalidProtocolBufferException {
206 170 DefaultTbQueueMsg msg = gson.fromJson(new String(((Data) data.getMessage().getBody()).getValue().getArray()), DefaultTbQueueMsg.class);
207 171 return decoder.decode(msg);
208 172 }
... ...
  1 +/**
  2 + * Copyright © 2016-2020 The Thingsboard Authors
  3 + *
  4 + * Licensed under the Apache License, Version 2.0 (the "License");
  5 + * you may not use this file except in compliance with the License.
  6 + * You may obtain a copy of the License at
  7 + *
  8 + * http://www.apache.org/licenses/LICENSE-2.0
  9 + *
  10 + * Unless required by applicable law or agreed to in writing, software
  11 + * distributed under the License is distributed on an "AS IS" BASIS,
  12 + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 + * See the License for the specific language governing permissions and
  14 + * limitations under the License.
  15 + */
  16 +package org.thingsboard.server.queue.common;
  17 +
  18 +import com.google.common.util.concurrent.ListeningExecutorService;
  19 +import com.google.common.util.concurrent.MoreExecutors;
  20 +import lombok.extern.slf4j.Slf4j;
  21 +import org.thingsboard.server.queue.TbQueueMsg;
  22 +
  23 +import java.util.concurrent.Executors;
  24 +import java.util.concurrent.TimeUnit;
  25 +
  26 +@Slf4j
  27 +public abstract class AbstractParallelTbQueueConsumerTemplate<R, T extends TbQueueMsg> extends AbstractTbQueueConsumerTemplate<R, T> {
  28 +
  29 + protected ListeningExecutorService consumerExecutor;
  30 +
  31 + public AbstractParallelTbQueueConsumerTemplate(String topic) {
  32 + super(topic);
  33 + }
  34 +
  35 + protected void initNewExecutor(int threadPoolSize) {
  36 + if (consumerExecutor != null) {
  37 + consumerExecutor.shutdown();
  38 + try {
  39 + consumerExecutor.awaitTermination(1, TimeUnit.MINUTES);
  40 + } catch (InterruptedException e) {
  41 + log.trace("Interrupted while waiting for consumer executor to stop");
  42 + }
  43 + }
  44 + consumerExecutor = MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(threadPoolSize));
  45 + }
  46 +
  47 + protected void shutdownExecutor() {
  48 + if (consumerExecutor != null) {
  49 + consumerExecutor.shutdownNow();
  50 + }
  51 + }
  52 +
  53 +}
... ...
... ... @@ -28,11 +28,13 @@ import java.util.List;
28 28 import java.util.Set;
29 29 import java.util.concurrent.locks.Lock;
30 30 import java.util.concurrent.locks.ReentrantLock;
  31 +import java.util.stream.Collectors;
31 32
32 33 @Slf4j
33 34 public abstract class AbstractTbQueueConsumerTemplate<R, T extends TbQueueMsg> implements TbQueueConsumer<T> {
34 35
35 36 private volatile boolean subscribed;
  37 + protected volatile boolean stopped = false;
36 38 protected volatile Set<TopicPartitionInfo> partitions;
37 39 protected final Lock consumerLock = new ReentrantLock();
38 40
... ... @@ -74,10 +76,12 @@ public abstract class AbstractTbQueueConsumerTemplate<R, T extends TbQueueMsg> i
74 76 log.debug("Failed to await subscription", e);
75 77 }
76 78 } else {
  79 + long pollStartTs = System.currentTimeMillis();
77 80 consumerLock.lock();
78 81 try {
79 82 if (!subscribed) {
80   - doSubscribe();
  83 + List<String> topicNames = partitions.stream().map(TopicPartitionInfo::getFullTopicName).collect(Collectors.toList());
  84 + doSubscribe(topicNames);
81 85 subscribed = true;
82 86 }
83 87
... ... @@ -95,6 +99,17 @@ public abstract class AbstractTbQueueConsumerTemplate<R, T extends TbQueueMsg> i
95 99 }
96 100 });
97 101 return result;
  102 + } else {
  103 + long pollDuration = System.currentTimeMillis() - pollStartTs;
  104 + if (pollDuration < durationInMillis) {
  105 + try {
  106 + Thread.sleep(durationInMillis - pollDuration);
  107 + } catch (InterruptedException e) {
  108 + if (!stopped) {
  109 + log.error("Failed to wait.", e);
  110 + }
  111 + }
  112 + }
98 113 }
99 114 } finally {
100 115 consumerLock.unlock();
... ... @@ -115,6 +130,7 @@ public abstract class AbstractTbQueueConsumerTemplate<R, T extends TbQueueMsg> i
115 130
116 131 @Override
117 132 public void unsubscribe() {
  133 + stopped = true;
118 134 consumerLock.lock();
119 135 try {
120 136 doUnsubscribe();
... ... @@ -127,7 +143,7 @@ public abstract class AbstractTbQueueConsumerTemplate<R, T extends TbQueueMsg> i
127 143
128 144 abstract protected T decode(R record) throws IOException;
129 145
130   - abstract protected void doSubscribe();
  146 + abstract protected void doSubscribe(List<String> topicNames);
131 147
132 148 abstract protected void doCommit();
133 149
... ...
... ... @@ -69,8 +69,7 @@ public class TbKafkaConsumerTemplate<T extends TbQueueMsg> extends AbstractTbQue
69 69 }
70 70
71 71 @Override
72   - protected void doSubscribe() {
73   - List<String> topicNames = partitions.stream().map(TopicPartitionInfo::getFullTopicName).collect(Collectors.toList());
  72 + protected void doSubscribe( List<String> topicNames) {
74 73 topicNames.forEach(admin::createTopicIfNotExists);
75 74 consumer.subscribe(topicNames);
76 75 }
... ...
... ... @@ -15,6 +15,7 @@
15 15 */
16 16 package org.thingsboard.server.queue.pubsub;
17 17
  18 +import com.amazonaws.services.sqs.model.Message;
18 19 import com.google.api.core.ApiFuture;
19 20 import com.google.api.core.ApiFutures;
20 21 import com.google.cloud.pubsub.v1.stub.GrpcSubscriberStub;
... ... @@ -35,11 +36,14 @@ import org.thingsboard.server.queue.TbQueueAdmin;
35 36 import org.thingsboard.server.queue.TbQueueConsumer;
36 37 import org.thingsboard.server.queue.TbQueueMsg;
37 38 import org.thingsboard.server.queue.TbQueueMsgDecoder;
  39 +import org.thingsboard.server.queue.common.AbstractParallelTbQueueConsumerTemplate;
  40 +import org.thingsboard.server.queue.common.AbstractTbQueueConsumerTemplate;
38 41 import org.thingsboard.server.queue.common.DefaultTbQueueMsg;
39 42
40 43 import java.io.IOException;
41 44 import java.util.ArrayList;
42 45 import java.util.Collections;
  46 +import java.util.LinkedHashSet;
43 47 import java.util.List;
44 48 import java.util.Objects;
45 49 import java.util.Set;
... ... @@ -47,10 +51,11 @@ import java.util.concurrent.CopyOnWriteArrayList;
47 51 import java.util.concurrent.ExecutionException;
48 52 import java.util.concurrent.ExecutorService;
49 53 import java.util.concurrent.Executors;
  54 +import java.util.concurrent.TimeUnit;
50 55 import java.util.stream.Collectors;
51 56
52 57 @Slf4j
53   -public class TbPubSubConsumerTemplate<T extends TbQueueMsg> implements TbQueueConsumer<T> {
  58 +public class TbPubSubConsumerTemplate<T extends TbQueueMsg> extends AbstractParallelTbQueueConsumerTemplate<PubsubMessage, T> {
54 59
55 60 private final Gson gson = new Gson();
56 61 private final TbQueueAdmin admin;
... ... @@ -58,23 +63,18 @@ public class TbPubSubConsumerTemplate<T extends TbQueueMsg> implements TbQueueCo
58 63 private final TbQueueMsgDecoder<T> decoder;
59 64 private final TbPubSubSettings pubSubSettings;
60 65
61   - private volatile boolean subscribed;
62   - private volatile Set<TopicPartitionInfo> partitions;
63 66 private volatile Set<String> subscriptionNames;
64 67 private final List<AcknowledgeRequest> acknowledgeRequests = new CopyOnWriteArrayList<>();
65 68
66   - private ExecutorService consumerExecutor;
67 69 private final SubscriberStub subscriber;
68   - private volatile boolean stopped;
69   -
70 70 private volatile int messagesPerTopic;
71 71
72 72 public TbPubSubConsumerTemplate(TbQueueAdmin admin, TbPubSubSettings pubSubSettings, String topic, TbQueueMsgDecoder<T> decoder) {
  73 + super(topic);
73 74 this.admin = admin;
74 75 this.pubSubSettings = pubSubSettings;
75 76 this.topic = topic;
76 77 this.decoder = decoder;
77   -
78 78 try {
79 79 SubscriberStubSettings subscriberStubSettings =
80 80 SubscriberStubSettings.newBuilder()
... ... @@ -84,89 +84,50 @@ public class TbPubSubConsumerTemplate<T extends TbQueueMsg> implements TbQueueCo
84 84 .setMaxInboundMessageSize(pubSubSettings.getMaxMsgSize())
85 85 .build())
86 86 .build();
87   -
88 87 this.subscriber = GrpcSubscriberStub.create(subscriberStubSettings);
89 88 } catch (IOException e) {
90 89 log.error("Failed to create subscriber.", e);
91 90 throw new RuntimeException("Failed to create subscriber.", e);
92 91 }
93   - stopped = false;
94 92 }
95 93
96 94 @Override
97   - public String getTopic() {
98   - return topic;
  95 + protected List<PubsubMessage> doPoll(long durationInMillis) {
  96 + try {
  97 + List<ReceivedMessage> messages = receiveMessages();
  98 + if (!messages.isEmpty()) {
  99 + return messages.stream().map(ReceivedMessage::getMessage).collect(Collectors.toList());
  100 + }
  101 + } catch (ExecutionException | InterruptedException e) {
  102 + if (stopped) {
  103 + log.info("[{}] Pub/Sub consumer is stopped.", topic);
  104 + } else {
  105 + log.error("Failed to receive messages", e);
  106 + }
  107 + }
  108 + return Collections.emptyList();
99 109 }
100 110
101 111 @Override
102   - public void subscribe() {
103   - partitions = Collections.singleton(new TopicPartitionInfo(topic, null, null, true));
104   - subscribed = false;
  112 + protected void doSubscribe(List<String> topicNames) {
  113 + subscriptionNames = new LinkedHashSet<>(topicNames);
  114 + subscriptionNames.forEach(admin::createTopicIfNotExists);
  115 + initNewExecutor(subscriptionNames.size() + 1);
  116 + messagesPerTopic = pubSubSettings.getMaxMessages() / subscriptionNames.size();
105 117 }
106 118
107 119 @Override
108   - public void subscribe(Set<TopicPartitionInfo> partitions) {
109   - this.partitions = partitions;
110   - subscribed = false;
  120 + protected void doCommit() {
  121 + acknowledgeRequests.forEach(subscriber.acknowledgeCallable()::futureCall);
  122 + acknowledgeRequests.clear();
111 123 }
112 124
113 125 @Override
114   - public void unsubscribe() {
115   - stopped = true;
116   - if (consumerExecutor != null) {
117   - consumerExecutor.shutdownNow();
118   - }
119   -
  126 + protected void doUnsubscribe() {
120 127 if (subscriber != null) {
121 128 subscriber.close();
122 129 }
123   - }
124   -
125   - @Override
126   - public List<T> poll(long durationInMillis) {
127   - if (!subscribed && partitions == null) {
128   - try {
129   - Thread.sleep(durationInMillis);
130   - } catch (InterruptedException e) {
131   - log.debug("Failed to await subscription", e);
132   - }
133   - } else {
134   - if (!subscribed) {
135   - subscriptionNames = partitions.stream().map(TopicPartitionInfo::getFullTopicName).collect(Collectors.toSet());
136   - subscriptionNames.forEach(admin::createTopicIfNotExists);
137   - consumerExecutor = Executors.newFixedThreadPool(subscriptionNames.size());
138   - messagesPerTopic = pubSubSettings.getMaxMessages() / subscriptionNames.size();
139   - subscribed = true;
140   - }
141   - List<ReceivedMessage> messages;
142   - try {
143   - messages = receiveMessages();
144   - if (!messages.isEmpty()) {
145   - List<T> result = new ArrayList<>();
146   - messages.forEach(msg -> {
147   - try {
148   - result.add(decode(msg.getMessage()));
149   - } catch (InvalidProtocolBufferException e) {
150   - log.error("Failed decode record: [{}]", msg);
151   - }
152   - });
153   - return result;
154   - }
155   - } catch (ExecutionException | InterruptedException e) {
156   - if (stopped) {
157   - log.info("[{}] Pub/Sub consumer is stopped.", topic);
158   - } else {
159   - log.error("Failed to receive messages", e);
160   - }
161   - }
162   - }
163   - return Collections.emptyList();
164   - }
165   -
166   - @Override
167   - public void commit() {
168   - acknowledgeRequests.forEach(subscriber.acknowledgeCallable()::futureCall);
169   - acknowledgeRequests.clear();
  130 + shutdownExecutor();
170 131 }
171 132
172 133 private List<ReceivedMessage> receiveMessages() throws ExecutionException, InterruptedException {
... ... @@ -211,6 +172,7 @@ public class TbPubSubConsumerTemplate<T extends TbQueueMsg> implements TbQueueCo
211 172 return transform.get();
212 173 }
213 174
  175 + @Override
214 176 public T decode(PubsubMessage message) throws InvalidProtocolBufferException {
215 177 DefaultTbQueueMsg msg = gson.fromJson(message.getData().toStringUtf8(), DefaultTbQueueMsg.class);
216 178 return decoder.decode(msg);
... ...
... ... @@ -23,9 +23,9 @@ import com.rabbitmq.client.GetResponse;
23 23 import lombok.extern.slf4j.Slf4j;
24 24 import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
25 25 import org.thingsboard.server.queue.TbQueueAdmin;
26   -import org.thingsboard.server.queue.TbQueueConsumer;
27 26 import org.thingsboard.server.queue.TbQueueMsg;
28 27 import org.thingsboard.server.queue.TbQueueMsgDecoder;
  28 +import org.thingsboard.server.queue.common.AbstractTbQueueConsumerTemplate;
29 29 import org.thingsboard.server.queue.common.DefaultTbQueueMsg;
30 30
31 31 import java.io.IOException;
... ... @@ -37,33 +37,26 @@ import java.util.concurrent.TimeoutException;
37 37 import java.util.stream.Collectors;
38 38
39 39 @Slf4j
40   -public class TbRabbitMqConsumerTemplate<T extends TbQueueMsg> implements TbQueueConsumer<T> {
  40 +public class TbRabbitMqConsumerTemplate<T extends TbQueueMsg> extends AbstractTbQueueConsumerTemplate<GetResponse, T> {
41 41
42 42 private final Gson gson = new Gson();
43 43 private final TbQueueAdmin admin;
44   - private final String topic;
45 44 private final TbQueueMsgDecoder<T> decoder;
46   - private final TbRabbitMqSettings rabbitMqSettings;
47 45 private final Channel channel;
48 46 private final Connection connection;
49 47
50   - private volatile Set<TopicPartitionInfo> partitions;
51   - private volatile boolean subscribed;
52 48 private volatile Set<String> queues;
53   - private volatile boolean stopped;
54 49
55 50 public TbRabbitMqConsumerTemplate(TbQueueAdmin admin, TbRabbitMqSettings rabbitMqSettings, String topic, TbQueueMsgDecoder<T> decoder) {
  51 + super(topic);
56 52 this.admin = admin;
57 53 this.decoder = decoder;
58   - this.topic = topic;
59   - this.rabbitMqSettings = rabbitMqSettings;
60 54 try {
61 55 connection = rabbitMqSettings.getConnectionFactory().newConnection();
62 56 } catch (IOException | TimeoutException e) {
63 57 log.error("Failed to create connection.", e);
64 58 throw new RuntimeException("Failed to create connection.", e);
65 59 }
66   -
67 60 try {
68 61 channel = connection.createChannel();
69 62 } catch (IOException e) {
... ... @@ -74,25 +67,42 @@ public class TbRabbitMqConsumerTemplate<T extends TbQueueMsg> implements TbQueue
74 67 }
75 68
76 69 @Override
77   - public String getTopic() {
78   - return topic;
  70 + protected List<GetResponse> doPoll(long durationInMillis) {
  71 + List<GetResponse> result = queues.stream()
  72 + .map(queue -> {
  73 + try {
  74 + return channel.basicGet(queue, false);
  75 + } catch (IOException e) {
  76 + log.error("Failed to get messages from queue: [{}]", queue);
  77 + throw new RuntimeException("Failed to get messages from queue.", e);
  78 + }
  79 + }).filter(Objects::nonNull).collect(Collectors.toList());
  80 + if (result.size() > 0) {
  81 + return result;
  82 + } else {
  83 + return Collections.emptyList();
  84 + }
79 85 }
80 86
81 87 @Override
82   - public void subscribe() {
83   - partitions = Collections.singleton(new TopicPartitionInfo(topic, null, null, true));
84   - subscribed = false;
  88 + protected void doSubscribe(List<String> topicNames) {
  89 + queues = partitions.stream()
  90 + .map(TopicPartitionInfo::getFullTopicName)
  91 + .collect(Collectors.toSet());
  92 + queues.forEach(admin::createTopicIfNotExists);
85 93 }
86 94
87 95 @Override
88   - public void subscribe(Set<TopicPartitionInfo> partitions) {
89   - this.partitions = partitions;
90   - subscribed = false;
  96 + protected void doCommit() {
  97 + try {
  98 + channel.basicAck(0, true);
  99 + } catch (IOException e) {
  100 + log.error("Failed to ack messages.", e);
  101 + }
91 102 }
92 103
93 104 @Override
94   - public void unsubscribe() {
95   - stopped = true;
  105 + protected void doUnsubscribe() {
96 106 if (channel != null) {
97 107 try {
98 108 channel.close();
... ... @@ -109,63 +119,6 @@ public class TbRabbitMqConsumerTemplate<T extends TbQueueMsg> implements TbQueue
109 119 }
110 120 }
111 121
112   - @Override
113   - public List<T> poll(long durationInMillis) {
114   - if (!subscribed && partitions == null) {
115   - try {
116   - Thread.sleep(durationInMillis);
117   - } catch (InterruptedException e) {
118   - log.debug("Failed to await subscription", e);
119   - }
120   - } else {
121   - if (!subscribed) {
122   - queues = partitions.stream()
123   - .map(TopicPartitionInfo::getFullTopicName)
124   - .collect(Collectors.toSet());
125   -
126   - queues.forEach(admin::createTopicIfNotExists);
127   - subscribed = true;
128   - }
129   -
130   - List<T> result = queues.stream()
131   - .map(queue -> {
132   - try {
133   - return channel.basicGet(queue, false);
134   - } catch (IOException e) {
135   - log.error("Failed to get messages from queue: [{}]", queue);
136   - throw new RuntimeException("Failed to get messages from queue.", e);
137   - }
138   - }).filter(Objects::nonNull).map(message -> {
139   - try {
140   - return decode(message);
141   - } catch (InvalidProtocolBufferException e) {
142   - log.error("Failed to decode message: [{}].", message);
143   - throw new RuntimeException("Failed to decode message.", e);
144   - }
145   - }).collect(Collectors.toList());
146   - if (result.size() > 0) {
147   - return result;
148   - }
149   - }
150   - try {
151   - Thread.sleep(durationInMillis);
152   - } catch (InterruptedException e) {
153   - if (!stopped) {
154   - log.error("Failed to wait.", e);
155   - }
156   - }
157   - return Collections.emptyList();
158   - }
159   -
160   - @Override
161   - public void commit() {
162   - try {
163   - channel.basicAck(0, true);
164   - } catch (IOException e) {
165   - log.error("Failed to ack messages.", e);
166   - }
167   - }
168   -
169 122 public T decode(GetResponse message) throws InvalidProtocolBufferException {
170 123 DefaultTbQueueMsg msg = gson.fromJson(new String(message.getBody()), DefaultTbQueueMsg.class);
171 124 return decoder.decode(msg);
... ...
... ... @@ -25,21 +25,17 @@ import com.amazonaws.services.sqs.model.Message;
25 25 import com.amazonaws.services.sqs.model.ReceiveMessageRequest;
26 26 import com.google.common.util.concurrent.Futures;
27 27 import com.google.common.util.concurrent.ListenableFuture;
28   -import com.google.common.util.concurrent.ListeningExecutorService;
29   -import com.google.common.util.concurrent.MoreExecutors;
30 28 import com.google.gson.Gson;
31 29 import com.google.protobuf.InvalidProtocolBufferException;
32 30 import lombok.Data;
33 31 import lombok.extern.slf4j.Slf4j;
34 32 import org.springframework.util.CollectionUtils;
35   -import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
36 33 import org.thingsboard.server.queue.TbQueueAdmin;
37   -import org.thingsboard.server.queue.TbQueueConsumer;
38 34 import org.thingsboard.server.queue.TbQueueMsg;
39 35 import org.thingsboard.server.queue.TbQueueMsgDecoder;
  36 +import org.thingsboard.server.queue.common.AbstractParallelTbQueueConsumerTemplate;
40 37 import org.thingsboard.server.queue.common.DefaultTbQueueMsg;
41 38
42   -import java.io.IOException;
43 39 import java.util.ArrayList;
44 40 import java.util.Collections;
45 41 import java.util.List;
... ... @@ -47,34 +43,28 @@ import java.util.Objects;
47 43 import java.util.Set;
48 44 import java.util.concurrent.CopyOnWriteArrayList;
49 45 import java.util.concurrent.ExecutionException;
50   -import java.util.concurrent.Executors;
51 46 import java.util.concurrent.TimeUnit;
52 47 import java.util.stream.Collectors;
53 48 import java.util.stream.Stream;
54 49
55 50 @Slf4j
56   -public class TbAwsSqsConsumerTemplate<T extends TbQueueMsg> implements TbQueueConsumer<T> {
  51 +public class TbAwsSqsConsumerTemplate<T extends TbQueueMsg> extends AbstractParallelTbQueueConsumerTemplate<Message, T> {
57 52
58 53 private static final int MAX_NUM_MSGS = 10;
59 54
60 55 private final Gson gson = new Gson();
61 56 private final TbQueueAdmin admin;
62 57 private final AmazonSQS sqsClient;
63   - private final String topic;
64 58 private final TbQueueMsgDecoder<T> decoder;
65 59 private final TbAwsSqsSettings sqsSettings;
66 60
67 61 private final List<AwsSqsMsgWrapper> pendingMessages = new CopyOnWriteArrayList<>();
68 62 private volatile Set<String> queueUrls;
69   - private volatile Set<TopicPartitionInfo> partitions;
70   - private ListeningExecutorService consumerExecutor;
71   - private volatile boolean subscribed;
72   - private volatile boolean stopped = false;
73 63
74 64 public TbAwsSqsConsumerTemplate(TbQueueAdmin admin, TbAwsSqsSettings sqsSettings, String topic, TbQueueMsgDecoder<T> decoder) {
  65 + super(topic);
75 66 this.admin = admin;
76 67 this.decoder = decoder;
77   - this.topic = topic;
78 68 this.sqsSettings = sqsSettings;
79 69
80 70 AWSCredentials awsCredentials = new BasicAWSCredentials(sqsSettings.getAccessKeyId(), sqsSettings.getSecretAccessKey());
... ... @@ -87,81 +77,64 @@ public class TbAwsSqsConsumerTemplate<T extends TbQueueMsg> implements TbQueueCo
87 77 }
88 78
89 79 @Override
90   - public String getTopic() {
91   - return topic;
  80 + protected void doSubscribe(List<String> topicNames) {
  81 + queueUrls = topicNames.stream().map(this::getQueueUrl).collect(Collectors.toSet());
  82 + initNewExecutor(queueUrls.size() * sqsSettings.getThreadsPerTopic() + 1);
92 83 }
93 84
94 85 @Override
95   - public void subscribe() {
96   - partitions = Collections.singleton(new TopicPartitionInfo(topic, null, null, true));
97   - subscribed = false;
  86 + protected List<Message> doPoll(long durationInMillis) {
  87 + if (!pendingMessages.isEmpty()) {
  88 + log.warn("Present {} non committed messages.", pendingMessages.size());
  89 + return Collections.emptyList();
  90 + }
  91 + int duration = (int) TimeUnit.MILLISECONDS.toSeconds(durationInMillis);
  92 + List<ListenableFuture<List<Message>>> futureList = queueUrls
  93 + .stream()
  94 + .map(url -> poll(url, duration))
  95 + .collect(Collectors.toList());
  96 + ListenableFuture<List<List<Message>>> futureResult = Futures.allAsList(futureList);
  97 + try {
  98 + return futureResult.get().stream()
  99 + .flatMap(List::stream)
  100 + .filter(Objects::nonNull)
  101 + .collect(Collectors.toList());
  102 + } catch (InterruptedException | ExecutionException e) {
  103 + if (stopped) {
  104 + log.info("[{}] Aws SQS consumer is stopped.", getTopic());
  105 + } else {
  106 + log.error("Failed to pool messages.", e);
  107 + }
  108 + return Collections.emptyList();
  109 + }
98 110 }
99 111
100 112 @Override
101   - public void subscribe(Set<TopicPartitionInfo> partitions) {
102   - this.partitions = partitions;
103   - subscribed = false;
  113 + public T decode(Message message) throws InvalidProtocolBufferException {
  114 + DefaultTbQueueMsg msg = gson.fromJson(message.getBody(), DefaultTbQueueMsg.class);
  115 + return decoder.decode(msg);
104 116 }
105 117
106 118 @Override
107   - public void unsubscribe() {
108   - stopped = true;
109   -
110   - if (sqsClient != null) {
111   - sqsClient.shutdown();
112   - }
113   - if (consumerExecutor != null) {
114   - consumerExecutor.shutdownNow();
115   - }
  119 + protected void doCommit() {
  120 + pendingMessages.forEach(msg ->
  121 + consumerExecutor.submit(() -> {
  122 + List<DeleteMessageBatchRequestEntry> entries = msg.getMessages()
  123 + .stream()
  124 + .map(message -> new DeleteMessageBatchRequestEntry(message.getMessageId(), message.getReceiptHandle()))
  125 + .collect(Collectors.toList());
  126 + sqsClient.deleteMessageBatch(msg.getUrl(), entries);
  127 + }));
  128 + pendingMessages.clear();
116 129 }
117 130
118 131 @Override
119   - public List<T> poll(long durationInMillis) {
120   - if (!subscribed && partitions == null) {
121   - try {
122   - Thread.sleep(durationInMillis);
123   - } catch (InterruptedException e) {
124   - log.debug("Failed to await subscription", e);
125   - }
126   - } else {
127   - if (!subscribed) {
128   - List<String> topicNames = partitions.stream().map(TopicPartitionInfo::getFullTopicName).collect(Collectors.toList());
129   - queueUrls = topicNames.stream().map(this::getQueueUrl).collect(Collectors.toSet());
130   - consumerExecutor = MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(queueUrls.size() * sqsSettings.getThreadsPerTopic() + 1));
131   - subscribed = true;
132   - }
133   -
134   - if (!pendingMessages.isEmpty()) {
135   - log.warn("Present {} non committed messages.", pendingMessages.size());
136   - return Collections.emptyList();
137   - }
138   -
139   - List<ListenableFuture<List<Message>>> futureList = queueUrls
140   - .stream()
141   - .map(url -> poll(url, (int) TimeUnit.MILLISECONDS.toSeconds(durationInMillis)))
142   - .collect(Collectors.toList());
143   - ListenableFuture<List<List<Message>>> futureResult = Futures.allAsList(futureList);
144   - try {
145   - return futureResult.get().stream()
146   - .flatMap(List::stream)
147   - .map(msg -> {
148   - try {
149   - return decode(msg);
150   - } catch (IOException e) {
151   - log.error("Failed to decode message: [{}]", msg);
152   - return null;
153   - }
154   - }).filter(Objects::nonNull)
155   - .collect(Collectors.toList());
156   - } catch (InterruptedException | ExecutionException e) {
157   - if (stopped) {
158   - log.info("[{}] Aws SQS consumer is stopped.", topic);
159   - } else {
160   - log.error("Failed to pool messages.", e);
161   - }
162   - }
  132 + protected void doUnsubscribe() {
  133 + stopped = true;
  134 + if (sqsClient != null) {
  135 + sqsClient.shutdown();
163 136 }
164   - return Collections.emptyList();
  137 + shutdownExecutor();
165 138 }
166 139
167 140 private ListenableFuture<List<Message>> poll(String url, int waitTimeSeconds) {
... ... @@ -194,25 +167,6 @@ public class TbAwsSqsConsumerTemplate<T extends TbQueueMsg> implements TbQueueCo
194 167 }, consumerExecutor);
195 168 }
196 169
197   - @Override
198   - public void commit() {
199   - pendingMessages.forEach(msg ->
200   - consumerExecutor.submit(() -> {
201   - List<DeleteMessageBatchRequestEntry> entries = msg.getMessages()
202   - .stream()
203   - .map(message -> new DeleteMessageBatchRequestEntry(message.getMessageId(), message.getReceiptHandle()))
204   - .collect(Collectors.toList());
205   - sqsClient.deleteMessageBatch(msg.getUrl(), entries);
206   - }));
207   -
208   - pendingMessages.clear();
209   - }
210   -
211   - public T decode(Message message) throws InvalidProtocolBufferException {
212   - DefaultTbQueueMsg msg = gson.fromJson(message.getBody(), DefaultTbQueueMsg.class);
213   - return decoder.decode(msg);
214   - }
215   -
216 170 @Data
217 171 private static class AwsSqsMsgWrapper {
218 172 private final String url;
... ...