feat: add ttl for session, task and result data in db (#452)

aneoconsulting · Jul 28, 2023 · a3dd2b8 · a3dd2b8
2 parents 176162f + 44be37c
commit a3dd2b8
Show file tree

Hide file tree

Showing 14 changed files with 90 additions and 20 deletions.
diff --git a/Adaptors/MongoDB/src/Common/IMongoDataModelMapping.cs b/Adaptors/MongoDB/src/Common/IMongoDataModelMapping.cs
@@ -1,4 +1,4 @@
-// This file is part of the ArmoniK project
+// This file is part of the ArmoniK project
 // 
 // Copyright (C) ANEO, 2021-2023. All rights reserved.
 // 
@@ -26,5 +26,6 @@ public interface IMongoDataModelMapping<T>
   string CollectionName { get; }
 
   Task InitializeIndexesAsync(IClientSessionHandle sessionHandle,
-                              IMongoCollection<T>  collection);
+                              IMongoCollection<T>  collection,
+                              Options.MongoDB      options);
 }
diff --git a/Adaptors/MongoDB/src/Common/MongoCollectionProvider.cs b/Adaptors/MongoDB/src/Common/MongoCollectionProvider.cs
@@ -149,7 +149,8 @@ await sessionProvider.Init(cancellationToken)
       try
       {
         await model.InitializeIndexesAsync(session,
-                                           output)
+                                           output,
+                                           options)
                    .ConfigureAwait(false);
         break;
       }

diff --git a/Adaptors/MongoDB/src/Object/ObjectDataModelMapping.cs b/Adaptors/MongoDB/src/Object/ObjectDataModelMapping.cs
@@ -48,7 +48,8 @@ public string Id
 
   /// <inheritdoc />
   public async Task InitializeIndexesAsync(IClientSessionHandle                     sessionHandle,
-                                           IMongoCollection<ObjectDataModelMapping> collection)
+                                           IMongoCollection<ObjectDataModelMapping> collection,
+                                           Options.MongoDB                          options)
   {
     var keyIndex      = Builders<ObjectDataModelMapping>.IndexKeys.Hashed(model => model.Key);
     var chunkIdxIndex = Builders<ObjectDataModelMapping>.IndexKeys.Hashed(model => model.ChunkIdx);

diff --git a/Adaptors/MongoDB/src/Options/MongoDB.cs b/Adaptors/MongoDB/src/Options/MongoDB.cs
@@ -52,7 +52,7 @@ public class MongoDB
 
   public string DatabaseName { get; set; } = "ArmoniK";
 
-  public TimeSpan DataRetention { get; set; } = TimeSpan.FromDays(15);
+  public TimeSpan DataRetention { get; set; } = TimeSpan.MaxValue;
 
   public TableStorage TableStorage { get; set; } = new();
 

diff --git a/Adaptors/MongoDB/src/Table/DataModel/Auth/AuthDataModelMapping.cs b/Adaptors/MongoDB/src/Table/DataModel/Auth/AuthDataModelMapping.cs
@@ -62,7 +62,8 @@ public string CollectionName
 
   /// <inheritdoc />
   public async Task InitializeIndexesAsync(IClientSessionHandle       sessionHandle,
-                                           IMongoCollection<AuthData> collection)
+                                           IMongoCollection<AuthData> collection,
+                                           Options.MongoDB            options)
   {
     var indexModels = new[]
                       {

diff --git a/Adaptors/MongoDB/src/Table/DataModel/Auth/RoleDataModelMapping.cs b/Adaptors/MongoDB/src/Table/DataModel/Auth/RoleDataModelMapping.cs
@@ -59,7 +59,8 @@ public string CollectionName
 
   /// <inheritdoc />
   public async Task InitializeIndexesAsync(IClientSessionHandle       sessionHandle,
-                                           IMongoCollection<RoleData> collection)
+                                           IMongoCollection<RoleData> collection,
+                                           Options.MongoDB            options)
   {
     var indexModels = new[]
                       {

diff --git a/Adaptors/MongoDB/src/Table/DataModel/Auth/UserDataModelMapping.cs b/Adaptors/MongoDB/src/Table/DataModel/Auth/UserDataModelMapping.cs
@@ -60,7 +60,8 @@ public string CollectionName
 
   /// <inheritdoc />
   public async Task InitializeIndexesAsync(IClientSessionHandle       sessionHandle,
-                                           IMongoCollection<UserData> collection)
+                                           IMongoCollection<UserData> collection,
+                                           Options.MongoDB            options)
   {
     var indexModels = new[]
                       {

diff --git a/Adaptors/MongoDB/src/Table/DataModel/IndexHelper.cs b/Adaptors/MongoDB/src/Table/DataModel/IndexHelper.cs
@@ -80,16 +80,21 @@ public static CreateIndexModel<T> CreateTextIndex<T>(Expression<Func<T, object?>
   /// <typeparam name="T">Type stored in database</typeparam>
   /// <param name="expr">Expression to select the field for the index</param>
   /// <param name="unique">Unicity constraint, default to false</param>
+  /// <param name="expireAfter">Setup document should expire</param>
   /// <returns>
   ///   The ascending index model
   /// </returns>
   public static CreateIndexModel<T> CreateAscendingIndex<T>(Expression<Func<T, object?>> expr,
-                                                            bool                         unique = false)
+                                                            bool                         unique      = false,
+                                                            TimeSpan?                    expireAfter = null)
     => new(Builders<T>.IndexKeys.Ascending(new ExpressionFieldDefinition<T>(expr)),
            new CreateIndexOptions
            {
              Name   = $"{expr.GetMember().Name}_1",
              Unique = unique,
+             ExpireAfter = expireAfter == TimeSpan.MaxValue
+                             ? null
+                             : expireAfter,
            });
 
   /// <summary>

diff --git a/Adaptors/MongoDB/src/Table/DataModel/PartitionDataModelMapping.cs b/Adaptors/MongoDB/src/Table/DataModel/PartitionDataModelMapping.cs
@@ -78,7 +78,8 @@ public string CollectionName
 
   /// <inheritdoc />
   public async Task InitializeIndexesAsync(IClientSessionHandle            sessionHandle,
-                                           IMongoCollection<PartitionData> collection)
+                                           IMongoCollection<PartitionData> collection,
+                                           Options.MongoDB                 options)
   {
     var indexModels = new[]
                       {

diff --git a/Adaptors/MongoDB/src/Table/DataModel/ResultDataModelMapping.cs b/Adaptors/MongoDB/src/Table/DataModel/ResultDataModelMapping.cs
@@ -70,13 +70,15 @@ public string CollectionName
 
   /// <inheritdoc />
   public async Task InitializeIndexesAsync(IClientSessionHandle     sessionHandle,
-                                           IMongoCollection<Result> collection)
+                                           IMongoCollection<Result> collection,
+                                           Options.MongoDB          options)
   {
     var indexModels = new[]
                       {
                         IndexHelper.CreateHashedIndex<Result>(model => model.SessionId),
                         IndexHelper.CreateHashedIndex<Result>(model => model.OwnerTaskId),
-                        IndexHelper.CreateAscendingIndex<Result>(model => model.CreationDate),
+                        IndexHelper.CreateAscendingIndex<Result>(model => model.CreationDate,
+                                                                 expireAfter: options.DataRetention),
                       };
 
     await collection.Indexes.CreateManyAsync(sessionHandle,

diff --git a/Adaptors/MongoDB/src/Table/DataModel/SessionDataModelMapping.cs b/Adaptors/MongoDB/src/Table/DataModel/SessionDataModelMapping.cs
@@ -100,11 +100,13 @@ public string CollectionName
 
   /// <inheritdoc />
   public async Task InitializeIndexesAsync(IClientSessionHandle          sessionHandle,
-                                           IMongoCollection<SessionData> collection)
+                                           IMongoCollection<SessionData> collection,
+                                           Options.MongoDB               options)
   {
     var indexModels = new[]
                       {
-                        IndexHelper.CreateAscendingIndex<SessionData>(model => model.CreationDate),
+                        IndexHelper.CreateAscendingIndex<SessionData>(model => model.CreationDate,
+                                                                      expireAfter: options.DataRetention),
                         IndexHelper.CreateAscendingIndex<SessionData>(model => model.CancellationDate),
                         IndexHelper.CreateHashedIndex<SessionData>(model => model.Status),
                         IndexHelper.CreateHashedIndex<SessionData>(model => model.Options.PartitionId),

diff --git a/Adaptors/MongoDB/src/Table/DataModel/TaskDataModelMapping.cs b/Adaptors/MongoDB/src/Table/DataModel/TaskDataModelMapping.cs
@@ -173,7 +173,8 @@ public string CollectionName
 
   /// <inheritdoc />
   public async Task InitializeIndexesAsync(IClientSessionHandle       sessionHandle,
-                                           IMongoCollection<TaskData> collection)
+                                           IMongoCollection<TaskData> collection,
+                                           Options.MongoDB            options)
   {
     var indexModels = new[]
                       {
@@ -182,7 +183,8 @@ public async Task InitializeIndexesAsync(IClientSessionHandle       sessionHandl
                         IndexHelper.CreateHashedIndex<TaskData>(model => model.SessionId),
                         IndexHelper.CreateHashedIndex<TaskData>(model => model.OwnerPodId),
                         IndexHelper.CreateHashedIndex<TaskData>(model => model.InitialTaskId),
-                        IndexHelper.CreateAscendingIndex<TaskData>(model => model.CreationDate),
+                        IndexHelper.CreateAscendingIndex<TaskData>(model => model.CreationDate,
+                                                                   expireAfter: options.DataRetention),
                         IndexHelper.CreateAscendingIndex<TaskData>(model => model.SubmittedDate),
                         IndexHelper.CreateAscendingIndex<TaskData>(model => model.StartDate),
                         IndexHelper.CreateAscendingIndex<TaskData>(model => model.EndDate),

diff --git a/Adaptors/MongoDB/tests/IndexTest.cs b/Adaptors/MongoDB/tests/IndexTest.cs
@@ -130,8 +130,9 @@ public void IndexCreation2ShouldSucceed()
                       {
                         IndexHelper.CreateHashedIndex<TaskData>(model => model.TaskId),
                         IndexHelper.CreateAscendingIndex<TaskData>(model => model.PodTtl),
-                        IndexHelper.CreateAscendingIndex<TaskData>(model => model.Options.MaxDuration),
-                        IndexHelper.CreateDescendingIndex<TaskData>(model => model.CreationDate),
+                        IndexHelper.CreateDescendingIndex<TaskData>(model => model.Options.MaxDuration),
+                        IndexHelper.CreateAscendingIndex<TaskData>(model => model.CreationDate,
+                                                                   expireAfter: TimeSpan.FromDays(1)),
                         IndexHelper.CreateTextIndex<TaskData>(model => model.OwnerPodId),
                       };
 
@@ -148,6 +149,56 @@ public void IndexCreation2ShouldSucceed()
                               .Count);
   }
 
+  [Test]
+  public void IndexCreationWithMaxExpireShouldSucceed()
+  {
+    var db         = provider_!.GetRequiredService<IMongoDatabase>();
+    var collection = db.GetCollection<TaskData>("Test");
+
+    var indexModels = new[]
+                      {
+                        IndexHelper.CreateAscendingIndex<TaskData>(model => model.CreationDate,
+                                                                   expireAfter: TimeSpan.MaxValue),
+                      };
+
+    collection.Indexes.CreateMany(indexModels);
+    foreach (var index in collection.Indexes.List()
+                                    .ToList())
+    {
+      Console.WriteLine(index);
+    }
+
+    Assert.AreEqual(indexModels.Length + 1,
+                    collection.Indexes.List()
+                              .ToList()
+                              .Count);
+  }
+
+  [Test]
+  public void IndexCreationWithNullExpireShouldSucceed()
+  {
+    var db         = provider_!.GetRequiredService<IMongoDatabase>();
+    var collection = db.GetCollection<TaskData>("Test");
+
+    var indexModels = new[]
+                      {
+                        IndexHelper.CreateAscendingIndex<TaskData>(model => model.CreationDate,
+                                                                   expireAfter: null),
+                      };
+
+    collection.Indexes.CreateMany(indexModels);
+    foreach (var index in collection.Indexes.List()
+                                    .ToList())
+    {
+      Console.WriteLine(index);
+    }
+
+    Assert.AreEqual(indexModels.Length + 1,
+                    collection.Indexes.List()
+                              .ToList()
+                              .Count);
+  }
+
   [Test]
   public void CombinedIndexCreationShouldSucceed()
   {

diff --git a/ProfilingTools/src/OpenTelemetryExporter/OpenTelemetryDataModelMapping.cs b/ProfilingTools/src/OpenTelemetryExporter/OpenTelemetryDataModelMapping.cs
@@ -1,4 +1,4 @@
-// This file is part of the ArmoniK project
+// This file is part of the ArmoniK project
 // 
 // Copyright (C) ANEO, 2021-2023. All rights reserved.
 // 
@@ -82,7 +82,8 @@ public string CollectionName
 
   /// <inheritdoc />
   public async Task InitializeIndexesAsync(IClientSessionHandle                sessionHandle,
-                                           IMongoCollection<OpenTelemetryData> collection)
+                                           IMongoCollection<OpenTelemetryData> collection,
+                                           Adapters.MongoDB.Options.MongoDB    options)
   {
     var sourceNameIndex  = Builders<OpenTelemetryData>.IndexKeys.Hashed(model => model.SourceName);
     var displayNameIndex = Builders<OpenTelemetryData>.IndexKeys.Hashed(model => model.DisplayName);