WIP: CalculatedFieldConsumer refactoring

2 years ago · c332e7373f
6 changed files with 133 additions and 21 deletions
--- a/application/src/main/java/org/thingsboard/server/service/cf/CalculatedFieldExecutionService.java
+++ b/application/src/main/java/org/thingsboard/server/service/cf/CalculatedFieldExecutionService.java
@ -20,6 +20,8 @@ import org.thingsboard.rule.engine.api.TimeseriesSaveRequest;
 import org.thingsboard.server.common.data.kv.TimeseriesSaveResult;
 import org.thingsboard.server.common.msg.queue.TbCallback;
 import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldEntityUpdateMsgProto;
+import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldLinkedTelemetryMsgProto;
+import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldTelemetryMsgProto;
 import org.thingsboard.server.gen.transport.TransportProtos.ComponentLifecycleMsgProto;
 import org.thingsboard.server.service.cf.telemetry.CalculatedFieldTelemetryUpdateRequest;

@ -36,6 +38,10 @@ public interface CalculatedFieldExecutionService {

    void pushRequestToQueue(AttributesSaveRequest request, List<Long> result);

+    void onTelemetryMsg(CalculatedFieldTelemetryMsgProto msg, TbCallback callback);
+
+    void onLinkedTelemetryMsg(CalculatedFieldLinkedTelemetryMsgProto linkedMsg, TbCallback callback);
+
 //    void pushEntityUpdateMsg(TransportProtos.CalculatedFieldEntityUpdateMsgProto proto, TbCallback callback);

    /*  ===================================================== */
--- a/application/src/main/java/org/thingsboard/server/service/cf/DefaultCalculatedFieldExecutionService.java
+++ b/application/src/main/java/org/thingsboard/server/service/cf/DefaultCalculatedFieldExecutionService.java
@ -75,6 +75,7 @@ import org.thingsboard.server.dao.timeseries.TimeseriesService;
 import org.thingsboard.server.gen.transport.TransportProtos.AttributeScopeProto;
 import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldEntityUpdateMsgProto;
 import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldIdProto;
+import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldLinkedTelemetryMsgProto;
 import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldTelemetryMsgProto;
 import org.thingsboard.server.gen.transport.TransportProtos.ComponentLifecycleEvent;
 import org.thingsboard.server.gen.transport.TransportProtos.ComponentLifecycleMsgProto;
@ -234,6 +235,18 @@ public class DefaultCalculatedFieldExecutionService extends AbstractPartitionBas
        return send;
    }

+    @Override
+    public void onTelemetryMsg(CalculatedFieldTelemetryMsgProto msg, TbCallback callback) {
+
+        callback.onSuccess();
+    }
+
+    @Override
+    public void onLinkedTelemetryMsg(CalculatedFieldLinkedTelemetryMsgProto linkedMsg, TbCallback callback) {
+
+        callback.onSuccess();
+    }
+
    @Override
    protected Map<TopicPartitionInfo, List<ListenableFuture<?>>> onAddedPartitions(Set<TopicPartitionInfo> addedPartitions) {
        var result = new HashMap<TopicPartitionInfo, List<ListenableFuture<?>>>();
--- a/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCalculatedFieldConsumerService.java
+++ b/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCalculatedFieldConsumerService.java
@ -21,6 +21,8 @@ import com.google.common.util.concurrent.MoreExecutors;
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
 import lombok.Data;
+import lombok.Getter;
+import lombok.Setter;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.ApplicationEventPublisher;
@ -32,12 +34,19 @@ import org.thingsboard.server.common.data.id.CalculatedFieldId;
 import org.thingsboard.server.common.data.id.EntityIdFactory;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.queue.QueueConfig;
+import org.thingsboard.server.common.msg.MsgType;
+import org.thingsboard.server.common.msg.TbActorMsg;
 import org.thingsboard.server.common.msg.queue.ServiceType;
 import org.thingsboard.server.common.msg.queue.TbCallback;
+import org.thingsboard.server.common.msg.rpc.ToDeviceRpcRequestActorMsg;
+import org.thingsboard.server.common.util.ProtoUtils;
 import org.thingsboard.server.dao.tenant.TbTenantProfileCache;
 import org.thingsboard.server.gen.transport.TransportProtos;
+import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldLinkedTelemetryMsgProto;
+import org.thingsboard.server.gen.transport.TransportProtos.CalculatedFieldTelemetryMsgProto;
 import org.thingsboard.server.gen.transport.TransportProtos.ToCalculatedFieldMsg;
 import org.thingsboard.server.gen.transport.TransportProtos.ToCalculatedFieldNotificationMsg;
+import org.thingsboard.server.gen.transport.TransportProtos.ToCoreMsg;
 import org.thingsboard.server.queue.TbQueueConsumer;
 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
 import org.thingsboard.server.queue.discovery.PartitionService;
@ -50,12 +59,20 @@ import org.thingsboard.server.service.cf.CalculatedFieldCache;
 import org.thingsboard.server.service.cf.CalculatedFieldExecutionService;
 import org.thingsboard.server.service.profile.TbAssetProfileCache;
 import org.thingsboard.server.service.profile.TbDeviceProfileCache;
+import org.thingsboard.server.service.queue.DefaultTbCoreConsumerService.PendingMsgHolder;
 import org.thingsboard.server.service.queue.consumer.MainQueueConsumerManager;
 import org.thingsboard.server.service.queue.processing.AbstractConsumerService;
+import org.thingsboard.server.service.queue.processing.IdMsgPair;
 import org.thingsboard.server.service.security.auth.jwt.settings.JwtSettingsService;

 import java.util.List;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;

@Service
@TbRuleEngineComponent
@ -132,7 +149,46 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractConsumerSer
    }

    private void processMsgs(List<TbProtoQueueMsg<ToCalculatedFieldMsg>> msgs, TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> consumer, CalculatedFieldQueueConfig config) throws Exception {
-
+        List<IdMsgPair<ToCalculatedFieldMsg>> orderedMsgList = msgs.stream().map(msg -> new IdMsgPair<>(UUID.randomUUID(), msg)).toList();
+        ConcurrentMap<UUID, TbProtoQueueMsg<ToCalculatedFieldMsg>> pendingMap = orderedMsgList.stream().collect(
+                Collectors.toConcurrentMap(IdMsgPair::getUuid, IdMsgPair::getMsg));
+        CountDownLatch processingTimeoutLatch = new CountDownLatch(1);
+        TbPackProcessingContext<TbProtoQueueMsg<ToCalculatedFieldMsg>> ctx = new TbPackProcessingContext<>(
+                processingTimeoutLatch, pendingMap, new ConcurrentHashMap<>());
+        PendingMsgHolder<ToCalculatedFieldMsg> pendingMsgHolder = new PendingMsgHolder<>();
+        Future<?> packSubmitFuture = consumersExecutor.submit(() -> {
+            orderedMsgList.forEach((element) -> {
+                UUID id = element.getUuid();
+                TbProtoQueueMsg<ToCalculatedFieldMsg> msg = element.getMsg();
+                log.trace("[{}] Creating main callback for message: {}", id, msg.getValue());
+                TbCallback callback = new TbPackCallback<>(id, ctx);
+                try {
+                    ToCalculatedFieldMsg toCfMsg = msg.getValue();
+                    pendingMsgHolder.setMsg(toCfMsg);
+                    if (toCfMsg.hasTelemetryMsg()) {
+                        log.trace("[{}] Forwarding regular telemetry message for processing {}", id, toCfMsg.getTelemetryMsg());
+                        forwardToCalculatedFieldService(toCfMsg.getTelemetryMsg(), callback);
+                    } else if (toCfMsg.hasLinkedTelemetryMsg()) {
+                        log.trace("[{}] Forwarding linked telemetry message for processing {}", id, toCfMsg.getLinkedTelemetryMsg());
+                        forwardToCalculatedFieldService(toCfMsg.getLinkedTelemetryMsg(), callback);
+                    }
+                } catch (Throwable e) {
+                    log.warn("[{}] Failed to process message: {}", id, msg, e);
+                    callback.onFailure(e);
+                }
+            });
+        });
+        if (!processingTimeoutLatch.await(packProcessingTimeout, TimeUnit.MILLISECONDS)) {
+            if (!packSubmitFuture.isDone()) {
+                packSubmitFuture.cancel(true);
+                log.info("Timeout to process message: {}", pendingMsgHolder.getMsg());
+            }
+            if (log.isDebugEnabled()) {
+                ctx.getAckMap().forEach((id, msg) -> log.debug("[{}] Timeout to process message: {}", id, msg.getValue()));
+            }
+            ctx.getFailedMap().forEach((id, msg) -> log.warn("[{}] Failed to process message: {}", id, msg.getValue()));
+        }
+        consumer.commit();
    }

    @Override
@ -193,6 +249,33 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractConsumerSer
 //        }
 //    }
 //
+
+    private void forwardToCalculatedFieldService(CalculatedFieldLinkedTelemetryMsgProto linkedMsg, TbCallback callback) {
+        var msg = linkedMsg.getMsg();
+        var tenantId = toTenantId(msg.getTenantIdMSB(), msg.getTenantIdLSB());
+        var calculatedFieldId = new CalculatedFieldId(new UUID(msg.getEntityIdMSB(), msg.getEntityIdLSB()));
+        ListenableFuture<?> future = calculatedFieldsExecutor.submit(() -> calculatedFieldExecutionService.onLinkedTelemetryMsg(linkedMsg, callback));
+        DonAsynchron.withCallback(future,
+                __ -> callback.onSuccess(),
+                t -> {
+                    log.warn("[{}] Failed to process calculated field message for calculated field [{}]", tenantId.getId(), calculatedFieldId.getId(), t);
+                    callback.onFailure(t);
+                });
+
+    }
+
+    private void forwardToCalculatedFieldService(CalculatedFieldTelemetryMsgProto msg, TbCallback callback) {
+        var tenantId = toTenantId(msg.getTenantIdMSB(), msg.getTenantIdLSB());
+        var calculatedFieldId = new CalculatedFieldId(new UUID(msg.getEntityIdMSB(), msg.getEntityIdLSB()));
+        ListenableFuture<?> future = calculatedFieldsExecutor.submit(() -> calculatedFieldExecutionService.onTelemetryMsg(msg, callback));
+        DonAsynchron.withCallback(future,
+                __ -> callback.onSuccess(),
+                t -> {
+                    log.warn("[{}] Failed to process calculated field message for calculated field [{}]", tenantId.getId(), calculatedFieldId.getId(), t);
+                    callback.onFailure(t);
+                });
+    }
+
    private void forwardToCalculatedFieldService(TransportProtos.ComponentLifecycleMsgProto msg, TbCallback callback) {
        var tenantId = toTenantId(msg.getTenantIdMSB(), msg.getTenantIdLSB());
        var calculatedFieldId = new CalculatedFieldId(new UUID(msg.getEntityIdMSB(), msg.getEntityIdLSB()));
--- a/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCoreConsumerService.java
+++ b/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCoreConsumerService.java
@ -269,7 +269,7 @@ public class DefaultTbCoreConsumerService extends AbstractConsumerService<ToCore
        CountDownLatch processingTimeoutLatch = new CountDownLatch(1);
        TbPackProcessingContext<TbProtoQueueMsg<ToCoreMsg>> ctx = new TbPackProcessingContext<>(
                processingTimeoutLatch, pendingMap, new ConcurrentHashMap<>());
-        PendingMsgHolder pendingMsgHolder = new PendingMsgHolder();
+        PendingMsgHolder<ToCoreMsg> pendingMsgHolder = new PendingMsgHolder<>();
        Future<?> packSubmitFuture = consumersExecutor.submit(() -> {
            orderedMsgList.forEach((element) -> {
                UUID id = element.getUuid();
@ -278,7 +278,7 @@ public class DefaultTbCoreConsumerService extends AbstractConsumerService<ToCore
                TbCallback callback = new TbPackCallback<>(id, ctx);
                try {
                    ToCoreMsg toCoreMsg = msg.getValue();
-                    pendingMsgHolder.setToCoreMsg(toCoreMsg);
+                    pendingMsgHolder.setMsg(toCoreMsg);
                    if (toCoreMsg.hasToSubscriptionMgrMsg()) {
                        log.trace("[{}] Forwarding message to subscription manager service {}", id, toCoreMsg.getToSubscriptionMgrMsg());
                        forwardToSubMgrService(toCoreMsg.getToSubscriptionMgrMsg(), callback);
@ -335,8 +335,7 @@ public class DefaultTbCoreConsumerService extends AbstractConsumerService<ToCore
        if (!processingTimeoutLatch.await(packProcessingTimeout, TimeUnit.MILLISECONDS)) {
            if (!packSubmitFuture.isDone()) {
                packSubmitFuture.cancel(true);
-                ToCoreMsg lastSubmitMsg = pendingMsgHolder.getToCoreMsg();
-                log.info("Timeout to process message: {}", lastSubmitMsg);
+                log.info("Timeout to process message: {}", pendingMsgHolder.getMsg());
            }
            if (log.isDebugEnabled()) {
                ctx.getAckMap().forEach((id, msg) -> log.debug("[{}] Timeout to process message: {}", id, msg.getValue()));
@ -346,12 +345,6 @@ public class DefaultTbCoreConsumerService extends AbstractConsumerService<ToCore
        consumer.commit();
    }

-    private static class PendingMsgHolder {
-        @Getter
-        @Setter
-        private volatile ToCoreMsg toCoreMsg;
-    }
-
    @Override
    protected ServiceType getServiceType() {
        return ServiceType.TB_CORE;
--- a/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbEdgeConsumerService.java
+++ b/application/src/main/java/org/thingsboard/server/service/queue/DefaultTbEdgeConsumerService.java
@ -137,7 +137,7 @@ public class DefaultTbEdgeConsumerService extends AbstractConsumerService<ToEdge
        CountDownLatch processingTimeoutLatch = new CountDownLatch(1);
        TbPackProcessingContext<TbProtoQueueMsg<ToEdgeMsg>> ctx = new TbPackProcessingContext<>(
                processingTimeoutLatch, pendingMap, new ConcurrentHashMap<>());
-        PendingMsgHolder pendingMsgHolder = new PendingMsgHolder();
+        PendingMsgHolder<ToEdgeMsg> pendingMsgHolder = new PendingMsgHolder<>();
        Future<?> submitFuture = consumersExecutor.submit(() -> {
            orderedMsgList.forEach((element) -> {
                UUID id = element.getUuid();
@ -145,7 +145,7 @@ public class DefaultTbEdgeConsumerService extends AbstractConsumerService<ToEdge
                TbCallback callback = new TbPackCallback<>(id, ctx);
                try {
                    ToEdgeMsg toEdgeMsg = msg.getValue();
-                    pendingMsgHolder.setToEdgeMsg(toEdgeMsg);
+                    pendingMsgHolder.setMsg(toEdgeMsg);
                    if (toEdgeMsg.hasEdgeNotificationMsg()) {
                        pushNotificationToEdge(toEdgeMsg.getEdgeNotificationMsg(), 0, packProcessingRetries, callback);
                    }
@ -161,20 +161,13 @@ public class DefaultTbEdgeConsumerService extends AbstractConsumerService<ToEdge
        if (!processingTimeoutLatch.await(packProcessingTimeout, TimeUnit.MILLISECONDS)) {
            if (!submitFuture.isDone()) {
                submitFuture.cancel(true);
-                ToEdgeMsg lastSubmitMsg = pendingMsgHolder.getToEdgeMsg();
-                log.info("Timeout to process message: {}", lastSubmitMsg);
+                log.info("Timeout to process message: {}", pendingMsgHolder.getMsg());
            }
            ctx.getFailedMap().forEach((id, msg) -> log.warn("[{}] Failed to process message: {}", id, msg.getValue()));
        }
        consumer.commit();
    }

-    private static class PendingMsgHolder {
-        @Getter
-        @Setter
-        private volatile ToEdgeMsg toEdgeMsg;
-    }
-
    @Override
    protected ServiceType getServiceType() {
        return ServiceType.TB_CORE;
--- a/application/src/main/java/org/thingsboard/server/service/queue/PendingMsgHolder.java
+++ b/application/src/main/java/org/thingsboard/server/service/queue/PendingMsgHolder.java
@ -0,0 +1,24 @@
+/**
+ * Copyright © 2016-2024 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 lombok.Getter;
+import lombok.Setter;
+
+public class PendingMsgHolder<T> {
+    @Getter @Setter
+    private volatile T msg;
+}