From 98c8e2801c23c7a5e008942989a0a24da22017df Mon Sep 17 00:00:00 2001
From: Ramkumar Chinchani <45800463+rchincha@users.noreply.github.com>
Date: Thu, 12 Sep 2024 08:26:43 -0700
Subject: [PATCH] fix: work around AWS S3 limits (#2629)

Fixes issue #2627

We get/put metadata in dynamodb and it appears there are limits enforced
by AWS S3 APIs.

https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_BatchGetItem.html

If you request more than 100 items, BatchGetItem returns a
ValidationException with the message "Too many items requested for the
BatchGetItem call."

Signed-off-by: Ramkumar Chinchani <rchincha@cisco.com>
---
 pkg/meta/dynamodb/dynamodb.go | 40 +++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 13 deletions(-)

diff --git a/pkg/meta/dynamodb/dynamodb.go b/pkg/meta/dynamodb/dynamodb.go
index 4bf96484..c0f80274 100644
--- a/pkg/meta/dynamodb/dynamodb.go
+++ b/pkg/meta/dynamodb/dynamodb.go
@@ -1970,24 +1970,38 @@ func (dwr DynamoDB) DeleteUserData(ctx context.Context) error {
 	return err
 }
 
+const AwsS3BatchLimit = 100
+
 func (dwr *DynamoDB) fetchImageMetaAttributesByDigest(ctx context.Context, digests []string,
 ) ([]map[string]types.AttributeValue, error) {
-	resp, err := dwr.Client.BatchGetItem(ctx, &dynamodb.BatchGetItemInput{
-		RequestItems: map[string]types.KeysAndAttributes{
-			dwr.ImageMetaTablename: {
-				Keys: getBatchImageKeys(digests),
+	// AWS S3 as a limit (=100) on number of keys that can retrieved in one
+	// request, so break it up
+	batchedResp := []map[string]types.AttributeValue{}
+
+	for start := 0; start < len(digests); {
+		size := min(len(digests)-start, AwsS3BatchLimit)
+		end := start + size
+
+		resp, err := dwr.Client.BatchGetItem(ctx, &dynamodb.BatchGetItemInput{
+			RequestItems: map[string]types.KeysAndAttributes{
+				dwr.ImageMetaTablename: {
+					Keys: getBatchImageKeys(digests[start:end]),
+				},
 			},
-		},
-	})
-	if err != nil {
-		return nil, err
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		if len(resp.Responses[dwr.ImageMetaTablename]) != size {
+			return nil, zerr.ErrImageMetaNotFound
+		}
+
+		batchedResp = append(batchedResp, resp.Responses[dwr.ImageMetaTablename]...)
+		start = end
 	}
 
-	if len(resp.Responses[dwr.ImageMetaTablename]) != len(digests) {
-		return nil, zerr.ErrImageMetaNotFound
-	}
-
-	return resp.Responses[dwr.ImageMetaTablename], nil
+	return batchedResp, nil
 }
 
 func getBatchImageKeys(digests []string) []map[string]types.AttributeValue {