fix(audit): Remove Task data from aggregator after a response has been responded or expires.

aggregator/internal/pkg/aggregator.go

@@ -12,7 +12,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/yetanotherco/aligned_layer/metrics"
 
-	"github.com/Layr-Labs/eigensdk-go/chainio/clients"
 	sdkclients "github.com/Layr-Labs/eigensdk-go/chainio/clients"
 	"github.com/Layr-Labs/eigensdk-go/logging"
 	"github.com/Layr-Labs/eigensdk-go/services/avsregistry"
@@ -61,12 +60,16 @@ type Aggregator struct {
 	// and can start from zero
 	batchesIdxByIdentifierHash map[[32]byte]uint32
 
-	// Stores the taskCreatedBlock for each batch bt batch index
+	// Stores the taskCreatedBlock for each batch by batch index
 	batchCreatedBlockByIdx map[uint32]uint64
 
 	// Stores the TaskResponse for each batch by batchIdentifierHash
 	batchDataByIdentifierHash map[[32]byte]BatchData
 
+	// Stores if a batch has been finalized, either by response or failure to respond
+	batchIsFinalizedByIdx map[uint32]struct{} // Id in the key list means it is finalized, using empty struct to save memory
+	batchIsFinalizedChan chan uint32
+
 	// This task index is to communicate with the local BLS
 	// Service.
 	// Note: In case of a reboot it can start from 0 again
@@ -106,6 +109,8 @@ func NewAggregator(aggregatorConfig config.AggregatorConfig) (*Aggregator, error
 	batchesIdxByIdentifierHash := make(map[[32]byte]uint32)
 	batchDataByIdentifierHash := make(map[[32]byte]BatchData)
 	batchCreatedBlockByIdx := make(map[uint32]uint64)
+	batchIsFinalizedByIdx := make(map[uint32]struct{})
+	batchIsFinalizedChan := make(chan uint32)
 
 	chainioConfig := sdkclients.BuildAllConfig{
 		EthHttpUrl:                 aggregatorConfig.BaseConfig.EthRpcUrl,
@@ -119,7 +124,7 @@ func NewAggregator(aggregatorConfig config.AggregatorConfig) (*Aggregator, error
 	aggregatorPrivateKey := aggregatorConfig.EcdsaConfig.PrivateKey
 
 	logger := aggregatorConfig.BaseConfig.Logger
-	clients, err := clients.BuildAll(chainioConfig, aggregatorPrivateKey, logger)
+	clients, err := sdkclients.BuildAll(chainioConfig, aggregatorPrivateKey, logger)
 	if err != nil {
 		logger.Errorf("Cannot create sdk clients", "err", err)
 		return nil, err
@@ -161,6 +166,8 @@ func NewAggregator(aggregatorConfig config.AggregatorConfig) (*Aggregator, error
 		batchesIdxByIdentifierHash: batchesIdxByIdentifierHash,
 		batchDataByIdentifierHash:  batchDataByIdentifierHash,
 		batchCreatedBlockByIdx:     batchCreatedBlockByIdx,
+		batchIsFinalizedByIdx:      batchIsFinalizedByIdx,
+		batchIsFinalizedChan:       batchIsFinalizedChan,
 		nextBatchIndex:             nextBatchIndex,
 		taskMutex:                  &sync.Mutex{},
 		walletMutex:                &sync.Mutex{},
@@ -171,6 +178,8 @@ func NewAggregator(aggregatorConfig config.AggregatorConfig) (*Aggregator, error
 		metrics:               aggregatorMetrics,
 	}
 
+	go aggregator.clearTasksFromMaps(garbageCollectorPeriod)
+
 	return &aggregator, nil
 }
 
@@ -202,17 +211,29 @@ func (agg *Aggregator) Start(ctx context.Context) error {
 				"taskIndex", blsAggServiceResp.TaskIndex)
 
 			go agg.handleBlsAggServiceResponse(blsAggServiceResp)
+		case taskIdx := <-agg.batchIsFinalizedChan:
+			agg.logger.Info("Batch is finalized", "taskIndex", taskIdx)
+			agg.finalizeBatchIdx(taskIdx)
 		}
 	}
 }
 
 const MaxSentTxRetries = 5
 
+const garbageCollectorPeriod = 60 * time.Second
+
+const BLS_AGG_SERVICE_TIMEOUT = 100 * time.Second
+
 func (agg *Aggregator) handleBlsAggServiceResponse(blsAggServiceResp blsagg.BlsAggregationServiceResponse) {
 	if blsAggServiceResp.Err != nil {
 		agg.taskMutex.Lock()
 		batchIdentifierHash := agg.batchesIdentifierHashByIdx[blsAggServiceResp.TaskIndex]
 		agg.logger.Error("BlsAggregationServiceResponse contains an error", "err", blsAggServiceResp.Err, "batchIdentifierHash", hex.EncodeToString(batchIdentifierHash[:]))
+		// Task errored, mark as finalized
+		// TODO: 
+		// Actually, if I remove this task from the maps, Agg will go and fetch it again from chain
+		// agg.batchIsFinalizedChan <- blsAggServiceResp.TaskIndex
+
 		agg.taskMutex.Unlock()
 		return
 	}
@@ -265,6 +286,11 @@ func (agg *Aggregator) handleBlsAggServiceResponse(blsAggServiceResp blsagg.BlsA
 			agg.logger.Info("Aggregator successfully responded to task",
 				"taskIndex", blsAggServiceResp.TaskIndex,
 				"batchIdentifierHash", "0x"+hex.EncodeToString(batchIdentifierHash[:]))
+
+			// Mark the batch as finalized
+			agg.batchIsFinalizedChan <- blsAggServiceResp.TaskIndex
+
+
 			return
 		}
 
@@ -278,6 +304,9 @@ func (agg *Aggregator) handleBlsAggServiceResponse(blsAggServiceResp blsagg.BlsA
 		"merkleRoot", "0x"+hex.EncodeToString(batchData.BatchMerkleRoot[:]),
 		"senderAddress", "0x"+hex.EncodeToString(batchData.SenderAddress[:]),
 		"batchIdentifierHash", "0x"+hex.EncodeToString(batchIdentifierHash[:]))
+
+	// Aggregator failed to respond to the task, mark the batch as finalized
+	agg.batchIsFinalizedChan <- blsAggServiceResp.TaskIndex
 }
 
 // / Sends response to contract and waits for transaction receipt
@@ -347,12 +376,17 @@ func (agg *Aggregator) AddNewTask(batchMerkleRoot [32]byte, senderAddress [20]by
 		BatchMerkleRoot: batchMerkleRoot,
 		SenderAddress:   senderAddress,
 	}
+	agg.logger.Info(
+		"Task Info added in aggregator:",
+		"Task", batchIndex,
+		"batchIdentifierHash", batchIdentifierHash,
+	)
 	agg.nextBatchIndex += 1
 
 	quorumNums := eigentypes.QuorumNums{eigentypes.QuorumNum(QUORUM_NUMBER)}
 	quorumThresholdPercentages := eigentypes.QuorumThresholdPercentages{eigentypes.QuorumThresholdPercentage(QUORUM_THRESHOLD)}
 
-	err := agg.blsAggregationService.InitializeNewTask(batchIndex, taskCreatedBlock, quorumNums, quorumThresholdPercentages, 100*time.Second)
+	err := agg.blsAggregationService.InitializeNewTask(batchIndex, taskCreatedBlock, quorumNums, quorumThresholdPercentages, BLS_AGG_SERVICE_TIMEOUT)
 	// FIXME(marian): When this errors, should we retry initializing new task? Logging fatal for now.
 	if err != nil {
 		agg.logger.Fatalf("BLS aggregation service error when initializing new task: %s", err)
@@ -362,3 +396,47 @@ func (agg *Aggregator) AddNewTask(batchMerkleRoot [32]byte, senderAddress [20]by
 	agg.AggregatorConfig.BaseConfig.Logger.Info("- Unlocked Resources: Adding new task")
 	agg.logger.Info("New task added", "batchIndex", batchIndex, "batchIdentifierHash", "0x"+hex.EncodeToString(batchIdentifierHash[:]))
 }
+
+// long-lived gorouting that periodically checks and removes finished Tasks from stored Maps
+func (agg *Aggregator) clearTasksFromMaps(period time.Duration) {
+	agg.AggregatorConfig.BaseConfig.Logger.Info("- Removing finalized Task Infos from Maps every %d seconds", period)
+
+	for {
+		time.Sleep(period)
+
+		agg.AggregatorConfig.BaseConfig.Logger.Info("Cleaning finalized tasks from maps")
+
+		// Reading batchIsFinalizedByIdx map without using a lock because worst case scenario is we miss a newly inserted value
+		// in which case, the value will be catched in the next iteration
+
+		for idx := range agg.batchIsFinalizedByIdx {
+			agg.AggregatorConfig.BaseConfig.Logger.Info("Cleaning up finalized task", "taskIndex", idx)
+
+			// Critical section inside anonymous function to ensure defer works
+			func() {
+				agg.taskMutex.Lock()
+				agg.AggregatorConfig.BaseConfig.Logger.Info("- Locked Resources: Removing Task Info from Aggregator")
+
+				defer func() {
+					agg.AggregatorConfig.BaseConfig.Logger.Info("- Unlocked Resources: Removed Task Info from Aggregator")
+					agg.taskMutex.Unlock()
+				}()
+
+				batchIdentifierHash := agg.batchesIdentifierHashByIdx[idx]
+
+				delete(agg.batchesIdxByIdentifierHash, batchIdentifierHash)
+				delete(agg.batchCreatedBlockByIdx, idx)
+				delete(agg.batchesIdentifierHashByIdx, idx)
+				delete(agg.batchDataByIdentifierHash, batchIdentifierHash)
+
+				delete(agg.batchIsFinalizedByIdx, idx)
+			}()
+		}
+
+	}
+}
+
+// called in the chan so no need for a mutex
+func (agg *Aggregator) finalizeBatchIdx(idx uint32) {
+	agg.batchIsFinalizedByIdx[idx] = struct{}{} //now the key is present, no need to waste memory on a value
+}