feat: Add endpoint on Polling Agent to kill the current task if task is in cancelling status

Common/src/Pollster/Pollster.cs

@@ -67,6 +67,7 @@ public class Pollster : IInitializable
  private readonly IWorkerStreamHandler workerStreamHandler_;
  private bool endLoopReached_;
  private HealthCheckResult? healthCheckFailedResult_;
+ public Func<Task>? StopCancelledTask;
  public string TaskProcessing;
 
  public Pollster(IPullQueueStorage pullQueueStorage,
@@ -295,6 +296,8 @@ void RecordError(Exception e)
  pollsterOptions_,
  cts);
 
+ StopCancelledTask = taskHandler.StopCancelledTask;
+
  var precondition = await taskHandler.AcquireTask()
  .ConfigureAwait(false);
 
@@ -311,6 +314,8 @@ await taskHandler.ExecuteTask()
  await taskHandler.PostProcessing()
  .ConfigureAwait(false);
 
+ StopCancelledTask = null;
+
  logger_.LogDebug("Task returned");
 
  // If the task was successful, we can remove a failure
@@ -331,7 +336,8 @@ await taskHandler.PostProcessing()
  }
  finally
  {
- TaskProcessing = string.Empty;
+ StopCancelledTask = null;
+ TaskProcessing = string.Empty;
  }
  }
  }

Common/src/Pollster/TaskHandler.cs

@@ -87,33 +87,35 @@ public TaskHandler(ISessionTable sessionTable,
  Injection.Options.Pollster pollsterOptions,
  CancellationTokenSource cancellationTokenSource)
  {
- sessionTable_ = sessionTable;
- taskTable_ = taskTable;
- resultTable_ = resultTable;
- messageHandler_ = messageHandler;
- taskProcessingChecker_ = taskProcessingChecker;
- submitter_ = submitter;
- dataPrefetcher_ = dataPrefetcher;
- workerStreamHandler_ = workerStreamHandler;
- activitySource_ = activitySource;
- agentHandler_ = agentHandler;
- logger_ = logger;
- cancellationTokenSource_ = cancellationTokenSource;
- ownerPodId_ = ownerPodId;
- ownerPodName_ = ownerPodName;
- taskData_ = null;
- sessionData_ = null;
+ sessionTable_ = sessionTable;
+ taskTable_ = taskTable;
+ resultTable_ = resultTable;
+ messageHandler_ = messageHandler;
+ taskProcessingChecker_ = taskProcessingChecker;
+ submitter_ = submitter;
+ dataPrefetcher_ = dataPrefetcher;
+ workerStreamHandler_ = workerStreamHandler;
+ activitySource_ = activitySource;
+ agentHandler_ = agentHandler;
+ logger_ = logger;
+ ownerPodId_ = ownerPodId;
+ ownerPodName_ = ownerPodName;
+ taskData_ = null;
+ sessionData_ = null;
  token_ = Guid.NewGuid()
  .ToString();
 
- workerConnectionCts_ = new CancellationTokenSource();
+ workerConnectionCts_ = new CancellationTokenSource();
+ cancellationTokenSource_ = new CancellationTokenSource();
 
- reg1_ = cancellationTokenSource_.Token.Register(() =>
- {
- logger_.LogWarning("Cancellation triggered, waiting {waitingTime} before cancelling task",
- pollsterOptions.GraceDelay);
- workerConnectionCts_.CancelAfter(pollsterOptions.GraceDelay);
- });
+ reg1_ = cancellationTokenSource.Token.Register(() => cancellationTokenSource_.Cancel());
+
+ cancellationTokenSource_.Token.Register(() =>
+ {
+ logger_.LogWarning("Cancellation triggered, waiting {waitingTime} before cancelling task",
+ pollsterOptions.GraceDelay);
+ workerConnectionCts_.CancelAfter(pollsterOptions.GraceDelay);
+ });
  workerConnectionCts_.Token.Register(() => logger_.LogWarning("Cancellation triggered, start to properly cancel task"));
  }
 
@@ -131,10 +133,32 @@ await messageHandler_.DisposeAsync()
  reg1_.Unregister();
  await reg1_.DisposeAsync()
  .ConfigureAwait(false);
+ cancellationTokenSource_.Dispose();
  workerConnectionCts_.Dispose();
  agent_?.Dispose();
  }
 
+ /// <summary>
+ /// Refresh task metadata and stop execution if current task should be cancelled
+ /// </summary>
+ /// <returns>
+ /// Task representing the asynchronous execution of the method
+ /// </returns>
+ public async Task StopCancelledTask()
+ {
+ if (taskData_?.Status is not null or TaskStatus.Cancelled or TaskStatus.Cancelling)
+ {
+ taskData_ = await taskTable_.ReadTaskAsync(messageHandler_.TaskId,
+ CancellationToken.None)
+ .ConfigureAwait(false);
+ if (taskData_.Status is TaskStatus.Cancelling)
+ {
+ logger_.LogWarning("Task has been cancelled, trigger cancellation from exterior.");
+ cancellationTokenSource_.Cancel();
+ }
+ }
+ }
+
  /// <summary>
  /// Acquisition of the task in the message given to the constructor
  /// </summary>
@@ -618,6 +642,12 @@ private async Task HandleErrorInternalAsync(Exception e,
  bool requeueIfUnavailable,
  CancellationToken cancellationToken)
  {
+ if (taskData.Status is TaskStatus.Cancelled or TaskStatus.Cancelling)
+ {
+ messageHandler_.Status = QueueMessageStatus.Processed;
+ return;
+ }
+
  if (cancellationToken.IsCancellationRequested || (requeueIfUnavailable && e is RpcException
  {
  StatusCode: StatusCode.Unavailable,

Common/tests/Helpers/ExceptionAsyncPipe.cs

@@ -0,0 +1,60 @@
+// This file is part of the ArmoniK project
+// 
+// Copyright (C) ANEO, 2021-2023. All rights reserved.
+// W. Kirschenmann <[email protected]>
+// J. Gurhem <[email protected]>
+// D. Dubuc <[email protected]>
+// L. Ziane Khodja <[email protected]>
+// F. Lemaitre <[email protected]>
+// S. Djebbar <[email protected]>
+// J. Fonseca <[email protected]>
+// D. Brasseur <[email protected]>
+// 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/>.
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+using ArmoniK.Api.gRPC.V1.Worker;
+using ArmoniK.Core.Common.Utils;
+
+namespace ArmoniK.Core.Common.Tests.Helpers;
+
+public class ExceptionAsyncPipe<T> : IAsyncPipe<ProcessReply, ProcessRequest>
+ where T : Exception, new()
+{
+ private readonly int delay_;
+
+ public ExceptionAsyncPipe(int delay)
+ => delay_ = delay;
+
+ public async Task<ProcessReply> ReadAsync(CancellationToken cancellationToken)
+ {
+ await Task.Delay(TimeSpan.FromMilliseconds(delay_),
+ cancellationToken)
+ .ConfigureAwait(false);
+ cancellationToken.ThrowIfCancellationRequested();
+ throw new T();
+ }
+
+ public Task WriteAsync(ProcessRequest message)
+ => Task.CompletedTask;
+
+ public Task WriteAsync(IEnumerable<ProcessRequest> message)
+ => Task.CompletedTask;
+
+ public Task CompleteAsync()
+ => Task.CompletedTask;
+}

Common/tests/Helpers/ExceptionWorkerStreamHandler.cs

@@ -0,0 +1,61 @@
+// This file is part of the ArmoniK project
+// 
+// Copyright (C) ANEO, 2021-2023. All rights reserved.
+// W. Kirschenmann <[email protected]>
+// J. Gurhem <[email protected]>
+// D. Dubuc <[email protected]>
+// L. Ziane Khodja <[email protected]>
+// F. Lemaitre <[email protected]>
+// S. Djebbar <[email protected]>
+// J. Fonseca <[email protected]>
+// D. Brasseur <[email protected]>
+// 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/>.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+using ArmoniK.Api.gRPC.V1.Worker;
+using ArmoniK.Core.Common.Storage;
+using ArmoniK.Core.Common.Stream.Worker;
+using ArmoniK.Core.Common.Utils;
+
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+
+namespace ArmoniK.Core.Common.Tests.Helpers;
+
+public class ExceptionWorkerStreamHandler<T> : IWorkerStreamHandler
+ where T : Exception, new()
+{
+ private readonly int delay_;
+
+ public ExceptionWorkerStreamHandler(int delay)
+ => delay_ = delay;
+
+ public Task<HealthCheckResult> Check(HealthCheckTag tag)
+ => Task.FromResult(HealthCheckResult.Healthy());
+
+ public Task Init(CancellationToken cancellationToken)
+ => Task.CompletedTask;
+
+ public void Dispose()
+ {
+ }
+
+ public IAsyncPipe<ProcessReply, ProcessRequest>? Pipe { get; private set; }
+
+ public void StartTaskProcessing(TaskData taskData,
+ CancellationToken cancellationToken)
+ => Pipe = new ExceptionAsyncPipe<T>(delay_);
+}

Common/tests/Pollster/PollsterTest.cs

@@ -501,6 +501,79 @@ await testServiceProvider.Pollster.Init(CancellationToken.None)
  testServiceProvider.Pollster.TaskProcessing);
  }
 
+ [Test]
+ public async Task CancelLongTaskShouldSucceed()
+ {
+ var mockPullQueueStorage = new Mock<IPullQueueStorage>();
+ var waitWorkerStreamHandler = new ExceptionWorkerStreamHandler<Exception>(15000);
+ var simpleAgentHandler = new SimpleAgentHandler();
+
+ using var testServiceProvider = new TestPollsterProvider(waitWorkerStreamHandler,
+ simpleAgentHandler,
+ mockPullQueueStorage.Object);
+
+ var tuple = await InitSubmitter(testServiceProvider.Submitter,
+ testServiceProvider.PartitionTable,
+ CancellationToken.None)
+ .ConfigureAwait(false);
+
+ mockPullQueueStorage.Setup(storage => storage.PullMessagesAsync(It.IsAny<int>(),
+ It.IsAny<CancellationToken>()))
+ .Returns(() => new List<IQueueMessageHandler>
+ {
+ new SimpleQueueMessageHandler
+ {
+ CancellationToken = CancellationToken.None,
+ Status = QueueMessageStatus.Waiting,
+ MessageId = Guid.NewGuid()
+ .ToString(),
+ TaskId = tuple.taskSubmitted,
+ },
+ }.ToAsyncEnumerable());
+
+ await testServiceProvider.Pollster.Init(CancellationToken.None)
+ .ConfigureAwait(false);
+
+ var source = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+
+ var mainLoopTask = testServiceProvider.Pollster.MainLoop(source.Token);
+
+ await Task.Delay(TimeSpan.FromMilliseconds(200),
+ CancellationToken.None)
+ .ConfigureAwait(false);
+
+ await testServiceProvider.TaskTable.CancelTaskAsync(new List<string>
+ {
+ tuple.taskSubmitted,
+ },
+ CancellationToken.None)
+ .ConfigureAwait(false);
+
+ await Task.Delay(TimeSpan.FromMilliseconds(200),
+ CancellationToken.None)
+ .ConfigureAwait(false);
+
+ await testServiceProvider.Pollster.StopCancelledTask!.Invoke()
+ .ConfigureAwait(false);
+
+ Assert.DoesNotThrowAsync(() => mainLoopTask);
+ Assert.False(testServiceProvider.Pollster.Failed);
+ Assert.True(source.Token.IsCancellationRequested);
+
+ Assert.AreEqual(TaskStatus.Cancelled,
+ (await testServiceProvider.TaskTable.GetTaskStatus(new[]
+ {
+ tuple.taskSubmitted,
+ },
+ CancellationToken.None)
+ .ConfigureAwait(false)).Single()
+ .Status);
+ Assert.AreEqual(string.Empty,
+ testServiceProvider.Pollster.TaskProcessing);
+ Assert.AreSame(string.Empty,
+ testServiceProvider.Pollster.TaskProcessing);
+ }
+
  public static IEnumerable ExecuteTooManyErrorShouldFailTestCase
  {
  get