From 1b565e01ce5caa1c652758a8be82943806e04194 Mon Sep 17 00:00:00 2001
From: Sergey Matvienko <smatvienko@thingsboard.io>
Date: Tue, 25 Jul 2023 06:54:48 +0200
Subject: [PATCH] Rule engine: ack all rate limited failures (draft)

---
 .../service/queue/TbMsgPackCallback.java       | 18 ++++++++++++++++++
 .../server/common/msg/queue/TbMsgCallback.java |  4 ++++
 2 files changed, 22 insertions(+)

diff --git a/application/src/main/java/org/thingsboard/server/service/queue/TbMsgPackCallback.java b/application/src/main/java/org/thingsboard/server/service/queue/TbMsgPackCallback.java
index ef2ba8798d..66364406c2 100644
--- a/application/src/main/java/org/thingsboard/server/service/queue/TbMsgPackCallback.java
+++ b/application/src/main/java/org/thingsboard/server/service/queue/TbMsgPackCallback.java
@@ -17,11 +17,13 @@ package org.thingsboard.server.service.queue;
 
 import io.micrometer.core.instrument.Timer;
 import lombok.extern.slf4j.Slf4j;
+import org.thingsboard.server.common.data.exception.ApiUsageLimitsExceededException;
 import org.thingsboard.server.common.data.id.RuleNodeId;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.msg.queue.RuleEngineException;
 import org.thingsboard.server.common.msg.queue.RuleNodeInfo;
 import org.thingsboard.server.common.msg.queue.TbMsgCallback;
+import org.thingsboard.server.common.msg.tools.TbRateLimitsException;
 
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
@@ -57,8 +59,24 @@ public class TbMsgPackCallback implements TbMsgCallback {
         ctx.onSuccess(id);
     }
 
+    @Override
+    public void onRateLimit(RuleEngineException e) {
+        log.debug("[{}] ON RATE LIMIT", id, e);
+        //TODO notify tenant on rate limit
+        if (failedMsgTimer != null) {
+            failedMsgTimer.record(System.currentTimeMillis() - startMsgProcessing, TimeUnit.MILLISECONDS);
+        }
+        ctx.onSuccess(id);
+    }
+    
     @Override
     public void onFailure(RuleEngineException e) {
+        Throwable cause = e.getCause();
+        if (cause instanceof TbRateLimitsException || cause instanceof ApiUsageLimitsExceededException) {
+            onRateLimit(e);
+            return;
+        }
+
         log.trace("[{}] ON FAILURE", id, e);
         if (failedMsgTimer != null) {
             failedMsgTimer.record(System.currentTimeMillis() - startMsgProcessing, TimeUnit.MILLISECONDS);
diff --git a/common/message/src/main/java/org/thingsboard/server/common/msg/queue/TbMsgCallback.java b/common/message/src/main/java/org/thingsboard/server/common/msg/queue/TbMsgCallback.java
index 3312c98b64..6cf298adb2 100644
--- a/common/message/src/main/java/org/thingsboard/server/common/msg/queue/TbMsgCallback.java
+++ b/common/message/src/main/java/org/thingsboard/server/common/msg/queue/TbMsgCallback.java
@@ -39,6 +39,10 @@ public interface TbMsgCallback {
 
     void onFailure(RuleEngineException e);
 
+    default void onRateLimit(RuleEngineException e) {
+        onFailure(e);
+    };
+
     /**
      * Returns 'true' if rule engine is expecting the message to be processed, 'false' otherwise.
      * message may no longer be valid, if the message pack is already expired/canceled/failed.