feat(cosmos): support custom document IDs and partition keys - #8699

src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs:369

• The /GrainType partition key path compatibility check currently rejects any custom IDocumentIdProvider (anything other than DefaultDocumentIdProvider). That blocks custom document id support for clusters using existing containers partitioned by /GrainType, even if the custom provider returns a compatible partition key (grainType). Consider relaxing this check (or checking compatibility based on the returned partition key) so that custom document ids can be used with legacy /GrainType containers.

                if (_partitionKeyPath == GRAINTYPE_PARTITION_KEY_PATH &&
                    _documentIdProvider is not DefaultDocumentIdProvider)

src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs:370

• This OrleansConfigurationException message still refers to a "Custom partition key provider", but the code now uses IDocumentIdProvider/DefaultDocumentIdProvider. Updating the message will make misconfiguration easier to understand.

                if (_partitionKeyPath == GRAINTYPE_PARTITION_KEY_PATH &&
                    _documentIdProvider is not DefaultDocumentIdProvider)
                    throw new OrleansConfigurationException("Custom partition key provider is not compatible with partition key path set to /GrainType");

src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs:178

• After successfully deleting state (DeleteStateOnClear=true), the code clears ETag/RecordExists but does not reset grainState.State. This differs from ResetGrainState used elsewhere in the provider and can leave stale state in the activation.

                {
                    try

src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs:202

• When DeleteStateOnClear is false, ClearStateAsync overwrites storage with an empty document and sets RecordExists=false, but it no longer resets grainState.State to a fresh instance. This can leave the activation with the pre-clear values even though the record is cleared.

                    await _executor.ExecuteOperation(static args =>
                    {
                        var (self, id, pk, requestOptions) = args;
                        return self._container.DeleteItemAsync<GrainStateEntity<T>>(id, pk, requestOptions);

src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs:168

• When DeleteStateOnClear is enabled and the caller has no ETag, this method returns without clearing the in-memory state (and without normalizing ETag). That can leave the activation holding stale state after a successful ClearStateAsync call. Consider resetting the grain state (without touching storage) before returning.

This issue also appears in the following locations of the same file:

• line 177
• line 199

    public async Task ClearStateAsync<T>(string grainType, GrainId grainId, IGrainState<T> grainState)
    {
        var (id, partitionKey) = await _documentIdProvider.GetDocumentIdentifiers(grainType, grainId);

        LogTraceClearingState(grainType, id, grainId, grainState.ETag, _options.DeleteStateOnClear, _options.ContainerName, partitionKey);

src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs:333

• The /GrainType partition-key-path guard currently rejects any custom IDocumentIdProvider, even if it only customizes the document id while still returning partitionKey == grainType. Since Cosmos uses GrainType as the partition key when the path is /GrainType, custom document ids can still be compatible as long as the partition key remains the grain type. Consider only rejecting custom partition-key providers here (and/or validating that the returned partition key equals grainType at runtime).

                if (_partitionKeyPath == GRAINTYPE_PARTITION_KEY_PATH &&
                    (_documentIdProvider is not DefaultDocumentIdProvider defaultProvider || defaultProvider.HasCustomPartitionKeyProvider))
                    throw new OrleansConfigurationException("Custom document id or partition key providers are not compatible with partition key path set to /GrainType");

src/Azure/Orleans.Persistence.Cosmos/HostingExtensions.cs:40

• The generic AddCosmosGrainStorage/AddCosmosGrainStorageAsDefault overloads no longer provide compile-time validation that TProvider implements IDocumentIdProvider or IPartitionKeyProvider (the constraint is now just class), so invalid types compile and fail only at runtime via AddIdentifierProvider. Consider adding a new constrained generic overload for IDocumentIdProvider and keeping legacy IPartitionKeyProvider-based overloads as [Obsolete] to preserve compile-time safety for the primary API.

    public static ISiloBuilder AddCosmosGrainStorageAsDefault<TProvider>(
        this ISiloBuilder builder,
        Action<CosmosGrainStorageOptions> configureOptions) where TProvider : class
    {
        return builder.AddCosmosGrainStorage<TProvider>(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, configureOptions);
    }

    /// <summary>
    /// Configure silo to use Azure Cosmos DB storage for grain storage using a custom document id provider
    /// or a legacy partition key provider.
    /// </summary>
    /// <typeparam name="TProvider">The document id or partition key provider.</typeparam>
    /// <param name="builder">The silo builder.</param>
    /// <param name="name">The storage provider name.</param>
    /// <param name="configureOptions">The delegate used to configure the provider.</param>
    public static ISiloBuilder AddCosmosGrainStorage<TProvider>(
        this ISiloBuilder builder,
        string name,
        Action<CosmosGrainStorageOptions> configureOptions) where TProvider : class

src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorage.cs:333

• This check rejects any custom IDocumentIdProvider when PartitionKeyPath is "/GrainType", even though customizing only the document id can still be compatible as long as the provider returns PartitionKey == grainType (Cosmos uses the GrainType property as the partition key, independent of document id).

Either relax this restriction for custom IDocumentIdProvider implementations, or validate/enforce that the returned partition key matches grainType when PartitionKeyPath is "/GrainType" (and update the exception message accordingly).

                if (_partitionKeyPath == GRAINTYPE_PARTITION_KEY_PATH &&
                    (_documentIdProvider is not DefaultDocumentIdProvider defaultProvider || defaultProvider.HasCustomPartitionKeyProvider))
                    throw new OrleansConfigurationException("Custom document id or partition key providers are not compatible with partition key path set to /GrainType");

src/Azure/Orleans.Persistence.Cosmos/HostingExtensions.cs:26

• The generic overloads now accept any class (no IDocumentIdProvider/IPartitionKeyProvider constraint), so mis-typed calls compile and only fail at runtime via AddIdentifierProvider’s ArgumentException. That’s a significant loss of type-safety/IDE guidance for a public API.

Consider restoring compile-time safety by splitting into two overloads (one constrained to IDocumentIdProvider, plus an [Obsolete] overload constrained to IPartitionKeyProvider for legacy), or introducing a shared marker interface. This will require regenerating the API surface file.

    public static ISiloBuilder AddCosmosGrainStorageAsDefault<TProvider>(
        this ISiloBuilder builder,
        Action<CosmosGrainStorageOptions> configureOptions) where TProvider : class
    {
        return builder.AddCosmosGrainStorage<TProvider>(ProviderConstants.DEFAULT_STORAGE_PROVIDER_NAME, configureOptions);