From e54b03466763979c29063f842fcb91f6d274e5c2 Mon Sep 17 00:00:00 2001
From: Reuben Bond <203839+ReubenBond@users.noreply.github.com>
Date: Tue, 31 Oct 2023 08:42:10 -0700
Subject: [PATCH] Cosmos DB: support custom partition key value providers
 (#8694)

---
 .../Orleans.Persistence.Cosmos/CosmosGrainStorage.cs | 12 +++++++++++-
 .../Orleans.Persistence.Cosmos/HostingExtensions.cs  |  6 +++---
 src/Orleans.Core/Utils/KeyedService.cs               |  9 +++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs b/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs
index 49db39df6a..cb140e60df 100644
--- 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);
     }
 }
\ No newline at end of file
diff --git a/src/Azure/Orleans.Persistence.Cosmos/HostingExtensions.cs b/src/Azure/Orleans.Persistence.Cosmos/HostingExtensions.cs
index 7294276944..624bc721b3 100644
--- 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
index 38ba370b4d..5d5edc1eb9 100644
--- 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>