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)