Fix StartProcessing if cancellation is requested before

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 <http://www.gnu.org/licenses/>.
 
@@ -190,7 +190,7 @@ await messageHandler_.DisposeIgnoreErrorAsync(logger_)
   /// <returns>
   ///   Task representing the asynchronous execution of the method
   /// </returns>
-  public async Task StopCancelledTask()
+  public async Task<bool> 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;
   }
 
   /// <summary>
@@ -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,

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<string>
                                                         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<TaskCanceledException>(() => exec);
+    Assert.That(() => exec,
+                Throws.InstanceOf<OperationCanceledException>());
 
     Assert.AreEqual(TaskStatus.Cancelling,
                     await testServiceProvider.TaskTable.GetTaskStatus(taskId)