[Encoding ]Request with pagination

common/aws/dynamodb/client.go

@@ -37,6 +37,11 @@ type Item = map[string]types.AttributeValue
 type Key = map[string]types.AttributeValue
 type ExpresseionValues = map[string]types.AttributeValue
 
+type QueryResult struct {
+	Items            []map[string]types.AttributeValue
+	LastEvaluatedKey map[string]types.AttributeValue
+}
+
 type Client struct {
 	dynamoClient *dynamodb.Client
 	logger       common.Logger
@@ -161,6 +166,34 @@ func (c *Client) QueryIndex(ctx context.Context, tableName string, indexName str
 	return response.Items, nil
 }
 
+// QueryIndexWithPagination returns all items in the index that match the given key
+// Results are limited to the given limit and the pagination token is returned
+func (c *Client) QueryIndexWithPagination(ctx context.Context, tableName string, indexName string, keyCondition string, expAttributeValues ExpresseionValues, limit int32, exclusiveStartKey map[string]types.AttributeValue) (QueryResult, error) {
+	queryInput := &dynamodb.QueryInput{
+		TableName:                 aws.String(tableName),
+		IndexName:                 aws.String(indexName),
+		KeyConditionExpression:    aws.String(keyCondition),
+		ExpressionAttributeValues: expAttributeValues,
+		Limit:                     &limit,
+	}
+
+	// If a pagination token was provided, set it as the ExclusiveStartKey
+	if exclusiveStartKey != nil {
+		queryInput.ExclusiveStartKey = exclusiveStartKey
+	}
+
+	response, err := c.dynamoClient.Query(ctx, queryInput)
+	if err != nil {
+		return QueryResult{}, err
+	}
+
+	// Return the items and the pagination token
+	return QueryResult{
+		Items:            response.Items,
+		LastEvaluatedKey: response.LastEvaluatedKey,
+	}, nil
+}
+
 func (c *Client) DeleteItem(ctx context.Context, tableName string, key Key) error {
 	_, err := c.dynamoClient.DeleteItem(ctx, &dynamodb.DeleteItemInput{Key: key, TableName: aws.String(tableName)})
 	if err != nil {

common/aws/dynamodb/client_test.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strconv"
 	"testing"
+	"time"
 
 	commonaws "github.com/Layr-Labs/eigenda/common/aws"
 	commondynamodb "github.com/Layr-Labs/eigenda/common/aws/dynamodb"
@@ -79,14 +81,48 @@ func teardown() {
 func createTable(t *testing.T, tableName string) {
 	ctx := context.Background()
 	tableDescription, err := test_utils.CreateTable(ctx, clientConfig, tableName, &dynamodb.CreateTableInput{
-		AttributeDefinitions: []types.AttributeDefinition{{
-			AttributeName: aws.String("MetadataKey"),
-			AttributeType: types.ScalarAttributeTypeS,
-		}},
-		KeySchema: []types.KeySchemaElement{{
-			AttributeName: aws.String("MetadataKey"),
-			KeyType:       types.KeyTypeHash,
-		}},
+		AttributeDefinitions: []types.AttributeDefinition{
+			{
+				AttributeName: aws.String("MetadataKey"),
+				AttributeType: types.ScalarAttributeTypeS,
+			},
+			{
+				AttributeName: aws.String("BlobStatus"),
+				AttributeType: types.ScalarAttributeTypeN, // Assuming BlobStatus is a numeric value
+			},
+			{
+				AttributeName: aws.String("RequestedAt"),
+				AttributeType: types.ScalarAttributeTypeN, // Assuming RequestedAt is a string representing a timestamp
+			},
+		},
+		KeySchema: []types.KeySchemaElement{
+			{
+				AttributeName: aws.String("MetadataKey"),
+				KeyType:       types.KeyTypeHash,
+			},
+		},
+		GlobalSecondaryIndexes: []types.GlobalSecondaryIndex{
+			{
+				IndexName: aws.String("StatusIndex"),
+				KeySchema: []types.KeySchemaElement{
+					{
+						AttributeName: aws.String("BlobStatus"),
+						KeyType:       types.KeyTypeHash,
+					},
+					{
+						AttributeName: aws.String("RequestedAt"),
+						KeyType:       types.KeyTypeRange, // Using RequestedAt as sort key
+					},
+				},
+				Projection: &types.Projection{
+					ProjectionType: types.ProjectionTypeAll, // ProjectionTypeAll means all attributes are projected into the index
+				},
+				ProvisionedThroughput: &types.ProvisionedThroughput{
+					ReadCapacityUnits:  aws.Int64(10),
+					WriteCapacityUnits: aws.Int64(10),
+				},
+			},
+		},
 		TableName: aws.String(tableName),
 		ProvisionedThroughput: &types.ProvisionedThroughput{
 			ReadCapacityUnits:  aws.Int64(10),
@@ -234,3 +270,148 @@ func TestBatchOperations(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Len(t, fetchedItem, 0)
 }
+
+func TestQueryIndex(t *testing.T) {
+	tableName := "ProcessingQueryIndex"
+	createTable(t, tableName)
+	indexName := "StatusIndex"
+
+	ctx := context.Background()
+	numItems := 30
+	items := make([]commondynamodb.Item, numItems)
+	for i := 0; i < numItems; i += 1 {
+		items[i] = commondynamodb.Item{
+			"MetadataKey": &types.AttributeValueMemberS{Value: fmt.Sprintf("key%d", i)},
+			"BlobKey":     &types.AttributeValueMemberS{Value: fmt.Sprintf("blob%d", i)},
+			"BlobSize":    &types.AttributeValueMemberN{Value: "123"},
+			"BlobStatus":  &types.AttributeValueMemberN{Value: "0"},
+			"RequestedAt": &types.AttributeValueMemberN{Value: strconv.FormatInt(time.Now().Unix(), 10)},
+		}
+	}
+	unprocessed, err := dynamoClient.PutItems(ctx, tableName, items)
+	assert.NoError(t, err)
+	assert.Len(t, unprocessed, 0)
+
+	queryResult, err := dynamoClient.QueryIndex(ctx, tableName, indexName, "BlobStatus = :status", commondynamodb.ExpresseionValues{
+		":status": &types.AttributeValueMemberN{
+			Value: "0",
+		}})
+	assert.NoError(t, err)
+	assert.Equal(t, len(queryResult), 30)
+}
+
+func TestQueryIndexPagination(t *testing.T) {
+	tableName := "ProcessingWithPagination"
+	createTable(t, tableName)
+	indexName := "StatusIndex"
+
+	ctx := context.Background()
+	numItems := 30
+	for i := 0; i < numItems; i += 1 {
+		requestedAt := time.Now().Add(-time.Duration(i) * time.Second).Unix()
+
+		// Create new item
+		item := commondynamodb.Item{
+			"MetadataKey": &types.AttributeValueMemberS{Value: fmt.Sprintf("key%d", i)},
+			"BlobKey":     &types.AttributeValueMemberS{Value: fmt.Sprintf("blob%d", i)},
+			"BlobSize":    &types.AttributeValueMemberN{Value: "123"},
+			"BlobStatus":  &types.AttributeValueMemberN{Value: "0"},
+			"RequestedAt": &types.AttributeValueMemberN{Value: strconv.FormatInt(requestedAt, 10)},
+		}
+		err := dynamoClient.PutItem(ctx, tableName, item)
+		assert.NoError(t, err)
+	}
+
+	queryResult, err := dynamoClient.QueryIndexWithPagination(ctx, tableName, indexName, "BlobStatus = :status", commondynamodb.ExpresseionValues{
+		":status": &types.AttributeValueMemberN{
+			Value: "0",
+		}}, 10, nil)
+	assert.NoError(t, err)
+	assert.Len(t, queryResult.Items, 10)
+	assert.Equal(t, "key29", queryResult.Items[0]["MetadataKey"].(*types.AttributeValueMemberS).Value)
+	assert.NotNil(t, queryResult.LastEvaluatedKey)
+	assert.Equal(t, "key20", queryResult.LastEvaluatedKey["MetadataKey"].(*types.AttributeValueMemberS).Value)
+	assert.Equal(t, "0", queryResult.LastEvaluatedKey["BlobStatus"].(*types.AttributeValueMemberN).Value)
+
+	// Get the next 10 items
+	queryResult, err = dynamoClient.QueryIndexWithPagination(ctx, tableName, indexName, "BlobStatus = :status", commondynamodb.ExpresseionValues{
+		":status": &types.AttributeValueMemberN{
+			Value: "0",
+		}}, 10, queryResult.LastEvaluatedKey)
+	assert.NoError(t, err)
+	assert.Len(t, queryResult.Items, 10)
+	assert.Equal(t, "key10", queryResult.LastEvaluatedKey["MetadataKey"].(*types.AttributeValueMemberS).Value)
+
+	// Get the last 10 items
+	queryResult, err = dynamoClient.QueryIndexWithPagination(ctx, tableName, indexName, "BlobStatus = :status", commondynamodb.ExpresseionValues{
+		":status": &types.AttributeValueMemberN{
+			Value: "0",
+		}}, 10, queryResult.LastEvaluatedKey)
+	assert.NoError(t, err)
+	assert.Len(t, queryResult.Items, 10)
+	assert.Equal(t, "key0", queryResult.LastEvaluatedKey["MetadataKey"].(*types.AttributeValueMemberS).Value)
+
+	// Empty result Since all items are processed
+	queryResult, err = dynamoClient.QueryIndexWithPagination(ctx, tableName, indexName, "BlobStatus = :status", commondynamodb.ExpresseionValues{
+		":status": &types.AttributeValueMemberN{
+			Value: "0",
+		}}, 10, queryResult.LastEvaluatedKey)
+	assert.NoError(t, err)
+	assert.Len(t, queryResult.Items, 0)
+	assert.Nil(t, queryResult.LastEvaluatedKey)
+}
+
+func TestQueryIndexWithPaginationForBatch(t *testing.T) {
+	tableName := "ProcessingWithPaginationForBatch"
+	createTable(t, tableName)
+	indexName := "StatusIndex"
+
+	ctx := context.Background()
+	numItems := 30
+	items := make([]commondynamodb.Item, numItems)
+	for i := 0; i < numItems; i += 1 {
+		items[i] = commondynamodb.Item{
+			"MetadataKey": &types.AttributeValueMemberS{Value: fmt.Sprintf("key%d", i)},
+			"BlobKey":     &types.AttributeValueMemberS{Value: fmt.Sprintf("blob%d", i)},
+			"BlobSize":    &types.AttributeValueMemberN{Value: "123"},
+			"BlobStatus":  &types.AttributeValueMemberN{Value: "0"},
+			"RequestedAt": &types.AttributeValueMemberN{Value: strconv.FormatInt(time.Now().Unix(), 10)},
+		}
+	}
+	unprocessed, err := dynamoClient.PutItems(ctx, tableName, items)
+	assert.NoError(t, err)
+	assert.Len(t, unprocessed, 0)
+
+	// Get First 10 items
+	queryResult, err := dynamoClient.QueryIndexWithPagination(ctx, tableName, indexName, "BlobStatus = :status", commondynamodb.ExpresseionValues{
+		":status": &types.AttributeValueMemberN{
+			Value: "0",
+		}}, 10, nil)
+	assert.NoError(t, err)
+	assert.Len(t, queryResult.Items, 10)
+
+	// Get the next 10 items
+	queryResult, err = dynamoClient.QueryIndexWithPagination(ctx, tableName, indexName, "BlobStatus = :status", commondynamodb.ExpresseionValues{
+		":status": &types.AttributeValueMemberN{
+			Value: "0",
+		}}, 10, queryResult.LastEvaluatedKey)
+	assert.NoError(t, err)
+	assert.Len(t, queryResult.Items, 10)
+
+	// Get the last 10 items
+	queryResult, err = dynamoClient.QueryIndexWithPagination(ctx, tableName, indexName, "BlobStatus = :status", commondynamodb.ExpresseionValues{
+		":status": &types.AttributeValueMemberN{
+			Value: "0",
+		}}, 10, queryResult.LastEvaluatedKey)
+	assert.NoError(t, err)
+	assert.Len(t, queryResult.Items, 10)
+
+	// Empty result Since all items are processed
+	queryResult, err = dynamoClient.QueryIndexWithPagination(ctx, tableName, indexName, "BlobStatus = :status", commondynamodb.ExpresseionValues{
+		":status": &types.AttributeValueMemberN{
+			Value: "0",
+		}}, 10, queryResult.LastEvaluatedKey)
+	assert.NoError(t, err)
+	assert.Len(t, queryResult.Items, 0)
+	assert.Nil(t, queryResult.LastEvaluatedKey)
+}

disperser/batcher/batcher.go

@@ -54,7 +54,8 @@ type Config struct {
 	BatchSizeMBLimit     uint
 	MaxNumRetriesPerBlob uint
 
-	TargetNumChunks uint
+	TargetNumChunks          uint
+	MaxBlobsToFetchFromStore int
 }
 
 type Batcher struct {
@@ -97,10 +98,11 @@ func NewBatcher(
 		uint64(config.BatchSizeMBLimit)*1024*1024, // convert to bytes
 	)
 	streamerConfig := StreamerConfig{
-		SRSOrder:               config.SRSOrder,
-		EncodingRequestTimeout: config.PullInterval,
-		EncodingQueueLimit:     config.EncodingRequestQueueSize,
-		TargetNumChunks:        config.TargetNumChunks,
+		SRSOrder:                 config.SRSOrder,
+		EncodingRequestTimeout:   config.PullInterval,
+		EncodingQueueLimit:       config.EncodingRequestQueueSize,
+		TargetNumChunks:          config.TargetNumChunks,
+		MaxBlobsToFetchFromStore: config.MaxBlobsToFetchFromStore,
 	}
 	encodingWorkerPool := workerpool.New(config.NumConnections)
 	encodingStreamer, err := NewEncodingStreamer(streamerConfig, queue, chainState, encoderClient, assignmentCoordinator, batchTrigger, encodingWorkerPool, metrics.EncodingStreamerMetrics, logger)

disperser/batcher/encoding_streamer.go

@@ -41,6 +41,9 @@ type StreamerConfig struct {
 
 	// TargetNumChunks is the target number of chunks per encoded blob
 	TargetNumChunks uint
+
+	// Maximum number of Blobs to fetch from store
+	MaxBlobsToFetchFromStore int
 }
 
 type EncodingStreamer struct {
@@ -62,6 +65,9 @@ type EncodingStreamer struct {
 
 	metrics *EncodingStreamerMetrics
 	logger  common.Logger
+
+	// Used to keep track of the last evaluated key for fetching metadatas
+	exclusiveStartKey *disperser.ExclusiveBlobStoreStartKey
 }
 
 type batch struct {
@@ -107,6 +113,7 @@ func NewEncodingStreamer(
 		encodingCtxCancelFuncs: make([]context.CancelFunc, 0),
 		metrics:                metrics,
 		logger:                 logger,
+		exclusiveStartKey:      nil,
 	}, nil
 }
 
@@ -175,7 +182,12 @@ func (e *EncodingStreamer) dedupRequests(metadatas []*disperser.BlobMetadata, re
 func (e *EncodingStreamer) RequestEncoding(ctx context.Context, encoderChan chan EncodingResultOrStatus) error {
 	stageTimer := time.Now()
 	// pull new blobs and send to encoder
-	metadatas, err := e.blobStore.GetBlobMetadataByStatus(ctx, disperser.Processing)
+	e.mu.Lock()
+	// TODO: Get Limit from Config
+	metadatas, newExclusiveStartKey, err := e.blobStore.GetBlobMetadataByStatusWithPagination(ctx, disperser.Processing, int32(e.StreamerConfig.MaxBlobsToFetchFromStore), e.exclusiveStartKey)
+	e.exclusiveStartKey = newExclusiveStartKey
+	e.mu.Lock()
+
 	if err != nil {
 		return fmt.Errorf("error getting blob metadatas: %w", err)
 	}

disperser/batcher/encoding_streamer_test.go

@@ -21,9 +21,10 @@ import (
 
 var (
 	streamerConfig = batcher.StreamerConfig{
-		SRSOrder:               300000,
-		EncodingRequestTimeout: 5 * time.Second,
-		EncodingQueueLimit:     100,
+		SRSOrder:                 300000,
+		EncodingRequestTimeout:   5 * time.Second,
+		EncodingQueueLimit:       100,
+		MaxBlobsToFetchFromStore: 10,
 	}
 )
 
@@ -294,9 +295,10 @@ func TestEncodingFailure(t *testing.T) {
 	sizeNotifier := batcher.NewEncodedSizeNotifier(make(chan struct{}, 1), 1e12)
 	workerpool := workerpool.New(5)
 	streamerConfig := batcher.StreamerConfig{
-		SRSOrder:               300000,
-		EncodingRequestTimeout: 5 * time.Second,
-		EncodingQueueLimit:     100,
+		SRSOrder:                 300000,
+		EncodingRequestTimeout:   5 * time.Second,
+		EncodingQueueLimit:       100,
+		MaxBlobsToFetchFromStore: 10,
 	}
 	metrics := batcher.NewMetrics("9100", logger)
 	encodingStreamer, err := batcher.NewEncodingStreamer(streamerConfig, blobStore, cst, encoderClient, asgn, sizeNotifier, workerpool, metrics.EncodingStreamerMetrics, logger)

disperser/cmd/batcher/config.go

@@ -51,6 +51,7 @@ func NewConfig(ctx *cli.Context) Config {
 			SRSOrder:                 ctx.GlobalInt(flags.SRSOrderFlag.Name),
 			MaxNumRetriesPerBlob:     ctx.GlobalUint(flags.MaxNumRetriesPerBlobFlag.Name),
 			TargetNumChunks:          ctx.GlobalUint(flags.TargetNumChunksFlag.Name),
+			MaxBlobsToFetchFromStore: ctx.GlobalInt(flags.MaxBlobsToFetchFromStoreFlag.Name),
 		},
 		TimeoutConfig: batcher.TimeoutConfig{
 			EncodingTimeout:    ctx.GlobalDuration(flags.EncodingTimeoutFlag.Name),

disperser/cmd/batcher/flags/flags.go

@@ -164,6 +164,14 @@ var (
 		EnvVar:   common.PrefixEnvVar(envVarPrefix, "TARGET_NUM_CHUNKS"),
 		Value:    0,
 	}
+
+	MaxBlobsToFetchFromStoreFlag = cli.IntFlag{
+		Name:     common.PrefixFlag(FlagPrefix, "max-blobs-to-fetch-from-store"),
+		Usage:    "Limit used to specify how many blobs to fetch from store at time when used with dynamodb pagination",
+		Required: false,
+		EnvVar:   common.PrefixEnvVar(envVarPrefix, "MAX_BLOBS_TO_FETCH_FROM_STORE"),
+		Value:    10,
+	}
 )
 
 var requiredFlags = []cli.Flag{