Explicit setting for UseProvisionedThroughput (#6328)

dotnet · Feb 22, 2020 · 38e1e86 · 38e1e86
1 parent 9fac0dd
commit 38e1e86
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 14 deletions.
diff --git a/src/AWS/Orleans.Persistence.DynamoDB/Options/DynamoDBStorageOptions.cs b/src/AWS/Orleans.Persistence.DynamoDB/Options/DynamoDBStorageOptions.cs
@@ -25,10 +25,15 @@ public class DynamoDBStorageOptions
         public string SecretKey { get; set; }
 
         /// <summary>
-        /// DynamoDB Service name 
+        /// DynamoDB Service name
         /// </summary>
         public string Service { get; set; }
 
+        /// <summary>
+        /// Use Provisioned Throughput for tables
+        /// </summary>
+        public bool UseProvisionedThroughput { get; set; } = true;
+
         /// <summary>
         /// Read capacity unit for DynamoDB storage
         /// </summary>
@@ -94,13 +99,16 @@ public void ValidateConfiguration()
                 throw new OrleansConfigurationException(
                     $"Configuration for DynamoDBGrainStorage {this.name} is invalid. {nameof(this.options.TableName)} is not valid.");
 
-            if (this.options.ReadCapacityUnits == 0)
-                throw new OrleansConfigurationException(
-                    $"Configuration for DynamoDBGrainStorage {this.name} is invalid. {nameof(this.options.ReadCapacityUnits)} is not valid.");
+            if (this.options.UseProvisionedThroughput)
+            {
+                if (this.options.ReadCapacityUnits == 0)
+                    throw new OrleansConfigurationException(
+                        $"Configuration for DynamoDBGrainStorage {this.name} is invalid. {nameof(this.options.ReadCapacityUnits)} is not valid.");
 
-            if (this.options.WriteCapacityUnits == 0)
-                throw new OrleansConfigurationException(
-                    $"Configuration for DynamoDBGrainStorage {this.name} is invalid. {nameof(this.options.WriteCapacityUnits)} is not valid.");
+                if (this.options.WriteCapacityUnits == 0)
+                    throw new OrleansConfigurationException(
+                        $"Configuration for DynamoDBGrainStorage {this.name} is invalid. {nameof(this.options.WriteCapacityUnits)} is not valid.");
+            }
         }
     }
 }
diff --git a/src/AWS/Orleans.Persistence.DynamoDB/Provider/DynamoDBGrainStorage.cs b/src/AWS/Orleans.Persistence.DynamoDB/Provider/DynamoDBGrainStorage.cs
@@ -83,7 +83,7 @@ public async Task Init(CancellationToken ct)
                 this.logger.LogInformation((int)ErrorCode.StorageProviderBase, $"AWS DynamoDB Grain Storage {this.name} is initializing: {initMsg}");
 
                 this.storage = new DynamoDBStorage(this.loggerFactory, this.options.Service, this.options.AccessKey, this.options.SecretKey,
-                 this.options.ReadCapacityUnits, this.options.WriteCapacityUnits);
+                 this.options.ReadCapacityUnits, this.options.WriteCapacityUnits, this.options.UseProvisionedThroughput);
 
                 await storage.InitializeTable(this.options.TableName,
                     new List<KeySchemaElement>
@@ -187,7 +187,7 @@ private async Task WriteStateInternal(IGrainState grainState, GrainStateRecord r
         {
             var fields = new Dictionary<string, AttributeValue>();
             if (this.options.TimeToLive.HasValue)
-            {                                
+            {
                 fields.Add(GRAIN_TTL_PROPERTY_NAME, new AttributeValue { N = ((DateTimeOffset)DateTime.UtcNow.Add(this.options.TimeToLive.Value)).ToUnixTimeSeconds().ToString() });
             }
 

diff --git a/src/AWS/Shared/Storage/DynamoDBStorage.cs b/src/AWS/Shared/Storage/DynamoDBStorage.cs
@@ -38,6 +38,7 @@ internal class DynamoDBStorage
         public const int DefaultWriteCapacityUnits = 5;
         private int readCapacityUnits = DefaultReadCapacityUnits;
         private int writeCapacityUnits = DefaultWriteCapacityUnits;
+        private readonly bool useProvisionedThroughput;
         private AmazonDynamoDBClient ddbClient;
         private ILogger Logger;
 
@@ -50,17 +51,20 @@ internal class DynamoDBStorage
         /// <param name="service"></param>
         /// <param name="readCapacityUnits"></param>
         /// <param name="writeCapacityUnits"></param>
+        /// <param name="useProvisionedThroughput"></param>
         public DynamoDBStorage(ILoggerFactory loggerFactory, string service,
             string accessKey = "", string secretKey = "",
             int readCapacityUnits = DefaultReadCapacityUnits,
-            int writeCapacityUnits = DefaultWriteCapacityUnits)
+            int writeCapacityUnits = DefaultWriteCapacityUnits,
+            bool useProvisionedThroughput = true)
         {
             if (service == null) throw new ArgumentNullException(nameof(service));
             this.accessKey = accessKey;
             this.secretKey = secretKey;
             this.service = service;
             this.readCapacityUnits = readCapacityUnits;
             this.writeCapacityUnits = writeCapacityUnits;
+            this.useProvisionedThroughput = useProvisionedThroughput;
             Logger = loggerFactory.CreateLogger<DynamoDBStorage>();
             CreateClient();
         }
@@ -127,14 +131,13 @@ private async Task<TableDescription> GetTableDescription(string tableName)
 
         private async Task CreateTable(string tableName, List<KeySchemaElement> keys, List<AttributeDefinition> attributes, List<GlobalSecondaryIndex> secondaryIndexes = null, string ttlAttributeName = null)
         {
-            var useProvisionedThroughput = readCapacityUnits > 0 && writeCapacityUnits > 0;
             var request = new CreateTableRequest
             {
                 TableName = tableName,
                 AttributeDefinitions = attributes,
                 KeySchema = keys,
-                BillingMode = useProvisionedThroughput ? BillingMode.PROVISIONED : BillingMode.PAY_PER_REQUEST,
-                ProvisionedThroughput = useProvisionedThroughput ? new ProvisionedThroughput
+                BillingMode = this.useProvisionedThroughput ? BillingMode.PROVISIONED : BillingMode.PAY_PER_REQUEST,
+                ProvisionedThroughput = this.useProvisionedThroughput ? new ProvisionedThroughput
                 {
                     ReadCapacityUnits = readCapacityUnits,
                     WriteCapacityUnits = writeCapacityUnits
@@ -143,7 +146,7 @@ private async Task CreateTable(string tableName, List<KeySchemaElement> keys, Li
 
             if (secondaryIndexes != null && secondaryIndexes.Count > 0)
             {
-                if (useProvisionedThroughput)
+                if (this.useProvisionedThroughput)
                 {
                     var indexThroughput = new ProvisionedThroughput {ReadCapacityUnits = readCapacityUnits, WriteCapacityUnits = writeCapacityUnits};
                     secondaryIndexes.ForEach(i =>