Throttle DynamoDB event migration based on provisioned capacity (#8468)

lib/backend/helpers.go

@@ -150,7 +150,7 @@ func RunWhileLocked(ctx context.Context, backend Backend, lockName string, ttl t
 			select {
 			case <-time.After(refreshAfter):
 				if err := lock.resetTTL(ctx, backend); err != nil {
-					defer cancelFunction()
+					cancelFunction()
 					log.Errorf("%v", err)
 					return
 				}

lib/events/dynamoevents/dynamoevents.go

@@ -1342,6 +1342,26 @@ func (l *Log) convertFieldsToDynamoMapFormat(ctx context.Context) error {
 	return trace.Wrap(l.migrateMatchingEvents(ctx, filterExpr, transformEvent))
 }
 
+func (l *Log) approximateOptimalMigrationWorkers() (int32, error) {
+	req := dynamodb.DescribeTableInput{TableName: aws.String(l.Tablename)}
+	table, err := l.svc.DescribeTable(&req)
+	if err != nil {
+		return 0, trace.Wrap(err)
+	}
+
+	// calculate the throughput, accounting for r/w bottlenecks
+	provisioned := table.Table.ProvisionedThroughput
+	if provisioned == nil || provisioned.ReadCapacityUnits == nil || provisioned.WriteCapacityUnits == nil {
+		return maxMigrationWorkers, nil
+	}
+	throughput := utils.MinInt64(*provisioned.ReadCapacityUnits, *provisioned.WriteCapacityUnits)
+
+	// divide throughput by batch size rounding upwards and then take 75% of that
+	optimalWorkers := (throughput + (DynamoBatchSize - 1)) / DynamoBatchSize * 3 / 4
+	clamped := utils.MinInt64(utils.MaxInt64(optimalWorkers, 1), maxMigrationWorkers)
+	return int32(clamped), nil
+}
+
 // migrateMatchingEvents walks existing events that match the given filter
 // expression and transforms them using the provided transform function.
 //
@@ -1358,7 +1378,12 @@ func (l *Log) migrateMatchingEvents(ctx context.Context, filterExpr string, tran
 	var startKey map[string]*dynamodb.AttributeValue
 	workerCounter := atomic.NewInt32(0)
 	totalProcessed := atomic.NewInt32(0)
-	workerErrors := make(chan error, maxMigrationWorkers)
+	migrationWorkers, err := l.approximateOptimalMigrationWorkers()
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	workerErrors := make(chan error, migrationWorkers)
 	workerBarrier := sync.WaitGroup{}
 
 	for {
@@ -1377,7 +1402,7 @@ func (l *Log) migrateMatchingEvents(ctx context.Context, filterExpr string, tran
 			// for any missed events after an appropriate grace period which is far worse.
 			ConsistentRead: aws.Bool(true),
 			// `DynamoBatchSize*maxMigrationWorkers` is the maximum concurrent event uploads.
-			Limit:            aws.Int64(DynamoBatchSize * maxMigrationWorkers),
+			Limit:            aws.Int64(DynamoBatchSize * int64(migrationWorkers)),
 			TableName:        aws.String(l.Tablename),
 			FilterExpression: aws.String(filterExpr),
 		}
@@ -1390,7 +1415,7 @@ func (l *Log) migrateMatchingEvents(ctx context.Context, filterExpr string, tran
 			return trace.Wrap(convertError(err))
 		}
 
-		writeRequests := make([]*dynamodb.WriteRequest, 0, DynamoBatchSize*maxMigrationWorkers)
+		writeRequests := make([]*dynamodb.WriteRequest, 0, DynamoBatchSize*migrationWorkers)
 
 		// For every item processed by this scan iteration we generate a write request.
 		for _, item := range scanOut.Items {
@@ -1420,7 +1445,7 @@ func (l *Log) migrateMatchingEvents(ctx context.Context, filterExpr string, tran
 			writeRequests = writeRequests[top:]
 
 			// Don't exceed maximum workers.
-			for workerCounter.Load() >= maxMigrationWorkers {
+			for workerCounter.Load() >= migrationWorkers {
 				select {
 				case <-time.After(time.Millisecond * 50):
 				case <-ctx.Done():

lib/utils/cmp.go

@@ -0,0 +1,33 @@
+/*
+Copyright 2021 Gravitational, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package utils
+
+// MinInt64 returns the numerically lesser of two signed 64-bit integers.
+func MinInt64(x, y int64) int64 {
+	if x > y {
+		return y
+	}
+	return x
+}
+
+// MaxInt64 returns the numerically greater of two signed 64-bit integers.
+func MaxInt64(x, y int64) int64 {
+	if x > y {
+		return x
+	}
+	return y
+}