Shard collections at DB init when sharding is required

Adaptors/MongoDB/src/Common/IMongoDataModelMapping.cs

@@ -26,6 +26,9 @@ public interface IMongoDataModelMapping<T>
  string CollectionName { get; }
 
  Task InitializeIndexesAsync(IClientSessionHandle sessionHandle,
- IMongoCollection<T> collection,
- Options.MongoDB options);
+ IMongoCollection<T> collection,
+ Options.MongoDB options);
+
+ Task ShardCollectionAsync(IClientSessionHandle sessionHandle,
+ Options.MongoDB options);
 }

Adaptors/MongoDB/src/Common/MongoCollectionProvider.cs

@@ -39,11 +39,11 @@ public class MongoCollectionProvider<TData, TModelMapping> : IInitializable, IAs
  private bool isInitialized_;
  private IMongoCollection<TData>? mongoCollection_;
 
- public MongoCollectionProvider(Options.MongoDB  options,
- SessionProvider  sessionProvider,
- IMongoDatabase  mongoDatabase,
+ public MongoCollectionProvider(Options.MongoDB options,
+ SessionProvider sessionProvider,
+ IMongoDatabase mongoDatabase,
  ILogger<IMongoCollection<TData>> logger,
- CancellationToken  cancellationToken = default)
+ CancellationToken cancellationToken = default)
  {
  if (options.DataRetention == TimeSpan.Zero)
  {
@@ -64,17 +64,17 @@ public MongoCollectionProvider(Options.MongoDB options,
  /// <inheritdoc />
  public Task<HealthCheckResult> Check(HealthCheckTag tag)
  => tag switch
-  {
-  HealthCheckTag.Startup or HealthCheckTag.Readiness => Task.FromResult(isInitialized_
-  ? HealthCheckResult.Healthy()
-  : HealthCheckResult.Unhealthy("MongoCollection not initialized yet.")),
-  HealthCheckTag.Liveness => Task.FromResult(isInitialized_ && mongoCollection_ is null
-  ? HealthCheckResult.Healthy()
-  : HealthCheckResult.Unhealthy("MongoCollection not initialized yet.")),
-  _ => throw new ArgumentOutOfRangeException(nameof(tag),
-  tag,
-  null),
-  };
+ {
+ HealthCheckTag.Startup or HealthCheckTag.Readiness => Task.FromResult(isInitialized_
+ ? HealthCheckResult.Healthy()
+ : HealthCheckResult.Unhealthy("MongoCollection not initialized yet.")),
+ HealthCheckTag.Liveness => Task.FromResult(isInitialized_ && mongoCollection_ is null
+ ? HealthCheckResult.Healthy()
+ : HealthCheckResult.Unhealthy("MongoCollection not initialized yet.")),
+ _ => throw new ArgumentOutOfRangeException(nameof(tag),
+ tag,
+ null),
+ };
 
  /// <inheritdoc />
  public async Task Init(CancellationToken cancellationToken)
@@ -87,13 +87,13 @@ public async Task Init(CancellationToken cancellationToken)
  isInitialized_ = true;
  }
 
- private static async Task<IMongoCollection<TData>> InitializeAsync(Options.MongoDB  options,
- SessionProvider  sessionProvider,
- IMongoDatabase  mongoDatabase,
+ private static async Task<IMongoCollection<TData>> InitializeAsync(Options.MongoDB options,
+ SessionProvider sessionProvider,
+ IMongoDatabase mongoDatabase,
  ILogger<IMongoCollection<TData>> logger,
- CancellationToken  cancellationToken = default)
+ CancellationToken cancellationToken = default)
  {
- var  model  = new TModelMapping();
+ var model = new TModelMapping();
  Exception? lastException = null;
 
  for (var collectionRetry = 1; collectionRetry < options.MaxRetries; collectionRetry++)
@@ -167,9 +167,41 @@ await Task.Delay(1000 * indexRetry,
  }
  }
 
+ if (options.Sharding){
+
+ for (var indexRetry = 1; indexRetry < options.MaxRetries; indexRetry++)
+ {
+ lastException = null;
+ try
+ {
+ await model.ShardCollectionAsync(session,
+ options)
+ .ConfigureAwait(false);
+ break;
+ }
+ catch (MongoCommandException ex) when (ex.CodeName == "IndexOptionsConflict")
+ {
+ logger.LogWarning(ex,
+ "Index options conflict for {CollectionName} collection",
+ model.CollectionName);
+ break;
+ }
+ catch (Exception ex)
+ {
+ lastException = ex;
+ logger.LogDebug(ex,
+ "Retrying to shard {CollectionName} collection",
+ model.CollectionName);
+ await Task.Delay(1000 * indexRetry,
+ cancellationToken)
+ .ConfigureAwait(false);
+ }
+ }
+ }
+
  if (lastException is not null)
  {
- throw new TimeoutException($"Init Indexes for {model.CollectionName}: Max retries reached",
+ throw new TimeoutException($"Init Index or shard for {model.CollectionName}: Max retries reached",
  lastException);
  }
 

Adaptors/MongoDB/src/Object/ObjectDataModelMapping.cs

@@ -78,4 +78,6 @@ await collection.Indexes.CreateManyAsync(sessionHandle,
  indexModels)
  .ConfigureAwait(false);
  }
+
+ public Task ShardCollectionAsync(IClientSessionHandle sessionHandle, Options.MongoDB options) => Task.CompletedTask;
 }

Adaptors/MongoDB/src/Options/MongoDB.cs

@@ -60,6 +60,11 @@ public class MongoDB
 
  public QueueStorage QueueStorage { get; set; } = new();
 
- public int MaxConnectionPoolSize { get; set; } = 500;
+ public int MaxConnectionPoolSize { get; set; } = 500;
+
  public TimeSpan ServerSelectionTimeout { get; set; } = TimeSpan.FromMinutes(2);
+
+ public bool Sharding { get; set; }
+
+ public string AuthSource { get; set; } = "";
 }

Adaptors/MongoDB/src/ServiceCollectionExt.cs

@@ -184,6 +184,10 @@ public static IServiceCollection AddMongoClient(this IServiceCollection services
  mongoOptions.DatabaseName);
  }
 
+ if (!string.IsNullOrEmpty(mongoOptions.AuthSource)) {
+ connectionString = $"{connectionString}?authSource={mongoOptions.AuthSource}";
+ }
+
  var settings = MongoClientSettings.FromUrl(new MongoUrl(connectionString));
  settings.AllowInsecureTls = mongoOptions.AllowInsecureTls;
  settings.UseTls = mongoOptions.Tls;

Adaptors/MongoDB/src/Table/DataModel/Auth/AuthDataModelMapping.cs

@@ -61,9 +61,9 @@ public string CollectionName
  => nameof(AuthData);
 
  /// <inheritdoc />
- public async Task InitializeIndexesAsync(IClientSessionHandle  sessionHandle,
+ public async Task InitializeIndexesAsync(IClientSessionHandle sessionHandle,
  IMongoCollection<AuthData> collection,
- Options.MongoDB  options)
+ Options.MongoDB options)
  {
  var indexModels = new[]
  {
@@ -77,4 +77,6 @@ await collection.Indexes.CreateManyAsync(sessionHandle,
  indexModels)
  .ConfigureAwait(false);
  }
+
+ public Task ShardCollectionAsync(IClientSessionHandle sessionHandle, Options.MongoDB options) => Task.CompletedTask;
 }

Adaptors/MongoDB/src/Table/DataModel/Auth/RoleDataModelMapping.cs

@@ -58,9 +58,9 @@ public string CollectionName
  => nameof(RoleData);
 
  /// <inheritdoc />
- public async Task InitializeIndexesAsync(IClientSessionHandle  sessionHandle,
+ public async Task InitializeIndexesAsync(IClientSessionHandle sessionHandle,
  IMongoCollection<RoleData> collection,
- Options.MongoDB  options)
+ Options.MongoDB options)
  {
  var indexModels = new[]
  {
@@ -73,4 +73,5 @@ await collection.Indexes.CreateManyAsync(sessionHandle,
  indexModels)
  .ConfigureAwait(false);
  }
+ public Task ShardCollectionAsync(IClientSessionHandle sessionHandle, Options.MongoDB options) => Task.CompletedTask;
 }

Adaptors/MongoDB/src/Table/DataModel/Auth/UserDataModelMapping.cs

@@ -59,9 +59,9 @@ public string CollectionName
  => nameof(UserData);
 
  /// <inheritdoc />
- public async Task InitializeIndexesAsync(IClientSessionHandle  sessionHandle,
+ public async Task InitializeIndexesAsync(IClientSessionHandle sessionHandle,
  IMongoCollection<UserData> collection,
- Options.MongoDB  options)
+ Options.MongoDB options)
  {
  var indexModels = new[]
  {
@@ -73,4 +73,6 @@ await collection.Indexes.CreateManyAsync(sessionHandle,
  indexModels)
  .ConfigureAwait(false);
  }
+
+ public Task ShardCollectionAsync(IClientSessionHandle sessionHandle, Options.MongoDB options) => Task.CompletedTask;
 }

Adaptors/MongoDB/src/Table/DataModel/PartitionDataModelMapping.cs

@@ -90,4 +90,6 @@ await collection.Indexes.CreateManyAsync(sessionHandle,
  indexModels)
  .ConfigureAwait(false);
  }
+
+ public Task ShardCollectionAsync(IClientSessionHandle sessionHandle, Options.MongoDB options) => Task.CompletedTask;
 }

Adaptors/MongoDB/src/Table/DataModel/ResultDataModelMapping.cs

@@ -76,9 +76,9 @@ public string CollectionName
  => nameof(Result);
 
  /// <inheritdoc />
- public async Task InitializeIndexesAsync(IClientSessionHandle  sessionHandle,
+ public async Task InitializeIndexesAsync(IClientSessionHandle sessionHandle,
  IMongoCollection<Result> collection,
- Options.MongoDB  options)
+ Options.MongoDB options)
  {
  var indexModels = new[]
  {
@@ -92,4 +92,11 @@ await collection.Indexes.CreateManyAsync(sessionHandle,
  indexModels)
  .ConfigureAwait(false);
  }
+
+ public async Task ShardCollectionAsync(IClientSessionHandle sessionHandle,
+ Options.MongoDB options)
+ {
+ await sessionHandle.shardCollection(options, CollectionName)
+ .ConfigureAwait(false);
+ }
 }

Adaptors/MongoDB/src/Table/DataModel/SessionDataModelMapping.cs

@@ -120,9 +120,9 @@ public string CollectionName
  => nameof(SessionData);
 
  /// <inheritdoc />
- public async Task InitializeIndexesAsync(IClientSessionHandle  sessionHandle,
+ public async Task InitializeIndexesAsync(IClientSessionHandle sessionHandle,
  IMongoCollection<SessionData> collection,
- Options.MongoDB  options)
+ Options.MongoDB options)
  {
  var indexModels = new[]
  {
@@ -137,4 +137,11 @@ await collection.Indexes.CreateManyAsync(sessionHandle,
  indexModels)
  .ConfigureAwait(false);
  }
+
+ public async Task ShardCollectionAsync(IClientSessionHandle sessionHandle,
+ Options.MongoDB options)
+ {
+ await sessionHandle.shardCollection(options, CollectionName)
+ .ConfigureAwait(false);
+ }
 }

Adaptors/MongoDB/src/Table/DataModel/ShardingExt.cs

@@ -0,0 +1,46 @@
+// 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.Collections.Generic;
+using System.Threading.Tasks;
+
+using MongoDB.Driver;
+using MongoDB.Bson;
+
+namespace ArmoniK.Core.Adapters.MongoDB.Table.DataModel;
+
+public static class ShardingExt
+{
+ public static async Task shardCollection(this IClientSessionHandle sessionHandle, Options.MongoDB options, string collectionName)
+ {
+ var adminDB = sessionHandle.Client.GetDatabase("admin");
+ var shardingCommandDict = new Dictionary<string, object>
+ {
+ { "shardCollection", $"{options.DatabaseName}.{collectionName}"},
+ {
+ "key",
+ new Dictionary<string, object>{
+ { "_id", "hashed"}
+ }
+ }
+ };
+
+ var shardingCommand = new BsonDocumentCommand<BsonDocument>(new BsonDocument(shardingCommandDict));
+ await adminDB.RunCommandAsync(shardingCommand)
+ .ConfigureAwait(false);
+ }
+}

Adaptors/MongoDB/src/Table/DataModel/TaskDataModelMapping.cs

@@ -25,6 +25,7 @@
 
 using MongoDB.Bson.Serialization;
 using MongoDB.Driver;
+using MongoDB.Bson;
 
 namespace ArmoniK.Core.Adapters.MongoDB.Table.DataModel;
 
@@ -184,9 +185,9 @@ public string CollectionName
 
 
  /// <inheritdoc />
- public async Task InitializeIndexesAsync(IClientSessionHandle  sessionHandle,
+ public async Task InitializeIndexesAsync(IClientSessionHandle sessionHandle,
  IMongoCollection<TaskData> collection,
- Options.MongoDB  options)
+ Options.MongoDB options)
  {
  var indexModels = new[]
  {
@@ -210,4 +211,11 @@ await collection.Indexes.CreateManyAsync(sessionHandle,
  indexModels)
  .ConfigureAwait(false);
  }
+
+ public async Task ShardCollectionAsync(IClientSessionHandle sessionHandle,
+ Options.MongoDB options) 
+ {
+ await sessionHandle.shardCollection(options, CollectionName)
+ .ConfigureAwait(false);
+ }
 }

ProfilingTools/src/OpenTelemetryExporter/OpenTelemetryDataModelMapping.cs

@@ -110,4 +110,5 @@ await collection.Indexes.CreateManyAsync(sessionHandle,
  indexModels)
  .ConfigureAwait(false);
  }
+ public Task ShardCollectionAsync(IClientSessionHandle sessionHandle, Adapters.MongoDB.Options.MongoDB options) => Task.CompletedTask;
 }