Use slice to batch writeRequests

Signed-off-by: Daniel Deluiggi <[email protected]>

Author: danielblando

pkg/ring/kv/dynamodb/dynamodb.go

@@ -171,28 +171,31 @@ func (kv dynamodbKV) Put(ctx context.Context, key dynamodbKey, data []byte) erro
 }
 
 func (kv dynamodbKV) Batch(ctx context.Context, put map[dynamodbKey][]byte, delete []dynamodbKey) error {
-	var writeRequests []*dynamodb.WriteRequest
+	writeRequestSize := len(put) + len(delete)
+	if writeRequestSize == 0 {
+		return nil
+	}
+
+	writeRequests := make([]*dynamodb.WriteRequest, writeRequestSize)
+	writeRequestsIndex := 0
 	for key, data := range put {
 		item := kv.generatePutItemRequest(key, data)
-		writeRequests = append(writeRequests, &dynamodb.WriteRequest{
+		writeRequests[writeRequestsIndex] = &dynamodb.WriteRequest{
 			PutRequest: &dynamodb.PutRequest{
 				Item: item,
 			},
-		})
+		}
+		writeRequestsIndex++
 	}
 
 	for _, key := range delete {
 		item := generateItemKey(key)
-
-		writeRequests = append(writeRequests, &dynamodb.WriteRequest{
+		writeRequests[writeRequestsIndex] = &dynamodb.WriteRequest{
 			DeleteRequest: &dynamodb.DeleteRequest{
 				Key: item,
 			},
-		})
-	}
-
-	if len(writeRequests) == 0 {
-		return nil
+		}
+		writeRequestsIndex++
 	}
 
 	input := &dynamodb.BatchWriteItemInput{
@@ -202,12 +205,15 @@ func (kv dynamodbKV) Batch(ctx context.Context, put map[dynamodbKey][]byte, dele
 	}
 
 	resp, err := kv.ddbClient.BatchWriteItemWithContext(ctx, input)
+	if err != nil {
+		return err
+	}
 
 	if resp.UnprocessedItems != nil && len(resp.UnprocessedItems) > 0 {
 		return fmt.Errorf("error processing batch request for %s requests", resp.UnprocessedItems)
 	}
 
-	return err
+	return nil
 }
 
 func (kv dynamodbKV) generatePutItemRequest(key dynamodbKey, data []byte) map[string]*dynamodb.AttributeValue {

pkg/ring/kv/dynamodb/dynamodb_test.go

@@ -2,6 +2,7 @@ package dynamodb
 
 import (
 	"context"
+	"fmt"
 	"strconv"
 	"testing"
 	"time"
@@ -60,13 +61,13 @@ func Test_Batch(t *testing.T) {
 	delete := []dynamodbKey{ddbKeyDelete}
 
 	ddbClientMock := &mockDynamodb{
-		batchWriteItem: func(input *dynamodb.BatchWriteItemInput) *dynamodb.BatchWriteItemOutput {
+		batchWriteItem: func(input *dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error) {
 			require.NotNil(t, input.RequestItems[tableName])
-			require.True(t, len(input.RequestItems[tableName]) == 2)
+			require.EqualValues(t, 2, len(input.RequestItems[tableName]))
 			require.True(t,
 				(checkPutRequestForItem(input.RequestItems[tableName][0], ddbKeyUpdate) || checkPutRequestForItem(input.RequestItems[tableName][1], ddbKeyUpdate)) &&
 					(checkDeleteRequestForItem(input.RequestItems[tableName][0], ddbKeyDelete) || checkDeleteRequestForItem(input.RequestItems[tableName][1], ddbKeyDelete)))
-			return &dynamodb.BatchWriteItemOutput{}
+			return &dynamodb.BatchWriteItemOutput{}, nil
 		},
 	}
 
@@ -84,7 +85,7 @@ func Test_EmptyBatch(t *testing.T) {
 	require.NoError(t, err)
 }
 
-func Test_Batch_Error(t *testing.T) {
+func Test_Batch_UnprocessedItems(t *testing.T) {
 	tableName := "TEST"
 	ddbKeyDelete := dynamodbKey{
 		primaryKey: "PKDelete",
@@ -93,14 +94,14 @@ func Test_Batch_Error(t *testing.T) {
 	delete := []dynamodbKey{ddbKeyDelete}
 
 	ddbClientMock := &mockDynamodb{
-		batchWriteItem: func(input *dynamodb.BatchWriteItemInput) *dynamodb.BatchWriteItemOutput {
+		batchWriteItem: func(input *dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error) {
 			return &dynamodb.BatchWriteItemOutput{
 				UnprocessedItems: map[string][]*dynamodb.WriteRequest{
 					tableName: {&dynamodb.WriteRequest{
 						PutRequest: &dynamodb.PutRequest{Item: generateItemKey(ddbKeyDelete)}},
 					},
 				},
-			}
+			}, nil
 		},
 	}
 
@@ -109,6 +110,25 @@ func Test_Batch_Error(t *testing.T) {
 	require.Errorf(t, err, "error processing batch dynamodb")
 }
 
+func Test_Batch_Error(t *testing.T) {
+	tableName := "TEST"
+	ddbKeyDelete := dynamodbKey{
+		primaryKey: "PKDelete",
+		sortKey:    "SKDelete",
+	}
+	delete := []dynamodbKey{ddbKeyDelete}
+
+	ddbClientMock := &mockDynamodb{
+		batchWriteItem: func(input *dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error) {
+			return &dynamodb.BatchWriteItemOutput{}, fmt.Errorf("mocked error")
+		},
+	}
+
+	ddb := newDynamodbClientMock(tableName, ddbClientMock, 5*time.Hour)
+	err := ddb.Batch(context.TODO(), nil, delete)
+	require.Errorf(t, err, "mocked error")
+}
+
 func checkPutRequestForItem(request *dynamodb.WriteRequest, key dynamodbKey) bool {
 	return request.PutRequest != nil &&
 		request.PutRequest.Item[primaryKey] != nil &&
@@ -127,7 +147,7 @@ func checkDeleteRequestForItem(request *dynamodb.WriteRequest, key dynamodbKey)
 
 type mockDynamodb struct {
 	putItem        func(input *dynamodb.PutItemInput) *dynamodb.PutItemOutput
-	batchWriteItem func(input *dynamodb.BatchWriteItemInput) *dynamodb.BatchWriteItemOutput
+	batchWriteItem func(input *dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error)
 
 	dynamodbiface.DynamoDBAPI
 }
@@ -137,7 +157,7 @@ func (m *mockDynamodb) PutItemWithContext(_ context.Context, input *dynamodb.Put
 }
 
 func (m *mockDynamodb) BatchWriteItemWithContext(_ context.Context, input *dynamodb.BatchWriteItemInput, _ ...request.Option) (*dynamodb.BatchWriteItemOutput, error) {
-	return m.batchWriteItem(input), nil
+	return m.batchWriteItem(input)
 }
 
 func newDynamodbClientMock(tableName string, mock *mockDynamodb, ttl time.Duration) *dynamodbKV {