Make provisioned read and write units configurable

DevOpsProUK · Jun 1, 2016 · 363d590 · 363d590
1 parent 14892f0
commit 363d590
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 6 deletions.
diff --git a/dynamodb/dynamo_lock_constants.go b/dynamodb/dynamo_lock_constants.go
@@ -18,3 +18,6 @@ const SLEEP_BETWEEN_TABLE_LOCK_ACQUIRE_ATTEMPTS = 10 * time.Second
 
 const DEFAULT_TABLE_NAME = "terragrunt_locks"
 const DEFAULT_AWS_REGION = "us-east-1"
+
+const DEFAULT_READ_CAPACITY_UNITS = 1
+const DEFAULT_WRITE_CAPACITY_UNITS = 1
diff --git a/dynamodb/dynamo_lock_table.go b/dynamodb/dynamo_lock_table.go
@@ -19,7 +19,7 @@ func createLockTableIfNecessary(tableName string, client *dynamodb.DynamoDB) err
 
 	if !tableExists {
 		util.Logger.Printf("Lock table %s does not exist in DynamoDB. Will need to create it just this first time.", tableName)
-		return createLockTable(tableName, client)
+		return createLockTable(tableName, DEFAULT_READ_CAPACITY_UNITS, DEFAULT_WRITE_CAPACITY_UNITS, client)
 	}
 
 	return nil
@@ -41,7 +41,7 @@ func lockTableExistsAndIsActive(tableName string, client *dynamodb.DynamoDB) (bo
 
 // Create a lock table in DynamoDB and wait until it is in "active" state. If the table already exists, merely wait
 // until it is in "active" state.
-func createLockTable(tableName string, client *dynamodb.DynamoDB) error {
+func createLockTable(tableName string, readCapacityUnits int, writeCapacityUnits int, client *dynamodb.DynamoDB) error {
 	util.Logger.Printf("Creating table %s in DynamoDB", tableName)
 
 	attributeDefinitions := []*dynamodb.AttributeDefinition{
@@ -56,7 +56,10 @@ func createLockTable(tableName string, client *dynamodb.DynamoDB) error {
 		TableName: aws.String(tableName),
 		AttributeDefinitions: attributeDefinitions,
 		KeySchema: keySchema,
-		ProvisionedThroughput: &dynamodb.ProvisionedThroughput{ReadCapacityUnits: aws.Int64(1), WriteCapacityUnits: aws.Int64(1)},
+		ProvisionedThroughput: &dynamodb.ProvisionedThroughput{
+			ReadCapacityUnits: aws.Int64(int64(readCapacityUnits)),
+			WriteCapacityUnits: aws.Int64(int64(writeCapacityUnits)),
+		},
 	})
 
 	if err != nil {

diff --git a/dynamodb/dynamo_lock_table_test.go b/dynamodb/dynamo_lock_table_test.go
@@ -177,8 +177,10 @@ func TestWriteItemToLockTableUntilSuccessItemAlreadyExistsButGetsDeleted(t *test
 func TestWriteItemToLockTableConcurrency(t * testing.T) {
 	t.Parallel()
 
-	// First, create a table
-	withLockTable(t, func(tableName string, client *dynamodb.DynamoDB) {
+	concurrency := 20
+
+	// First, create a table with enough read and write units to ensure this test doesn't get throttled
+	withLockTableProvisionedUnits(t, concurrency, concurrency, func(tableName string, client *dynamodb.DynamoDB) {
 		itemId := uniqueId()
 
 		// Use a WaitGroup to ensure the test doesn't exit before all goroutines finish.
@@ -190,7 +192,7 @@ func TestWriteItemToLockTableConcurrency(t * testing.T) {
 
 		// Launch a bunch of goroutines who will all try to create the same item at more or less the same time.
 		// Only one of the goroutines should succeed.
-		for i := 0; i < 20; i++ {
+		for i := 0; i < concurrency; i++ {
 			waitGroup.Add(1)
 			go func() {
 				defer waitGroup.Done()

diff --git a/dynamodb/dynamo_lock_test_utils.go b/dynamodb/dynamo_lock_test_utils.go
@@ -90,5 +90,16 @@ func withLockTable(t *testing.T, action func(tableName string, client *dynamodb.
 	assert.Nil(t, err)
 	defer cleanupTable(t, tableName, client)
 
+	action(tableName, client)
+}
+
+func withLockTableProvisionedUnits(t *testing.T, readCapacityUnits int, writeCapacityUnits int, action func(tableName string, client *dynamodb.DynamoDB)) {
+	client := createDynamoDbClientForTest(t)
+	tableName := uniqueTableNameForTest()
+
+	err := createLockTable(tableName, readCapacityUnits, writeCapacityUnits, client)
+	assert.Nil(t, err)
+	defer cleanupTable(t, tableName, client)
+
 	action(tableName, client)
 }