Cosmos DB: support custom partition key value providers (dotnet#8694)

benjaminpetit · Nov 2, 2023 · e54b034 · e54b034
1 parent e2c8d77
commit e54b034
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs b/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs
@@ -404,6 +404,16 @@ public static class CosmosStorageFactory
     public static IGrainStorage Create(IServiceProvider services, string name)
     {
         var optionsMonitor = services.GetRequiredService<IOptionsMonitor<CosmosGrainStorageOptions>>();
-        return ActivatorUtilities.CreateInstance<CosmosGrainStorage>(services, name, optionsMonitor.Get(name));
+        var partitionKeyProvider = services.GetServiceByName<IPartitionKeyProvider>(name)
+            ?? services.GetRequiredService<IPartitionKeyProvider>();
+        var loggerFactory = services.GetRequiredService<ILoggerFactory>();
+        var clusterOptions = services.GetRequiredService<IOptions<ClusterOptions>>();
+        return new CosmosGrainStorage(
+            name,
+            optionsMonitor.Get(name),
+            loggerFactory,
+            services,
+            clusterOptions,
+            partitionKeyProvider);
     }
 }
diff --git a/src/Azure/Orleans.Persistence.Cosmos/HostingExtensions.cs b/src/Azure/Orleans.Persistence.Cosmos/HostingExtensions.cs
@@ -36,7 +36,7 @@ public static ISiloBuilder AddCosmosGrainStorage<TPartitionKeyProvider>(
         string name,
         Action<CosmosGrainStorageOptions> configureOptions) where TPartitionKeyProvider : class, IPartitionKeyProvider
     {
-        builder.Services.TryAddSingleton<IPartitionKeyProvider, TPartitionKeyProvider>();
+        builder.Services.AddSingletonNamedService<IPartitionKeyProvider, TPartitionKeyProvider>(name);
         builder.Services.AddCosmosGrainStorage(name, configureOptions);
         return builder;
     }
@@ -129,7 +129,7 @@ public static ISiloBuilder AddCosmosGrainStorage<TPartitionKeyProvider>(
         string name,
         Action<OptionsBuilder<CosmosGrainStorageOptions>>? configureOptions = null) where TPartitionKeyProvider : class, IPartitionKeyProvider
     {
-        builder.Services.TryAddSingleton<IPartitionKeyProvider, TPartitionKeyProvider>();
+        builder.Services.AddSingletonNamedService<IPartitionKeyProvider, TPartitionKeyProvider>(name);
         builder.Services.AddCosmosGrainStorage(name, configureOptions);
         return builder;
     }
@@ -162,7 +162,7 @@ public static ISiloBuilder AddCosmosGrainStorage(
     {
         if (customPartitionKeyProviderType != null)
         {
-            builder.Services.TryAddSingleton(typeof(IPartitionKeyProvider), customPartitionKeyProviderType);
+            builder.Services.AddSingletonNamedService<IPartitionKeyProvider>(name, customPartitionKeyProviderType);
         }
 
         builder.Services.AddCosmosGrainStorage(name, configureOptions);

diff --git a/src/Orleans.Core/Utils/KeyedService.cs b/src/Orleans.Core/Utils/KeyedService.cs
@@ -205,6 +205,15 @@ public static IServiceCollection AddSingletonNamedService<TService>(this IServic
             return collection.AddSingletonKeyedService<string, TService>(name, factory);
         }
 
+        /// <summary>
+        /// Register a singleton named service
+        /// </summary>
+        public static IServiceCollection AddSingletonNamedService<TService>(this IServiceCollection collection, string name, Type implementationType)
+            where TService : class
+        {
+            return collection.AddSingletonKeyedService<string, TService>(name, (sp, name) => (TService)ActivatorUtilities.CreateInstance(sp, implementationType));
+        }
+
         /// <summary>
         /// Register a singleton named service
         /// </summary>