From 5c3c5e4d8abe1e366aa996215f1b645a36561d50 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Wed, 1 Nov 2023 14:05:51 -0700 Subject: [PATCH] Support customization of both document id and partition key for Cosmos DB grain persistence --- .../CosmosGrainStorage.cs | 32 +++----- .../DefaultDocumentIdProvider.cs | 40 ++++++++++ .../HostingExtensions.cs | 78 ++++++++++--------- .../IDocumentIdProvider.cs | 15 ++++ .../IPartitionKeyProvider.cs | 20 ----- .../PersistenceProviderTests_Cosmos.cs | 4 +- 6 files changed, 107 insertions(+), 82 deletions(-) create mode 100644 src/Azure/Orleans.Persistence.Cosmos/DefaultDocumentIdProvider.cs create mode 100644 src/Azure/Orleans.Persistence.Cosmos/IDocumentIdProvider.cs delete mode 100644 src/Azure/Orleans.Persistence.Cosmos/IPartitionKeyProvider.cs diff --git a/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs b/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs index cb140e60df..afab664aa7 100644 --- a/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs +++ b/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs @@ -3,14 +3,12 @@ using System.Diagnostics; using Microsoft.Extensions.DependencyInjection; using Orleans.Storage; -using static Orleans.Persistence.Cosmos.CosmosIdSanitizer; namespace Orleans.Persistence.Cosmos; internal class CosmosGrainStorage : IGrainStorage, ILifecycleParticipant { private const string ANY_ETAG = "*"; - private const string KEY_STRING_SEPARATOR = "__"; private const string GRAINTYPE_PARTITION_KEY_PATH = "/GrainType"; private readonly ILogger _logger; private readonly CosmosGrainStorageOptions _options; @@ -18,8 +16,8 @@ internal class CosmosGrainStorage : IGrainStorage, ILifecycleParticipant clusterOptions, - IPartitionKeyProvider partitionKeyProvider - ) + IDocumentIdProvider documentIdProvider) { _logger = loggerFactory.CreateLogger(); _options = options; _name = name; _serviceProvider = serviceProvider; _serviceId = clusterOptions.Value.ServiceId; - _partitionKeyProvider = partitionKeyProvider; _executor = options.OperationExecutor; _partitionKeyPath = _options.PartitionKeyPath; + _documentIdProvider = documentIdProvider; } public async Task ReadStateAsync(string grainType, GrainId grainId, IGrainState grainState) { - var id = GetKeyString(grainId); - var partitionKey = await BuildPartitionKey(grainType, grainId); + var (id, partitionKey) = await _documentIdProvider.GetDocumentIdentifiers(grainType, grainId); if (_logger.IsEnabled(LogLevel.Trace)) { @@ -105,9 +101,7 @@ public async Task ReadStateAsync(string grainType, GrainId grainId, IGrainSta public async Task WriteStateAsync(string grainType, GrainId grainId, IGrainState grainState) { - var id = GetKeyString(grainId); - - var partitionKey = await BuildPartitionKey(grainType, grainId); + var (id, partitionKey) = await _documentIdProvider.GetDocumentIdentifiers(grainType, grainId); if (_logger.IsEnabled(LogLevel.Trace)) { @@ -188,8 +182,7 @@ public async Task WriteStateAsync(string grainType, GrainId grainId, IGrainSt public async Task ClearStateAsync(string grainType, GrainId grainId, IGrainState grainState) { - var id = GetKeyString(grainId); - var partitionKey = await BuildPartitionKey(grainType, grainId); + var (id, partitionKey) = await _documentIdProvider.GetDocumentIdentifiers(grainType, grainId); if (_logger.IsEnabled(LogLevel.Trace)) { _logger.LogTrace( @@ -262,11 +255,6 @@ public void Participate(ISiloLifecycle lifecycle) lifecycle.Subscribe(OptionFormattingUtilities.Name(_name), _options.InitStage, Init); } - private string GetKeyString(GrainId grainId) => $"{Sanitize(_serviceId)}{KEY_STRING_SEPARATOR}{Sanitize(grainId.Type.ToString()!)}{SeparatorChar}{Sanitize(grainId.Key.ToString()!)}"; - - private ValueTask BuildPartitionKey(string grainType, GrainId grainId) => - _partitionKeyProvider.GetPartitionKey(grainType, grainId); - private async Task Init(CancellationToken ct) { var stopWatch = Stopwatch.StartNew(); @@ -368,7 +356,7 @@ private async Task TryCreateResources() var container = containerResponse.Resource; _partitionKeyPath = container.PartitionKeyPath; if (_partitionKeyPath == GRAINTYPE_PARTITION_KEY_PATH && - _partitionKeyProvider is not DefaultPartitionKeyProvider) + _documentIdProvider is not DefaultDocumentIdProvider) throw new OrleansConfigurationException("Custom partition key provider is not compatible with partition key path set to /GrainType"); } @@ -404,8 +392,8 @@ public static class CosmosStorageFactory public static IGrainStorage Create(IServiceProvider services, string name) { var optionsMonitor = services.GetRequiredService>(); - var partitionKeyProvider = services.GetServiceByName(name) - ?? services.GetRequiredService(); + var documentIdProvider = services.GetServiceByName(name) + ?? services.GetRequiredService(); var loggerFactory = services.GetRequiredService(); var clusterOptions = services.GetRequiredService>(); return new CosmosGrainStorage( @@ -414,6 +402,6 @@ public static IGrainStorage Create(IServiceProvider services, string name) loggerFactory, services, clusterOptions, - partitionKeyProvider); + documentIdProvider); } } \ No newline at end of file diff --git a/src/Azure/Orleans.Persistence.Cosmos/DefaultDocumentIdProvider.cs b/src/Azure/Orleans.Persistence.Cosmos/DefaultDocumentIdProvider.cs new file mode 100644 index 0000000000..a36a7b16d6 --- /dev/null +++ b/src/Azure/Orleans.Persistence.Cosmos/DefaultDocumentIdProvider.cs @@ -0,0 +1,40 @@ +using static Orleans.Persistence.Cosmos.CosmosIdSanitizer; + +namespace Orleans.Persistence.Cosmos; + +/// +/// The default implementation of . +/// +public sealed class DefaultDocumentIdProvider : IDocumentIdProvider +{ + private const string KEY_STRING_SEPARATOR = "__"; + private readonly ClusterOptions _options; + + /// + /// Initializes a new instance of the class. + /// + /// The cluster options. + public DefaultDocumentIdProvider(IOptions options) + { + _options = options.Value; + } + + /// + public ValueTask<(string DocumentId, string PartitionKey)> GetDocumentIdentifiers(string stateName, GrainId grainId) => new((GetId(stateName, grainId), GetPartitionKey(stateName, grainId))); + + /// + /// Gets the id for the specified grain state document. + /// + /// The state name. + /// The grain id. + /// The document id. + public string GetId(string stateName, GrainId grainId) => $"{Sanitize(_options.ServiceId)}{KEY_STRING_SEPARATOR}{Sanitize(grainId.Type.ToString()!)}{SeparatorChar}{Sanitize(grainId.Key.ToString()!)}"; + + /// + /// Gets the Cosmos DB partition key for the specified grain state document. + /// + /// The state name. + /// The grain id. + /// The document partition key. + public string GetPartitionKey(string stateName, GrainId grainId) => Sanitize(stateName); +} \ 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 624bc721b3..034108944d 100644 --- a/src/Azure/Orleans.Persistence.Cosmos/HostingExtensions.cs +++ b/src/Azure/Orleans.Persistence.Cosmos/HostingExtensions.cs @@ -3,6 +3,7 @@ using Orleans.Storage; using Orleans.Providers; using Orleans.Persistence.Cosmos; +using Orleans.Configuration.Internal; namespace Orleans.Hosting; @@ -12,65 +13,65 @@ namespace Orleans.Hosting; public static class HostingExtensions { /// - /// Configure silo to use Azure Cosmos DB storage as the default grain storage using a custom Partition Key Provider. + /// Configure silo to use Azure Cosmos DB storage as the default grain storage using a custom document id provider. /// - /// The custom partition key provider type. + /// The document id provider. /// The silo builder. /// The delegate used to configure the provider. - public static ISiloBuilder AddCosmosGrainStorageAsDefault( + public static ISiloBuilder AddCosmosGrainStorageAsDefault( this ISiloBuilder builder, - Action configureOptions) where TPartitionKeyProvider : class, IPartitionKeyProvider + Action configureOptions) where TDocumentIdProvider : class, IDocumentIdProvider { - return builder.AddCosmosGrainStorage(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, configureOptions); + return builder.AddCosmosGrainStorage(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, configureOptions); } /// - /// Configure silo to use Azure Cosmos DB storage for grain storage using a custom Partition Key Provider. + /// Configure silo to use Azure Cosmos DB storage for grain storage using a custom document id provider. /// - /// The custom partition key provider type. + /// The document id provider. /// The silo builder. /// The storage provider name. /// The delegate used to configure the provider. - public static ISiloBuilder AddCosmosGrainStorage( + public static ISiloBuilder AddCosmosGrainStorage( this ISiloBuilder builder, string name, - Action configureOptions) where TPartitionKeyProvider : class, IPartitionKeyProvider + Action configureOptions) where TDocumentIdProvider : class, IDocumentIdProvider { - builder.Services.AddSingletonNamedService(name); + builder.Services.AddSingletonNamedService(name); builder.Services.AddCosmosGrainStorage(name, configureOptions); return builder; } /// - /// Configure silo to use Azure Cosmos DB storage as the default grain storage using a custom Partition Key Provider. + /// Configure silo to use Azure Cosmos DB storage as the default grain storage using a custom document id provider. /// /// The silo builder. /// The delegate used to configure the provider. - /// The custom partition key provider type. + /// The document id provider. public static ISiloBuilder AddCosmosGrainStorageAsDefault( this ISiloBuilder builder, Action configureOptions, - Type customPartitionKeyProviderType) + Type customDocumentIdProviderType) { - return builder.AddCosmosGrainStorage(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, configureOptions, customPartitionKeyProviderType); + return builder.AddCosmosGrainStorage(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, configureOptions, customDocumentIdProviderType); } /// - /// Configure silo to use Azure Cosmos DB storage for grain storage using a custom Partition Key Provider. + /// Configure silo to use Azure Cosmos DB storage for grain storage using a custom document id provider. /// /// The silo builder. /// The storage provider name. /// The delegate used to configure the provider. - /// The custom partition key provider type. + /// The document id provider. public static ISiloBuilder AddCosmosGrainStorage( this ISiloBuilder builder, string name, Action configureOptions, - Type customPartitionKeyProviderType) + Type customDocumentIdProviderType) { - if (customPartitionKeyProviderType != null) + if (customDocumentIdProviderType != null) { - builder.Services.TryAddSingleton(typeof(IPartitionKeyProvider), customPartitionKeyProviderType); + builder.Services.TryAddSingleton(typeof(IDocumentIdProvider), customDocumentIdProviderType); } builder.Services.AddCosmosGrainStorage(name, configureOptions); @@ -105,51 +106,51 @@ public static ISiloBuilder AddCosmosGrainStorage( } /// - /// Configure silo to use Azure Cosmos DB storage as the default grain storage using a custom Partition Key Provider. + /// Configure silo to use Azure Cosmos DB storage as the default grain storage using a custom document id provider. /// - /// The custom partition key provider type. + /// The document id provider. /// The silo builder. /// The delegate used to configure the provider. - public static ISiloBuilder AddCosmosGrainStorageAsDefault( + public static ISiloBuilder AddCosmosGrainStorageAsDefault( this ISiloBuilder builder, - Action>? configureOptions = null) where TPartitionKeyProvider : class, IPartitionKeyProvider + Action>? configureOptions = null) where TDocumentIdProvider : class, IDocumentIdProvider { - return builder.AddCosmosGrainStorage(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, configureOptions); + return builder.AddCosmosGrainStorage(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, configureOptions); } /// - /// Configure silo to use Azure Cosmos DB storage for grain storage using a custom Partition Key Provider. + /// Configure silo to use Azure Cosmos DB storage for grain storage using a custom document id provider. /// - /// The custom partition key provider type. + /// The document id provider. /// The silo builder. /// The storage provider name. /// The delegate used to configure the provider. - public static ISiloBuilder AddCosmosGrainStorage( + public static ISiloBuilder AddCosmosGrainStorage( this ISiloBuilder builder, string name, - Action>? configureOptions = null) where TPartitionKeyProvider : class, IPartitionKeyProvider + Action>? configureOptions = null) where TDocumentIdProvider : class, IDocumentIdProvider { - builder.Services.AddSingletonNamedService(name); + builder.Services.AddSingletonNamedService(name); builder.Services.AddCosmosGrainStorage(name, configureOptions); return builder; } /// - /// Configure silo to use Azure Cosmos DB storage as the default grain storage using a custom Partition Key Provider. + /// Configure silo to use Azure Cosmos DB storage as the default grain storage using a custom document id provider. /// /// The silo builder. - /// The custom partition key provider type. + /// The document id provider. /// The delegate used to configure the provider. public static ISiloBuilder AddCosmosGrainStorageAsDefault( this ISiloBuilder builder, - Type customPartitionKeyProviderType, + Type customDocumentIdProviderType, Action>? configureOptions = null) { - return builder.AddCosmosGrainStorage(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, customPartitionKeyProviderType, configureOptions); + return builder.AddCosmosGrainStorage(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, customDocumentIdProviderType, configureOptions); } /// - /// Configure silo to use Azure Cosmos DB storage for grain storage using a custom Partition Key Provider. + /// Configure silo to use Azure Cosmos DB storage for grain storage using a custom document id provider. /// /// The silo builder. /// The storage provider name. @@ -157,12 +158,12 @@ public static ISiloBuilder AddCosmosGrainStorageAsDefault( public static ISiloBuilder AddCosmosGrainStorage( this ISiloBuilder builder, string name, - Type customPartitionKeyProviderType, + Type customDocumentIdProviderType, Action>? configureOptions = null) { - if (customPartitionKeyProviderType != null) + if (customDocumentIdProviderType != null) { - builder.Services.AddSingletonNamedService(name, customPartitionKeyProviderType); + builder.Services.AddSingletonNamedService(name, customDocumentIdProviderType); } builder.Services.AddCosmosGrainStorage(name, configureOptions); @@ -252,7 +253,8 @@ public static IServiceCollection AddCosmosGrainStorage( name)); services.ConfigureNamedOptionForLogging(name); services.TryAddSingleton(sp => sp.GetServiceByName(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME)); - services.TryAddSingleton(); + services.TryAddSingleton(); + services.TryAddFromExisting(); return services.AddSingletonNamedService(name, CosmosStorageFactory.Create) .AddSingletonNamedService(name, (s, n) => (ILifecycleParticipant)s.GetRequiredServiceByName(n)); } diff --git a/src/Azure/Orleans.Persistence.Cosmos/IDocumentIdProvider.cs b/src/Azure/Orleans.Persistence.Cosmos/IDocumentIdProvider.cs new file mode 100644 index 0000000000..e535a46b4b --- /dev/null +++ b/src/Azure/Orleans.Persistence.Cosmos/IDocumentIdProvider.cs @@ -0,0 +1,15 @@ +namespace Orleans.Persistence.Cosmos; + +/// +/// Gets document and partition identifiers for grain state documents. +/// +public interface IDocumentIdProvider +{ + /// + /// Gets the document identifier for the specified grain. + /// + /// The grain state name. + /// The grain identifier. + /// The document id and partition key. + ValueTask<(string DocumentId, string PartitionKey)> GetDocumentIdentifiers(string stateName, GrainId grainId); +} diff --git a/src/Azure/Orleans.Persistence.Cosmos/IPartitionKeyProvider.cs b/src/Azure/Orleans.Persistence.Cosmos/IPartitionKeyProvider.cs deleted file mode 100644 index ea44a14077..0000000000 --- a/src/Azure/Orleans.Persistence.Cosmos/IPartitionKeyProvider.cs +++ /dev/null @@ -1,20 +0,0 @@ -namespace Orleans.Persistence.Cosmos; - -/// -/// Creates a partition key for the provided grain. -/// -public interface IPartitionKeyProvider -{ - /// - /// Creates a partition key for the provided grain. - /// - /// The grain type. - /// The grain identifier. - /// The partition key. - ValueTask GetPartitionKey(string grainType, GrainId grainId); -} - -internal class DefaultPartitionKeyProvider : IPartitionKeyProvider -{ - public ValueTask GetPartitionKey(string grainType, GrainId grainId) => new(CosmosIdSanitizer.Sanitize(grainType)); -} \ No newline at end of file diff --git a/test/Extensions/Tester.Cosmos/PersistenceProviderTests_Cosmos.cs b/test/Extensions/Tester.Cosmos/PersistenceProviderTests_Cosmos.cs index d41dec12ec..9daeb20d8f 100644 --- a/test/Extensions/Tester.Cosmos/PersistenceProviderTests_Cosmos.cs +++ b/test/Extensions/Tester.Cosmos/PersistenceProviderTests_Cosmos.cs @@ -46,10 +46,10 @@ private async Task InitializeStorage() options.ConfigureTestDefaults(); - var pkProvider = new DefaultPartitionKeyProvider(); var clusterOptions = Options.Create(new ClusterOptions { ClusterId = _clusterId, ServiceId = _serviceId }); + var idProvider = new DefaultDocumentIdProvider(clusterOptions); - var store = ActivatorUtilities.CreateInstance(providerRuntime.ServiceProvider, options, clusterOptions, "TestStorage", pkProvider); + var store = ActivatorUtilities.CreateInstance(providerRuntime.ServiceProvider, options, clusterOptions, "TestStorage", idProvider); var lifecycle = ActivatorUtilities.CreateInstance(providerRuntime.ServiceProvider); store.Participate(lifecycle); await lifecycle.OnStart();