diff --git a/Common/src/Pollster/TaskHandler.cs b/Common/src/Pollster/TaskHandler.cs
index 9ae15eb1f..846e20c9b 100644
--- a/Common/src/Pollster/TaskHandler.cs
+++ b/Common/src/Pollster/TaskHandler.cs
@@ -1,17 +1,17 @@
// This file is part of the ArmoniK project
-//
+//
// Copyright (C) ANEO, 2021-2024. All rights reserved.
-//
+//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
-//
+//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY, without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
-//
+//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see .
@@ -190,7 +190,7 @@ await messageHandler_.DisposeIgnoreErrorAsync(logger_)
///
/// Task representing the asynchronous execution of the method
///
- public async Task StopCancelledTask()
+ public async Task StopCancelledTask()
{
using var measure = functionExecutionMetrics_.CountAndTime();
using var activity = activitySource_.StartActivityFromParent(activityContext_,
@@ -214,8 +214,12 @@ await lateCts_.CancelAsync()
messageHandler_.Status = QueueMessageStatus.Cancelled;
await ReleaseTaskHandler()
.ConfigureAwait(false);
+
+ return true;
}
}
+
+ return false;
}
///
@@ -766,6 +770,23 @@ await taskTable_.StartTask(taskData_,
"Task already in a final state, removing it from the queue");
throw;
}
+ // If an ArmoniKException is thrown, check if this was because
+ // the task has been canceled. Use standard error management otherwise.
+ catch (ArmoniKException e)
+ {
+ if (await StopCancelledTask()
+ .ConfigureAwait(false))
+ {
+ earlyCts_.Token.ThrowIfCancellationRequested();
+ }
+ else
+ {
+ await HandleErrorRequeueAsync(e,
+ taskData_,
+ earlyCts_.Token)
+ .ConfigureAwait(false);
+ }
+ }
catch (Exception e)
{
await HandleErrorRequeueAsync(e,
diff --git a/Common/tests/Pollster/TaskHandlerTest.cs b/Common/tests/Pollster/TaskHandlerTest.cs
index 032c18999..77a3fd07d 100644
--- a/Common/tests/Pollster/TaskHandlerTest.cs
+++ b/Common/tests/Pollster/TaskHandlerTest.cs
@@ -1542,7 +1542,7 @@ await testServiceProvider.TaskTable.GetTaskStatus(taskId)
}
[Test]
- public async Task CancelLongTaskShouldSucceed()
+ public async Task CancelLongTaskShouldSucceed([Values] bool before)
{
var sqmh = new SimpleQueueMessageHandler
{
@@ -1573,7 +1573,9 @@ public async Task CancelLongTaskShouldSucceed()
await testServiceProvider.TaskHandler.PreProcessing()
.ConfigureAwait(false);
- var exec = testServiceProvider.TaskHandler.ExecuteTask();
+ var exec = before
+ ? testServiceProvider.TaskHandler.ExecuteTask()
+ : Task.CompletedTask;
// Cancel task for test
@@ -1587,6 +1589,11 @@ await testServiceProvider.TaskTable.CancelTaskAsync(new List
CancellationToken.None)
.ConfigureAwait(false);
+ if (!before)
+ {
+ exec = testServiceProvider.TaskHandler.ExecuteTask();
+ }
+
await Task.Delay(TimeSpan.FromMilliseconds(200))
.ConfigureAwait(false);
@@ -1600,7 +1607,8 @@ await testServiceProvider.TaskHandler.StopCancelledTask()
await testServiceProvider.TaskHandler.StopCancelledTask()
.ConfigureAwait(false);
- Assert.ThrowsAsync(() => exec);
+ Assert.That(() => exec,
+ Throws.InstanceOf());
Assert.AreEqual(TaskStatus.Cancelling,
await testServiceProvider.TaskTable.GetTaskStatus(taskId)