Support different throughput options (#60)

* basic changes to handle throughput modes.

* Revert breaking change

* Update ci.yml

* Update publish.yml

.github/workflows/ci.yml

@@ -15,12 +15,12 @@ jobs:
     - name: Setup .NET Core
       uses: actions/setup-dotnet@v1
       with:
-        dotnet-version: 5.0.301
+        dotnet-version: 6.0.x
     # - name: Azure Cosmos Emulator
     #   uses: galvesribeiro/AzureCosmosAction@v1.0.0
     # - name: Start CosmosDB Emulator
     #   run: .\StartEmulator.cmd
     - name: Build
       run: dotnet build --configuration Release
     # - name: Test
-    #   run: dotnet test --configuration Release --no-build
+    #   run: dotnet test --configuration Release --no-build

.github/workflows/publish.yml

@@ -15,7 +15,7 @@ jobs:
     - name: Setup .NET Core
       uses: actions/setup-dotnet@v1
       with:
-        dotnet-version: 5.0.301
+        dotnet-version: 6.0.x
     - name: Pack
       working-directory: src/Orleans.Clustering.CosmosDB
       run: dotnet pack --configuration Release -p:Version=${GITHUB_REF##*/v}
@@ -52,4 +52,4 @@ jobs:
         tag_name: ${{ github.ref }}
         release_name: Release ${{ github.ref }}
         draft: false
-        prerelease: false
+        prerelease: false

.gitignore

@@ -9,6 +9,9 @@
 *.userosscache
 *.sln.docstates
 
+**/CosmosDBTestSecrets.json
+**/cosmosdbtestsecrets.json
+
 # User-specific files (MonoDevelop/Xamarin Studio)
 *.userprefs
 

src/Orleans.Clustering.CosmosDB/CosmosDBMembershipTable.cs

@@ -9,6 +9,7 @@
 using System.Linq;
 using System.Net;
 using System.Threading.Tasks;
+using Orleans.Clustering.CosmosDB.Extensions;
 
 namespace Orleans.Clustering.CosmosDB
 {
@@ -402,15 +403,12 @@ private async Task TryCreateCosmosDBResources()
             containerProperties.IndexingPolicy.ExcludedPaths.Add(new ExcludedPath { Path = "/\"SuspectingTimes\"/*" });
             containerProperties.IndexingPolicy.ExcludedPaths.Add(new ExcludedPath { Path = "/StartTime/*" });
             containerProperties.IndexingPolicy.ExcludedPaths.Add(new ExcludedPath { Path = "/IAmAliveTime/*" });
-            //var consistency = this._options.GetConsistencyLevel();
-            //if (consistency.HasValue)
-            //{
-            //    containerProperties.IndexingPolicy.IndexingMode = consistency.Value;
-            //}
             containerProperties.IndexingPolicy.IndexingMode = IndexingMode.Consistent;
 
+
+
             await dbResponse.CreateContainerIfNotExistsAsync(
-                containerProperties, this._options.CollectionThroughput);
+                containerProperties, this._options.GetThroughputProperties());
         }
 
         public async Task CleanupDefunctSiloEntries(DateTimeOffset beforeDate)

src/Orleans.Clustering.CosmosDB/Extensions/CosmosDBClusteringOptionsExtensions.cs

@@ -0,0 +1,18 @@
+using System;
+using Microsoft.Azure.Cosmos;
+using Orleans.Clustering.CosmosDB.Models;
+
+namespace Orleans.Clustering.CosmosDB.Extensions;
+
+internal static class CosmosDBClusteringOptionsExtensions
+{
+    internal static ThroughputProperties GetThroughputProperties(this CosmosDBClusteringOptions options) =>
+        options.ThroughputMode switch
+        {
+            ThroughputMode.Manual => ThroughputProperties.CreateManualThroughput(options.CollectionThroughput),
+            ThroughputMode.Autoscale => ThroughputProperties.CreateAutoscaleThroughput(
+                options.CollectionThroughput == 400 ? 4000 : options.CollectionThroughput),
+            ThroughputMode.Serverless => null,
+            _ => throw new ArgumentOutOfRangeException(nameof(options.ThroughputMode), $"There is no setup for throughput mode {options.ThroughputMode}")
+        };
+}

src/Orleans.Clustering.CosmosDB/Models/ThroughputMode.cs

@@ -0,0 +1,8 @@
+namespace Orleans.Clustering.CosmosDB.Models;
+
+public enum ThroughputMode
+{
+    Manual,
+    Autoscale,
+    Serverless
+}

src/Orleans.Clustering.CosmosDB/Options/CosmosDBClusteringOptions.cs

@@ -1,6 +1,7 @@
 using Microsoft.Azure.Cosmos;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Converters;
+using Orleans.Clustering.CosmosDB.Models;
 
 namespace Orleans.Clustering.CosmosDB
 {
@@ -10,15 +11,49 @@ public class CosmosDBClusteringOptions
         private const string ORLEANS_CLUSTER_COLLECTION = "OrleansCluster";
         private const int ORLEANS_CLUSTER_COLLECTION_THROUGHPUT = 400;
 
+        /// <summary>
+        /// The <see cref="CosmosClient"/> used for clustering.
+        /// </summary>
         public CosmosClient Client { get; set; }
+
+        /// <summary>
+        /// Gets or sets the cosmos account endpoint URI. This can be retrieved from the Overview section of the Azure Portal.
+        /// This is required if you are authenticating using tokens.
+        /// <remarks>
+        /// In the form of https://{databaseaccount}.documents.azure.com:443/, see: https://docs.microsoft.com/en-us/rest/api/cosmos-db/cosmosdb-resource-uri-syntax-for-rest
+        /// </remarks>
+        /// </summary>
         public string AccountEndpoint { get; set; }
+
+        /// <summary>
+        /// The account key used for a cosmos DB account.
+        /// </summary>
         [Redact]
         public string AccountKey { get; set; }
+
+        /// <summary>
+        /// Tries to create the database and container used for clustering if it does not exist.
+        /// </summary>
         public bool CanCreateResources { get; set; }
+
+        /// <summary>
+        /// The name of the database to use for clustering information.
+        /// </summary>
         public string DB { get; set; } = ORLEANS_DB;
+
+        /// <summary>
+        /// The name of the collection/container to use to store clustering information.
+        /// </summary>
         public string Collection { get; set; } = ORLEANS_CLUSTER_COLLECTION;
+
+        /// <summary>
+        /// The RU throughput used for the collection/container storing clustering information.
+        /// </summary>
         public int CollectionThroughput { get; set; } = ORLEANS_CLUSTER_COLLECTION_THROUGHPUT;
 
+        /// <summary>
+        /// The connection mode to use when connecting to the azure cosmos DB service.
+        /// </summary>
         [JsonConverter(typeof(StringEnumConverter))]
         public ConnectionMode ConnectionMode { get; set; } = ConnectionMode.Direct;
 
@@ -27,6 +62,12 @@ public class CosmosDBClusteringOptions
         /// </summary>
         public bool DropDatabaseOnInit { get; set; }
 
+        /// <summary>
+        /// The throughput mode to use for the collection/container used for clustering.
+        /// </summary>
+        /// <remarks>If the throughput mode is set to Autoscale then the <see cref="CollectionThroughput"/> will need to be at least 4000 RUs.</remarks>
+        public ThroughputMode ThroughputMode { get; set; } = ThroughputMode.Manual;
+
         // TODO: Consistency level for emulator (defaults to Session; https://docs.microsoft.com/en-us/azure/cosmos-db/local-emulator)
         internal IndexingMode? GetConsistencyLevel() => !string.IsNullOrWhiteSpace(this.AccountEndpoint) && this.AccountEndpoint.Contains("localhost") ? (IndexingMode?)IndexingMode.None : null;
     }

src/Orleans.Persistence.CosmosDB/CosmosDBGrainStorage.cs

@@ -16,6 +16,7 @@
 using System.Net;
 using System.Threading;
 using System.Threading.Tasks;
+using Orleans.Persistence.CosmosDB.Extensions;
 
 namespace Orleans.Persistence.CosmosDB
 {
@@ -160,7 +161,7 @@ public async Task ReadStateAsync(string grainType, GrainReference grainReference
                     grainState.State = Activator.CreateInstance(grainState.State.GetType());
                     grainState.RecordExists = false;
                 }
-                
+
                 grainState.ETag = doc.Resource.ETag;
             }
             catch (CosmosException dce)
@@ -317,7 +318,7 @@ public async Task<List<GrainReference>> LookupAsync<K>(string grainType, string
                 var response = await ExecuteWithRetries(async () =>
                 {
                     var pk = new PartitionKey(grainType);
-                    
+
                     var query = this._container.GetItemQueryIterator<GrainStateEntity>(
                         new QueryDefinition($"SELECT * FROM c WHERE c.State.{indexedField} = @key").WithParameter("@key", key),
                         requestOptions: new QueryRequestOptions { PartitionKey = pk }
@@ -406,12 +407,12 @@ private static bool IsNumericType(Type o)
 
         private async Task TryCreateCosmosDBResources()
         {
-            var offerThroughput =
-                    this._options.DatabaseThroughput >= 400
-                    ? (int?)this._options.DatabaseThroughput
-                    : null;
+            var dbThroughput =
+                this._options.DatabaseThroughput >= 400
+                ? (int?)this._options.DatabaseThroughput
+                : null;
 
-            var dbResponse = await this._cosmos.CreateDatabaseIfNotExistsAsync(this._options.DB, offerThroughput);
+            var dbResponse = await this._cosmos.CreateDatabaseIfNotExistsAsync(this._options.DB, dbThroughput);
             var db = dbResponse.Database;
 
             var stateCollection = new ContainerProperties(this._options.Collection, DEFAULT_PARTITION_KEY_PATH);
@@ -432,7 +433,7 @@ private async Task TryCreateCosmosDBResources()
             for (var retry = 0; retry <= maxRetries; ++retry)
             {
                 var collResponse = await db.CreateContainerIfNotExistsAsync(
-                    stateCollection, offerThroughput);
+                    stateCollection, this._options.GetThroughputProperties());
 
                 if (collResponse.StatusCode == HttpStatusCode.OK || collResponse.StatusCode == HttpStatusCode.Created)
                 {

src/Orleans.Persistence.CosmosDB/Extensions/CosmosDBStorageOptionsExtensions.cs

@@ -0,0 +1,18 @@
+using System;
+using Microsoft.Azure.Cosmos;
+using Orleans.Persistence.CosmosDB.Models;
+
+namespace Orleans.Persistence.CosmosDB.Extensions;
+
+internal static class CosmosDBStorageOptionsExtensions
+{
+    internal static ThroughputProperties GetThroughputProperties(this CosmosDBStorageOptions options) =>
+        options.ThroughputMode switch
+        {
+            ThroughputMode.Manual => ThroughputProperties.CreateManualThroughput(options.CollectionThroughput),
+            ThroughputMode.Autoscale => ThroughputProperties.CreateAutoscaleThroughput(
+                options.CollectionThroughput == 400 ? 4000 : options.CollectionThroughput),
+            ThroughputMode.Serverless => null,
+            _ => throw new ArgumentOutOfRangeException(nameof(options.ThroughputMode), $"There is no setup for throughput mode {options.ThroughputMode}")
+        };
+}

src/Orleans.Persistence.CosmosDB/Models/ThroughputMode.cs

@@ -0,0 +1,8 @@
+namespace Orleans.Persistence.CosmosDB.Models;
+
+public enum ThroughputMode
+{
+    Manual,
+    Autoscale,
+    Serverless
+}

src/Orleans.Persistence.CosmosDB/Options/CosmosDBStorageOptions.cs

@@ -3,6 +3,7 @@
 using Newtonsoft.Json.Converters;
 using Orleans.Runtime;
 using System.Collections.Generic;
+using Orleans.Persistence.CosmosDB.Models;
 
 namespace Orleans.Persistence.CosmosDB
 {
@@ -60,6 +61,12 @@ public class CosmosDBStorageOptions
         /// </summary>
         public int InitStage { get; set; } = DEFAULT_INIT_STAGE;
 
+        /// <summary>
+        /// The throughput mode to use for the collection/container used for clustering.
+        /// </summary>
+        /// <remarks>If the throughput mode is set to Autoscale then the <see cref="CollectionThroughput"/> will need to be at least 4000 RUs.</remarks>
+        public ThroughputMode ThroughputMode { get; set; } = ThroughputMode.Manual;
+
         public const int DEFAULT_INIT_STAGE = ServiceLifecycleStage.ApplicationServices;
 
         // TODO: Consistency level for emulator (defaults to Session; https://docs.microsoft.com/en-us/azure/cosmos-db/local-emulator)

test/Orleans.CosmosDB.Tests/Extensions/CosmosDBClusteringOptionsExtensionsTests.cs

@@ -0,0 +1,77 @@
+using Orleans.Clustering.CosmosDB;
+using Orleans.Clustering.CosmosDB.Extensions;
+using Orleans.Clustering.CosmosDB.Models;
+using Xunit;
+
+namespace Orleans.CosmosDB.Tests.Extensions;
+
+public class CosmosDBClusteringOptionsExtensionsTests
+{
+    [Fact]
+    public void GetThroughputProperties_ManualMode_SetsThroughputCorrectly()
+    {
+        //Arrange
+        var options = new CosmosDBClusteringOptions
+        {
+            ThroughputMode = ThroughputMode.Manual
+        };
+
+        //Act
+        var properties = options.GetThroughputProperties();
+
+        //Assert
+        Assert.Equal(400, properties.Throughput);
+        Assert.Null(properties.AutoscaleMaxThroughput);
+    }
+
+    [Fact]
+    public void GetThroughputProperties_AutoscaleModeDefaultThroughput_CreatesAutoscaleThroughputWith4000RUs()
+    {
+        //Arrange
+        var options = new CosmosDBClusteringOptions
+        {
+            ThroughputMode = ThroughputMode.Autoscale
+        };
+
+        //Act
+        var properties = options.GetThroughputProperties();
+
+        //Assert
+        Assert.Equal(4000, properties.AutoscaleMaxThroughput);
+        Assert.Null(properties.Throughput);
+    }
+
+    [Fact]
+    public void GetThroughputProperties_AutoscaleModeCustomThroughput_CreatesAutoscaleThroughputWithCustomRUs()
+    {
+        //Arrange
+        var options = new CosmosDBClusteringOptions
+        {
+            ThroughputMode = ThroughputMode.Autoscale,
+            CollectionThroughput = 6000
+        };
+
+        //Act
+        var properties = options.GetThroughputProperties();
+
+        //Assert
+        Assert.Equal(6000, properties.AutoscaleMaxThroughput);
+        Assert.Null(properties.Throughput);
+    }
+
+    [Fact]
+    public void GetThroughputProperties_ServerlessMode_ReturnsNullThroughput()
+    {
+        //Arrange
+        var options = new CosmosDBClusteringOptions
+        {
+            ThroughputMode = ThroughputMode.Serverless
+        };
+
+        //Act
+        var properties = options.GetThroughputProperties();
+
+        //Assert
+        Assert.Null(properties);
+    }
+}