fix: resubmit task when agent tries to acquire a task in Retried

Adaptors/Memory/src/PushQueueStorage.cs

@@ -78,7 +78,7 @@ private Task PushMessagesAsync(IEnumerable<MessageData> messages,
  if (!id2Handlers_.TryAdd(messageHandler.TaskId!,
  messageHandler))
  {
- throw new InvalidOperationException("Duplicate messageId found.");
+ throw new InvalidOperationException($"Duplicate messageId: {messageHandler.TaskId} found.");
  }
  }
 

Adaptors/Memory/src/TaskTable.cs

@@ -64,7 +64,7 @@ public Task CreateTasks(IEnumerable<TaskData> tasks,
  if (!taskId2TaskData_.TryAdd(taskData.TaskId,
  taskData))
  {
- throw new ArmoniKException($"Tasks '{taskData.TaskId}' already exists");
+ throw new TaskAlreadyExistsException($"Tasks '{taskData.TaskId}' already exists");
  }
 
  var session = session2TaskIds_.GetOrAdd(taskData.SessionId,

Adaptors/MongoDB/src/TaskTable.cs

@@ -75,9 +75,17 @@ public async Task CreateTasks(IEnumerable<TaskData> tasks,
 
  var taskCollection = taskCollectionProvider_.Get();
 
- await taskCollection.InsertManyAsync(tasks.Select(taskData => taskData),
- cancellationToken: cancellationToken)
- .ConfigureAwait(false);
+ try
+ {
+ await taskCollection.InsertManyAsync(tasks.Select(taskData => taskData),
+ cancellationToken: cancellationToken)
+ .ConfigureAwait(false);
+ }
+ catch (MongoBulkWriteException<TaskData> e)
+ {
+ throw new TaskAlreadyExistsException("Task already exists",
+ e);
+ }
  }
 
  /// <inheritdoc />

Common/src/Exceptions/TaskAlreadyExistsException.cs

@@ -0,0 +1,40 @@
+// 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/>.
+
+using System;
+
+namespace ArmoniK.Core.Common.Exceptions;
+
+[Serializable]
+public class TaskAlreadyExistsException : ArmoniKException
+{
+ public TaskAlreadyExistsException()
+ {
+ }
+
+ public TaskAlreadyExistsException(string message)
+ : base(message)
+ {
+ }
+
+ public TaskAlreadyExistsException(string message,
+ Exception innerException)
+ : base(message,
+ innerException)
+ {
+ }
+}

Common/src/Pollster/AcquisitionStatus.cs

@@ -77,6 +77,27 @@ public enum AcquisitionStatus
  /// </summary>
  TaskIsRetried,
 
+ /// <summary>
+ /// Task not acquired because its status is <see cref="TaskStatus.Retried" />. Moreover, the retried task is
+ /// <see cref="TaskStatus.Creating" />
+ /// Retried task finalization is required.
+ /// </summary>
+ TaskIsRetriedAndRetryIsCreating,
+
+ /// <summary>
+ /// Task not acquired because its status is <see cref="TaskStatus.Retried" />. Moreover, the retried task was not found
+ /// in the database
+ /// Retried task creation and submission is required.
+ /// </summary>
+ TaskIsRetriedAndRetryIsNotFound,
+
+ /// <summary>
+ /// Task not acquired because its status is <see cref="TaskStatus.Retried" />. Moreover, the retried task is
+ /// <see cref="TaskStatus.Submitted" />
+ /// Reinsertion in the queue may be required.
+ /// </summary>
+ TaskIsRetriedAndRetryIsSubmitted,
+
  /// <summary>
  /// Task not acquired because its status is <see cref="TaskStatus.Processing" /> but the other pod does not seem to be
  /// processing it

Common/src/Pollster/TaskHandler.cs

@@ -16,6 +16,7 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
 using System.Net.Http;
@@ -445,6 +446,79 @@ await submitter_.CompleteTaskAsync(taskData_,
  case TaskStatus.Retried:
  logger_.LogInformation("Task is in retry ; retry task should be executed");
  messageHandler_.Status = QueueMessageStatus.Poisonous;
+ var retryId = taskData_.RetryId();
+
+ try
+ {
+ var retryData = await taskTable_.ReadTaskAsync(retryId,
+ lateCts_.Token)
+ .ConfigureAwait(false);
+ if (retryData.Status is TaskStatus.Creating or TaskStatus.Submitted)
+ {
+ logger_.LogWarning("Retried task {task} is in {status}; trying to finalize task creation",
+ retryId,
+ retryData.Status);
+ await submitter_.FinalizeTaskCreation(new List<TaskCreationRequest>
+ {
+ new(retryId,
+ retryData.PayloadId,
+ retryData.Options,
+ retryData.ExpectedOutputIds,
+ retryData.DataDependencies),
+ },
+ sessionData_,
+ taskData_.TaskId,
+ CancellationToken.None)
+ .ConfigureAwait(false);
+ return retryData.Status == TaskStatus.Submitted
+ ? AcquisitionStatus.TaskIsRetriedAndRetryIsSubmitted
+ : AcquisitionStatus.TaskIsRetriedAndRetryIsCreating;
+ }
+ }
+ catch (TaskNotFoundException)
+ {
+ logger_.LogWarning("Retried task {task} was not found in the database; resubmit it",
+ retryId);
+
+ try
+ {
+ await taskTable_.RetryTask(taskData_,
+ CancellationToken.None)
+ .ConfigureAwait(false);
+ }
+ catch (TaskAlreadyExistsException)
+ {
+ logger_.LogWarning("Retried task {task} already exists; finalize creation if needed",
+ retryId);
+ }
+ finally
+ {
+ var retryData = await taskTable_.ReadTaskAsync(retryId,
+ lateCts_.Token)
+ .ConfigureAwait(false);
+ if (retryData.Status is TaskStatus.Creating or TaskStatus.Submitted)
+ {
+ logger_.LogWarning("Retried task {task} is in {status}; trying to finalize task creation",
+ retryId,
+ retryData.Status);
+ await submitter_.FinalizeTaskCreation(new List<TaskCreationRequest>
+ {
+ new(retryId,
+ retryData.PayloadId,
+ retryData.Options,
+ retryData.ExpectedOutputIds,
+ retryData.DataDependencies),
+ },
+ sessionData_,
+ taskData_.TaskId,
+ CancellationToken.None)
+ .ConfigureAwait(false);
+ }
+ }
+
+ return AcquisitionStatus.TaskIsRetriedAndRetryIsNotFound;
+ }
+
  return AcquisitionStatus.TaskIsRetried;
  case TaskStatus.Unspecified:
  default:

Common/src/Storage/TaskData.cs

@@ -204,4 +204,14 @@ public static implicit operator Application(TaskData taskData)
  taskData.Options.ApplicationNamespace,
  taskData.Options.ApplicationVersion,
  taskData.Options.ApplicationService);
+
+ /// <summary>
+ /// Compute the Id of the new task if this task should be retried
+ /// Should be deterministic as it can be called several time
+ /// </summary>
+ /// <returns>
+ /// Id of the retried task
+ /// </returns>
+ public string RetryId()
+ => InitialTaskId + $"###{RetryOfIds.Count + 1}";
 }

Common/src/Storage/TaskTableExtensions.cs

@@ -484,7 +484,7 @@ public static async Task<string> RetryTask(this ITaskTable taskTable,
  TaskData taskData,
  CancellationToken cancellationToken = default)
  {
- var newTaskId = taskData.InitialTaskId + $"###{taskData.RetryOfIds.Count + 1}";
+ var newTaskId = taskData.RetryId();
 
  var newTaskRetryOfIds = new List<string>(taskData.RetryOfIds)
  {

Common/tests/Helpers/TestTaskHandlerProvider.cs

@@ -22,7 +22,6 @@
 using System.Threading;
 
 using ArmoniK.Api.Common.Options;
-using ArmoniK.Core.Adapters.Memory;
 using ArmoniK.Core.Adapters.MongoDB;
 using ArmoniK.Core.Base;
 using ArmoniK.Core.Common.gRPC.Services;
@@ -50,19 +49,22 @@ namespace ArmoniK.Core.Common.Tests.Helpers;
 
 public class TestTaskHandlerProvider : IDisposable
 {
- private const string DatabaseName = "ArmoniK_TestDB";
- private static readonly ActivitySource ActivitySource = new("ArmoniK.Core.Common.Tests.TestTaskHandlerProvider");
- private readonly WebApplication app_;
- private readonly IMongoClient client_;
- private readonly LoggerFactory loggerFactory_;
- private readonly IObjectStorage objectStorage_;
- public readonly IPartitionTable PartitionTable;
- public readonly IResultTable ResultTable;
- private readonly IMongoRunner runner_;
- public readonly ISessionTable SessionTable;
- public readonly ISubmitter Submitter;
- public readonly TaskHandler TaskHandler;
- public readonly ITaskTable TaskTable;
+ private const string DatabaseName = "ArmoniK_TestDB";
+ private static readonly ActivitySource ActivitySource = new("ArmoniK.Core.Common.Tests.TestTaskHandlerProvider");
+ private readonly WebApplication app_;
+ private readonly IMongoClient client_;
+ public readonly ILogger Logger;
+ private readonly LoggerFactory loggerFactory_;
+ private readonly IObjectStorage objectStorage_;
+ public readonly IPartitionTable PartitionTable;
+ public readonly IPushQueueStorage PushQueueStorage;
+ public readonly IResultTable ResultTable;
+ private readonly IMongoRunner runner_;
+ public readonly ISessionTable SessionTable;
+ public readonly ISubmitter Submitter;
+ public readonly TaskHandler TaskHandler;
+ public readonly ITaskTable TaskTable;
+
 
  public IHostApplicationLifetime Lifetime;
 
@@ -139,6 +141,7 @@ public TestTaskHandlerProvider(IWorkerStreamHandler workerStreamHandler,
 
  loggerFactory_ = new LoggerFactory();
  loggerFactory_.AddProvider(new ConsoleForwardingLoggerProvider());
+ Logger = loggerFactory_.CreateLogger("root");
 
  var builder = WebApplication.CreateBuilder();
 
@@ -155,7 +158,7 @@ public TestTaskHandlerProvider(IWorkerStreamHandler workerStreamHandler,
  Injection.Options.Submitter.SettingSection)
  .AddOption<Injection.Options.Pollster>(builder.Configuration,
  Injection.Options.Pollster.SettingSection)
- .AddSingleton<IPushQueueStorage, PushQueueStorage>()
+ .AddSingleton<IPushQueueStorage, SimplePushQueueStorage>()
  .AddSingleton<MeterHolder>()
  .AddSingleton<AgentIdentifier>()
  .AddSingleton<ExceptionManager.Options>()
@@ -206,14 +209,15 @@ public TestTaskHandlerProvider(IWorkerStreamHandler workerStreamHandler,
 
  app_ = builder.Build();
 
- ResultTable = app_.Services.GetRequiredService<IResultTable>();
- TaskTable = app_.Services.GetRequiredService<ITaskTable>();
- PartitionTable = app_.Services.GetRequiredService<IPartitionTable>();
- SessionTable = app_.Services.GetRequiredService<ISessionTable>();
- Submitter = app_.Services.GetRequiredService<ISubmitter>();
- TaskHandler = app_.Services.GetRequiredService<TaskHandler>();
- Lifetime = app_.Services.GetRequiredService<IHostApplicationLifetime>();
- objectStorage_ = app_.Services.GetRequiredService<IObjectStorage>();
+ ResultTable = app_.Services.GetRequiredService<IResultTable>();
+ TaskTable = app_.Services.GetRequiredService<ITaskTable>();
+ PartitionTable = app_.Services.GetRequiredService<IPartitionTable>();
+ SessionTable = app_.Services.GetRequiredService<ISessionTable>();
+ Submitter = app_.Services.GetRequiredService<ISubmitter>();
+ TaskHandler = app_.Services.GetRequiredService<TaskHandler>();
+ Lifetime = app_.Services.GetRequiredService<IHostApplicationLifetime>();
+ objectStorage_ = app_.Services.GetRequiredService<IObjectStorage>();
+ PushQueueStorage = app_.Services.GetRequiredService<IPushQueueStorage>();
 
  ResultTable.Init(CancellationToken.None)
  .Wait();