dotnet / dotnet/orleans

Support hierarchical partition keys in Cosmos DB grain storage

Open
#9,899 2 comments 1 reaction 0 assignees View on GitHub
Dominant language
C#
Stars
10.9k
Forks
2.1k
Avg merge
14h 42m
Merged PRs (30d)
354

Description

## Summary

The current `IPartitionKeyProvider` interface returns a `string`, which limits Cosmos DB grain storage to single-level partition keys. With [hierarchical partition keys](https://learn.microsoft.com/en-us/azure/cosmos-db/hierarchical-partition-keys) now GA in Cosmos DB, it would be valuable to support multi-level partitioning for Orleans grain state.

## Background

I'm returning to Orleans after a hiatus and starting a new project that would benefit from hierarchical partitioning. The current implementation wraps the string result in `new PartitionKey(partitionKey)`:

```csharp
// Current implementation in CosmosGrainStorage.cs
private ValueTask BuildPartitionKey(string grainType, GrainId grainId) =>
_partitionKeyProvider.GetPartitionKey(grainType, grainId);

// Later used as:
var pk = new PartitionKey(partitionKey);
```

## Proposal

Add a virtual default interface method (DIM) to `IPartitionKeyProvider` that returns a `PartitionKey` directly, allowing implementations to construct hierarchical keys:

```csharp
public interface IPartitionKeyProvider
{
// Existing method (unchanged for backwards compatibility)
ValueTask GetPartitionKey(string grainType, GrainId grainId);

// New DIM for hierarchical partition key support
ValueTask GetPartitionKeyValue(string grainType, GrainId grainId)
=> new(new PartitionKey(GetPartitionKey(grainType, grainId).Result));
}
```

Then update `CosmosGrainStorage` to call `GetPartitionKeyValue()` instead of wrapping the string manually.

### Example Usage

```csharp
public class TenantAwarePartitionKeyProvider : IPartitionKeyProvider
{
public ValueTask GetPartitionKey(string grainType, GrainId grainId)
=> new(grainType); // Legacy fallback

public ValueTask GetPartitionKeyValue(string grainType, GrainId grainId)
{
var tenantId = ExtractTenantFromGrainId(grainId);
// Hierarchical: TenantId -> GrainType -> Region
return new(new PartitionKeyBuilder()
.Add(tenantId)
.Add(grainType)
.Build());
}
}
```

## Benefits

1. **Multi-tenant scenarios**: Partition first by tenant, then by grain type
2. **Better data distribution**: Avoid hot partitions in high-volume scenarios
3. **Backwards compatible**: Existing implementations continue to work unchanged via the DIM
4. **Cosmos DB best practices**: Aligns with [Microsoft's guidance on partition key design](https://learn.microsoft.com/en-us/azure/cosmos-db/partitioning-overview#choose-a-partition-key)

## Considerations

- `CosmosStorageOptions.PartitionKeyPath` would need to support hierarchical paths (e.g., `/TenantId/GrainType`) (we should be able to define GrainStateEntity as an `object` w/ nested properties rather than a string)
- Container creation logic may need updates for hierarchical key configuration
- Documentation updates for multi-level partitioning patterns

## Environment

- Orleans version: 10.x
- Target: `Orleans.Persistence.Cosmos`

Happy to submit a PR if this direction makes sense.

Contributor guide

Open the contributing guide

Research direction

Start with IPartitionKeyProvider and CosmosGrainStorage.cs to trace the current string partition-key flow and the proposed direct PartitionKey path. Then inspect CosmosStorageOptions.PartitionKeyPath, GrainStateEntity, and container creation logic for hierarchical configuration. Done means supporting hierarchical keys while preserving existing providers, with the related configuration and documentation updates identified in the issue.

Written by the indexing model from the issue text.

Assessment

Tech stack
azure, csharp
Domain
databases, distributed-systems
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.