Cancel currently running tasks when the job is discarded

application/src/main/java/org/thingsboard/server/service/cf/ctx/state/KafkaCalculatedFieldStateService.java

@@ -73,7 +73,7 @@ public class KafkaCalculatedFieldStateService extends AbstractCalculatedFieldSta
                 .queueKey(queueKey)
                 .topic(partitionService.getTopic(queueKey))
                 .pollInterval(pollInterval)
-                .msgPackProcessor((msgs, consumer, config) -> {
+                .msgPackProcessor((msgs, consumer, consumerKey, config) -> {
                     for (TbProtoQueueMsg<CalculatedFieldStateProto> msg : msgs) {
                         try {
                             if (msg.getValue() != null) {

application/src/main/java/org/thingsboard/server/service/job/DummyJobProcessor.java

@@ -82,6 +82,7 @@ public class DummyJobProcessor implements JobProcessor {
                 .processingTimeMs(configuration.getTaskProcessingTimeMs())
                 .errors(errors)
                 .failAlways(failAlways)
+                .processingTimeoutMs(configuration.getTaskProcessingTimeoutMs())
                 .build();
     }
 

application/src/main/java/org/thingsboard/server/service/job/task/DummyTaskProcessor.java

@@ -15,13 +15,16 @@
  */
 package org.thingsboard.server.service.job.task;
 
-import lombok.RequiredArgsConstructor;
+import org.apache.commons.lang3.tuple.Pair;
 import org.thingsboard.server.common.data.job.JobType;
 import org.thingsboard.server.common.data.job.task.DummyTask;
 import org.thingsboard.server.common.data.job.task.DummyTaskResult;
+import org.thingsboard.server.common.data.job.task.Task;
 import org.thingsboard.server.queue.task.TaskProcessor;
 
-@RequiredArgsConstructor
+import java.util.Map;
+import java.util.concurrent.Future;
+
 public class DummyTaskProcessor extends TaskProcessor<DummyTask, DummyTaskResult> {
 
     @Override
@@ -40,8 +43,12 @@ public class DummyTaskProcessor extends TaskProcessor<DummyTask, DummyTaskResult
     }
 
     @Override
-    public long getTaskProcessingTimeout() {
-        return 2000;
+    public long getProcessingTimeout(DummyTask task) {
+        return task.getProcessingTimeoutMs() > 0 ? task.getProcessingTimeoutMs() : 2000;
+    }
+
+    public Map<Object, Pair<Task<DummyTaskResult>, Future<DummyTaskResult>>> getCurrentTasks() {
+        return currentTasks;
     }
 
     @Override

application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCalculatedFieldConsumerService.java

@@ -141,7 +141,7 @@ public class DefaultTbCalculatedFieldConsumerService extends AbstractPartitionBa
         }
     }
 
-    private void processMsgs(List<TbProtoQueueMsg<ToCalculatedFieldMsg>> msgs, TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> consumer, QueueConfig config) throws Exception {
+    private void processMsgs(List<TbProtoQueueMsg<ToCalculatedFieldMsg>> msgs, TbQueueConsumer<TbProtoQueueMsg<ToCalculatedFieldMsg>> consumer, Object consumerKey, QueueConfig 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));

application/src/main/java/org/thingsboard/server/service/queue/DefaultTbCoreConsumerService.java

@@ -255,7 +255,7 @@ public class DefaultTbCoreConsumerService extends AbstractConsumerService<ToCore
                 .collect(Collectors.toSet()));
     }
 
-    private void processMsgs(List<TbProtoQueueMsg<ToCoreMsg>> msgs, TbQueueConsumer<TbProtoQueueMsg<ToCoreMsg>> consumer, QueueConfig config) throws Exception {
+    private void processMsgs(List<TbProtoQueueMsg<ToCoreMsg>> msgs, TbQueueConsumer<TbProtoQueueMsg<ToCoreMsg>> consumer, Object consumerKey, QueueConfig config) throws Exception {
         List<IdMsgPair<ToCoreMsg>> orderedMsgList = msgs.stream().map(msg -> new IdMsgPair<>(UUID.randomUUID(), msg)).toList();
         ConcurrentMap<UUID, TbProtoQueueMsg<ToCoreMsg>> pendingMap = orderedMsgList.stream().collect(
                 Collectors.toConcurrentMap(IdMsgPair::getUuid, IdMsgPair::getMsg));

application/src/main/java/org/thingsboard/server/service/queue/DefaultTbEdgeConsumerService.java

@@ -126,7 +126,7 @@ public class DefaultTbEdgeConsumerService extends AbstractConsumerService<ToEdge
         mainConsumer.update(partitions);
     }
 
-    private void processMsgs(List<TbProtoQueueMsg<ToEdgeMsg>> msgs, TbQueueConsumer<TbProtoQueueMsg<ToEdgeMsg>> consumer, QueueConfig edgeQueueConfig) throws InterruptedException {
+    private void processMsgs(List<TbProtoQueueMsg<ToEdgeMsg>> msgs, TbQueueConsumer<TbProtoQueueMsg<ToEdgeMsg>> consumer, Object consumerKey, QueueConfig edgeQueueConfig) throws InterruptedException {
         List<IdMsgPair<ToEdgeMsg>> orderedMsgList = msgs.stream().map(msg -> new IdMsgPair<>(UUID.randomUUID(), msg)).toList();
         ConcurrentMap<UUID, TbProtoQueueMsg<ToEdgeMsg>> pendingMap = orderedMsgList.stream().collect(
                 Collectors.toConcurrentMap(IdMsgPair::getUuid, IdMsgPair::getMsg));

application/src/main/java/org/thingsboard/server/service/queue/ruleengine/TbRuleEngineQueueConsumerManager.java

@@ -127,6 +127,7 @@ public class TbRuleEngineQueueConsumerManager extends MainQueueConsumerManager<T
     @Override
     protected void processMsgs(List<TbProtoQueueMsg<ToRuleEngineMsg>> msgs,
                                TbQueueConsumer<TbProtoQueueMsg<ToRuleEngineMsg>> consumer,
+                               Object consumerKey,
                                Queue queue) throws Exception {
         TbRuleEngineSubmitStrategy submitStrategy = getSubmitStrategy(queue);
         TbRuleEngineProcessingStrategy ackStrategy = getProcessingStrategy(queue);

application/src/test/java/org/thingsboard/server/service/job/JobManagerTest.java

@@ -172,6 +172,32 @@ public class JobManagerTest extends AbstractControllerTest {
         });
     }
 
+    @Test
+    public void testCancelJob_whileTaskRunning() throws Exception {
+        JobId jobId = submitJob(DummyJobConfiguration.builder()
+                .successfulTasksCount(1)
+                .taskProcessingTimeMs(TimeUnit.HOURS.toMillis(1))
+                .taskProcessingTimeoutMs(TimeUnit.HOURS.toMillis(1))
+                .build()).getId();
+        await().atMost(TIMEOUT, TimeUnit.SECONDS).untilAsserted(() -> {
+            assertThat(taskProcessor.getCurrentTasks()).isNotEmpty();
+            Job job = findJobById(jobId);
+            assertThat(job.getStatus()).isEqualTo(JobStatus.RUNNING);
+            assertThat(job.getResult().getTotalCount()).isEqualTo(1);
+            assertThat(job.getResult().getCompletedCount()).isZero();
+        });
+
+        cancelJob(jobId);
+
+        await().atMost(TIMEOUT, TimeUnit.SECONDS).untilAsserted(() -> {
+            Job job = findJobById(jobId);
+            assertThat(job.getStatus()).isEqualTo(JobStatus.CANCELLED);
+            assertThat(job.getResult().getTotalCount()).isEqualTo(1);
+            assertThat(job.getResult().getDiscardedCount()).isEqualTo(1);
+            assertThat(job.getResult().getFailedCount()).isZero();
+        });
+    }
+
     @Test
     public void testCancelJob_simulateTaskProcessorRestart() throws Exception {
         int tasksCount = 10;

application/src/test/java/org/thingsboard/server/service/queue/ruleengine/TbRuleEngineStrategyTest.java

@@ -238,7 +238,7 @@ public class TbRuleEngineStrategyTest {
                 .map(this::toProto)
                 .toList();
 
-        consumerManager.processMsgs(protoMsgs, consumer, queue);
+        consumerManager.processMsgs(protoMsgs, consumer, queueKey, queue);
 
         processingData.forEach(data -> {
             verify(actorContext, times(data.attempts)).tell(argThat(msg ->

common/data/src/main/java/org/thingsboard/server/common/data/job/DummyJobConfiguration.java

@@ -38,6 +38,7 @@ public class DummyJobConfiguration extends JobConfiguration {
     private int permanentlyFailedTasksCount;
     private List<String> errors;
     private int retries;
+    private long taskProcessingTimeoutMs;
 
     private String generalError;
     private int submittedTasksBeforeGeneralError;

common/data/src/main/java/org/thingsboard/server/common/data/job/task/DummyTask.java

@@ -36,6 +36,7 @@ public class DummyTask extends Task<DummyTaskResult> {
 
     private int number;
     private long processingTimeMs;
+    private long processingTimeoutMs;
     private List<String> errors; // errors for each attempt
     private boolean failAlways;
 

common/edqs/src/main/java/org/thingsboard/server/edqs/processor/EdqsProcessor.java

@@ -120,7 +120,7 @@ public class EdqsProcessor implements TbQueueHandler<TbProtoQueueMsg<ToEdqsMsg>,
                 .queueKey(new QueueKey(ServiceType.EDQS, config.getEventsTopic()))
                 .topic(topicService.buildTopicName(config.getEventsTopic()))
                 .pollInterval(config.getPollInterval())
-                .msgPackProcessor((msgs, consumer, config) -> {
+                .msgPackProcessor((msgs, consumer, consumerKey, config) -> {
                     for (TbProtoQueueMsg<ToEdqsMsg> queueMsg : msgs) {
                         if (consumer.isStopped()) {
                             return;

common/edqs/src/main/java/org/thingsboard/server/edqs/state/KafkaEdqsStateService.java

@@ -91,7 +91,7 @@ public class KafkaEdqsStateService implements EdqsStateService {
                 .queueKey(new QueueKey(ServiceType.EDQS, config.getStateTopic()))
                 .topic(topicService.buildTopicName(config.getStateTopic()))
                 .pollInterval(config.getPollInterval())
-                .msgPackProcessor((msgs, consumer, config) -> {
+                .msgPackProcessor((msgs, consumer, consumerKey, config) -> {
                     for (TbProtoQueueMsg<ToEdqsMsg> queueMsg : msgs) {
                         try {
                             ToEdqsMsg msg = queueMsg.getValue();

common/queue/src/main/java/org/thingsboard/server/queue/common/PartitionedQueueResponseTemplate.java

@@ -74,7 +74,7 @@ public class PartitionedQueueResponseTemplate<Request extends TbQueueMsg, Respon
                 .queueKey(key + "-requests")
                 .topic(requestsTopic)
                 .pollInterval(pollInterval)
-                .msgPackProcessor((requests, consumer, config) -> processRequests(requests, consumer))
+                .msgPackProcessor((requests, consumer, consumerKey, config) -> processRequests(requests, consumer))
                 .consumerCreator((config, tpi) -> consumerCreator.apply(tpi))
                 .consumerExecutor(consumerExecutor)
                 .scheduler(scheduler)

common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/MainQueueConsumerManager.java

@@ -203,7 +203,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
         log.info("[{}] Launching consumer", consumerTask.getKey());
         Future<?> consumerLoop = consumerExecutor.submit(() -> {
             ThingsBoardThreadFactory.updateCurrentThreadName(consumerTask.getKey().toString());
-            consumerLoop(consumerTask.getConsumer());
+            consumerLoop(consumerTask.getKey(), consumerTask.getConsumer());
             log.info("[{}] Consumer stopped", consumerTask.getKey());
 
             try {
@@ -218,7 +218,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
         consumerTask.setTask(consumerLoop);
     }
 
-    private void consumerLoop(TbQueueConsumer<M> consumer) {
+    private void consumerLoop(Object consumerKey, TbQueueConsumer<M> consumer) {
         try {
             while (!stopped && !consumer.isStopped()) {
                 try {
@@ -226,7 +226,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
                     if (msgs.isEmpty()) {
                         continue;
                     }
-                    processMsgs(msgs, consumer, config);
+                    processMsgs(msgs, consumer, consumerKey, config);
                 } catch (Exception e) {
                     if (!consumer.isStopped()) {
                         log.warn("Failed to process messages from queue", e);
@@ -250,9 +250,9 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
         }
     }
 
-    protected void processMsgs(List<M> msgs, TbQueueConsumer<M> consumer, C config) throws Exception {
+    protected void processMsgs(List<M> msgs, TbQueueConsumer<M> consumer, Object consumerKey, C config) throws Exception {
         log.trace("Processing {} messages", msgs.size());
-        msgPackProcessor.process(msgs, consumer, config);
+        msgPackProcessor.process(msgs, consumer, consumerKey, config);
         log.trace("Processed {} messages", msgs.size());
     }
 
@@ -273,7 +273,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
     }
 
     public interface MsgPackProcessor<M extends TbQueueMsg, C extends QueueConfig> {
-        void process(List<M> msgs, TbQueueConsumer<M> consumer, C config) throws Exception;
+        void process(List<M> msgs, TbQueueConsumer<M> consumer, Object consumerKey, C config) throws Exception;
     }
 
     public interface ConsumerWrapper<M extends TbQueueMsg> {

common/queue/src/main/java/org/thingsboard/server/queue/task/TaskProcessor.java

@@ -17,6 +17,7 @@ package org.thingsboard.server.queue.task;
 
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
+import org.apache.commons.lang3.tuple.Pair;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -25,6 +26,7 @@ import org.thingsboard.common.util.JacksonUtil;
 import org.thingsboard.common.util.SetCache;
 import org.thingsboard.common.util.ThingsBoardThreadFactory;
 import org.thingsboard.server.common.data.id.EntityId;
+import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.job.JobType;
 import org.thingsboard.server.common.data.job.task.Task;
 import org.thingsboard.server.common.data.job.task.TaskResult;
@@ -42,14 +44,18 @@ import org.thingsboard.server.queue.discovery.event.PartitionChangeEvent;
 import org.thingsboard.server.queue.settings.TasksQueueConfig;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Predicate;
 
 public abstract class TaskProcessor<T extends Task<R>, R extends TaskResult> {
 
@@ -68,6 +74,8 @@ public abstract class TaskProcessor<T extends Task<R>, R extends TaskResult> {
     private MainQueueConsumerManager<TbProtoQueueMsg<TaskProto>, QueueConfig> taskConsumer;
     private final ExecutorService taskExecutor = Executors.newCachedThreadPool(ThingsBoardThreadFactory.forName(getJobType().name().toLowerCase() + "-task-processor"));
 
+    protected final Map<Object, Pair<Task<R>, Future<R>>> currentTasks = new ConcurrentHashMap<>();
+
     private final SetCache<String> discarded = new SetCache<>(TimeUnit.MINUTES.toMillis(60));
     private final SetCache<String> failed = new SetCache<>(TimeUnit.MINUTES.toMillis(60));
 
@@ -104,21 +112,24 @@ public abstract class TaskProcessor<T extends Task<R>, R extends TaskResult> {
                 if (event.getEvent() == ComponentLifecycleEvent.STOPPED) {
                     log.info("Adding job {} ({}) to discarded", entityId, tasksKey);
                     addToDiscarded(tasksKey);
+                    cancelRunningTasks(tasksKey);
                 } else if (event.getEvent() == ComponentLifecycleEvent.FAILED) {
                     log.info("Adding job {} ({}) to failed", entityId, tasksKey);
                     failed.add(tasksKey);
+                    cancelRunningTasks(tasksKey);
                 }
             }
             case TENANT -> {
                 if (event.getEvent() == ComponentLifecycleEvent.DELETED) {
-                    deletedTenants.add(entityId.getId());
                     log.info("Adding tenant {} to deleted", entityId);
+                    deletedTenants.add(entityId.getId());
+                    cancelRunningTasks((TenantId) entityId);
                 }
             }
         }
     }
 
-    private void processMsgs(List<TbProtoQueueMsg<TaskProto>> msgs, TbQueueConsumer<TbProtoQueueMsg<TaskProto>> consumer, QueueConfig queueConfig) throws Exception {
+    private void processMsgs(List<TbProtoQueueMsg<TaskProto>> msgs, TbQueueConsumer<TbProtoQueueMsg<TaskProto>> consumer, Object consumerKey, QueueConfig queueConfig) throws Exception {
         for (TbProtoQueueMsg<TaskProto> msg : msgs) {
             try {
                 @SuppressWarnings("unchecked")
@@ -135,7 +146,7 @@ public abstract class TaskProcessor<T extends Task<R>, R extends TaskResult> {
                     continue;
                 }
 
-                processTask(task);
+                processTask(task, consumerKey);
             } catch (InterruptedException e) {
                 throw e;
             } catch (Exception e) {
@@ -145,30 +156,39 @@ public abstract class TaskProcessor<T extends Task<R>, R extends TaskResult> {
         consumer.commit();
     }
 
-    private void processTask(T task) throws InterruptedException {
+    private void processTask(T task, Object consumerKey) throws InterruptedException {
         task.setAttempt(task.getAttempt() + 1);
         log.debug("Processing task: {}", task);
         Future<R> future = null;
         try {
             long startNs = System.nanoTime();
+            long timeoutMs = getProcessingTimeout(task);
+
             future = taskExecutor.submit(() -> process(task));
+            currentTasks.put(consumerKey, Pair.of(task, future));
+
             R result;
             try {
-                result = future.get(getTaskProcessingTimeout(), TimeUnit.MILLISECONDS);
+                result = future.get(timeoutMs, TimeUnit.MILLISECONDS);
             } catch (ExecutionException e) {
                 throw e.getCause();
             } catch (TimeoutException e) {
-                throw new TimeoutException("Timeout after " + getTaskProcessingTimeout() + " ms");
+                throw new TimeoutException("Timeout after " + timeoutMs + " ms");
             }
+
             long timingNs = System.nanoTime() - startNs;
             log.info("Processed task in {} ms: {}", timingNs / 1000000.0, task);
             reportTaskResult(task, result);
         } catch (InterruptedException e) {
             throw e;
+        } catch (CancellationException e) {
+            if (!failed.contains(task.getKey()) && !deletedTenants.contains(task.getTenantId().getId())) {
+                reportTaskDiscarded(task);
+            }
         } catch (Throwable e) {
             log.error("Failed to process task (attempt {}): {}", task.getAttempt(), task, e);
             if (task.getAttempt() <= task.getRetries()) {
-                processTask(task);
+                processTask(task, consumerKey);
             } else {
                 reportTaskFailure(task, e);
             }
@@ -176,11 +196,31 @@ public abstract class TaskProcessor<T extends Task<R>, R extends TaskResult> {
             if (future != null && !future.isDone()) {
                 future.cancel(true);
             }
+            currentTasks.remove(consumerKey);
         }
     }
 
     public abstract R process(T task) throws Exception;
 
+    private void cancelRunningTasks(String tasksKey) {
+        cancelRunningTasks(task -> task.getKey().equals(tasksKey));
+    }
+
+    private void cancelRunningTasks(TenantId tenantId) {
+        cancelRunningTasks(task -> task.getTenantId().equals(tenantId));
+    }
+
+    private void cancelRunningTasks(Predicate<Task<R>> filter) {
+        currentTasks.values().forEach(entry -> {
+            Task<R> task = entry.getKey();
+            Future<R> future = entry.getValue();
+            if (filter.test(task)) {
+                log.debug("Cancelling running task {}", task);
+                future.cancel(true);
+            }
+        });
+    }
+
     private void reportTaskFailure(T task, Throwable error) {
         R taskResult = task.toFailed(error);
         reportTaskResult(task, taskResult);
@@ -215,7 +255,7 @@ public abstract class TaskProcessor<T extends Task<R>, R extends TaskResult> {
         taskExecutor.shutdownNow();
     }
 
-    public abstract long getTaskProcessingTimeout();
+    public abstract long getProcessingTimeout(T task);
 
     public abstract JobType getJobType();