feat: add retry to get batch of s3 (#958)

Co-authored-by: Tatu <[email protected]>
Co-authored-by: Mariano Nicolini <[email protected]>
Co-authored-by: Urix <[email protected]>

Author: 4 people

Makefile

@@ -244,7 +244,7 @@ batcher_send_sp1_burst:
 		--proving_system SP1 \
 		--proof ../../scripts/test_files/sp1/sp1_fibonacci.proof \
 		--vm_program ../../scripts/test_files/sp1/sp1_fibonacci.elf \
-		--repetitions 15 \
+		--repetitions $(BURST_SIZE) \
 		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \
 		--rpc_url $(RPC_URL) \
 		--payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS)
@@ -271,7 +271,7 @@ batcher_send_risc0_burst:
 		--proof ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.proof \
         --vm_program ../../scripts/test_files/risc_zero/fibonacci_proof_generator/fibonacci_id.bin \
         --public_input ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.pub \
-        --repetitions 15 \
+        --repetitions $(BURST_SIZE) \
 		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \
 		--rpc_url $(RPC_URL) \
 		--payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS)

operator/pkg/operator.go

@@ -55,7 +55,9 @@ type Operator struct {
 }
 
 const (
-	BatchDownloadTimeout = 1 * time.Minute
+	BatchDownloadTimeout    = 1 * time.Minute
+	BatchDownloadMaxRetries = 3
+	BatchDownloadRetryDelay = 5 * time.Second
 )
 
 func NewOperatorFromConfig(configuration config.OperatorConfig) (*Operator, error) {
@@ -183,7 +185,7 @@ func (o *Operator) handleNewBatchLog(newBatchLog *servicemanager.ContractAligned
 		hex.EncodeToString(signedTaskResponse.BatchMerkleRoot[:]),
 		hex.EncodeToString(signedTaskResponse.SenderAddress[:]),
 	)
-	
+
 	o.aggRpcClient.SendSignedTaskResponseToAggregator(&signedTaskResponse)
 }
 
@@ -199,7 +201,7 @@ func (o *Operator) ProcessNewBatchLog(newBatchLog *servicemanager.ContractAligne
 	ctx, cancel := context.WithTimeout(context.Background(), BatchDownloadTimeout)
 	defer cancel()
 
-	verificationDataBatch, err := o.getBatchFromS3(ctx, newBatchLog.BatchDataPointer, newBatchLog.BatchMerkleRoot)
+	verificationDataBatch, err := o.getBatchFromDataService(ctx, newBatchLog.BatchDataPointer, newBatchLog.BatchMerkleRoot, BatchDownloadMaxRetries, BatchDownloadRetryDelay)
 	if err != nil {
 		o.Logger.Errorf("Could not get proofs from S3 bucket: %v", err)
 		return err

operator/pkg/s3.go

@@ -5,25 +5,58 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"time"
 
 	"github.com/ugorji/go/codec"
 
 	"github.com/yetanotherco/aligned_layer/operator/merkle_tree"
 )
 
-func (o *Operator) getBatchFromS3(ctx context.Context, batchURL string, expectedMerkleRoot [32]byte) ([]VerificationData, error) {
-	o.Logger.Infof("Getting batch from S3..., batchURL: %s", batchURL)
+func (o *Operator) getBatchFromDataService(ctx context.Context, batchURL string, expectedMerkleRoot [32]byte, maxRetries int, retryDelay time.Duration) ([]VerificationData, error) {
+	o.Logger.Infof("Getting batch from data service, batchURL: %s", batchURL)
+
+	var resp *http.Response
+	var err error
+	var req *http.Request
+
+	for attempt := 0; attempt < maxRetries; attempt++ {
+		if attempt > 0 {
+			o.Logger.Infof("Waiting for %s before retrying data fetch (attempt %d of %d)", retryDelay, attempt+1, maxRetries)
+			select {
+			case <-time.After(retryDelay):
+				// Wait before retrying
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			}
+			retryDelay *= 2 // Exponential backoff. Ex: 5s, 10s, 20s
+		}
 
-	req, err := http.NewRequestWithContext(ctx, "GET", batchURL, nil)
+		req, err = http.NewRequestWithContext(ctx, "GET", batchURL, nil)
+		if err != nil {
+			return nil, err
+		}
 
-	if err != nil {
-		return nil, err
+		resp, err = http.DefaultClient.Do(req)
+		if err == nil && resp.StatusCode == http.StatusOK {
+			break // Successful request, exit retry loop
+		}
+
+		if resp != nil {
+			err := resp.Body.Close()
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		o.Logger.Warnf("Error fetching batch from data service - (attempt %d): %v", attempt+1, err)
 	}
 
-	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return nil, err
 	}
+
+	// At this point, the HTTP request was successfull.
+
 	defer func(Body io.ReadCloser) {
 		err := Body.Close()
 		if err != nil {
@@ -33,7 +66,7 @@ func (o *Operator) getBatchFromS3(ctx context.Context, batchURL string, expected
 
 	// Check if the response is OK
 	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("error getting Proof Head from S3: %s", resp.Status)
+		return nil, fmt.Errorf("error getting batch from data service: %s", resp.Status)
 	}
 
 	contentLength := resp.ContentLength
@@ -58,8 +91,8 @@ func (o *Operator) getBatchFromS3(ctx context.Context, batchURL string, expected
 
 	// Checks if downloaded merkle root is the same as the expected one
 	o.Logger.Infof("Verifying batch merkle tree...")
-	merkle_root_check := merkle_tree.VerifyMerkleTreeBatch(batchBytes, uint(len(batchBytes)), expectedMerkleRoot)
-	if !merkle_root_check {
+	merkleRootCheck := merkle_tree.VerifyMerkleTreeBatch(batchBytes, uint(len(batchBytes)), expectedMerkleRoot)
+	if !merkleRootCheck {
 		return nil, fmt.Errorf("merkle root check failed")
 	}
 	o.Logger.Infof("Batch merkle tree verified")