fix(audit): Periodically Remove oldTask data from aggregator

aggregator/internal/pkg/aggregator.go

@@ -60,7 +60,7 @@ 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
@@ -178,6 +178,8 @@ func NewAggregator(aggregatorConfig config.AggregatorConfig) (*Aggregator, error
 		telemetry:             aggregatorTelemetry,
 	}
 
+	go aggregator.clearTasksFromMaps(garbageCollectorPeriod, garbageCollectorTasksAge)
+
 	return &aggregator, nil
 }
 
@@ -215,6 +217,11 @@ func (agg *Aggregator) Start(ctx context.Context) error {
 
 const MaxSentTxRetries = 5
 
+const garbageCollectorPeriod = time.Second * 150 //TODO change to time.Day * 1
+const garbageCollectorTasksAge = uint64(10)     //TODO change to 2592000, 1 month of blocks
+
+const BLS_AGG_SERVICE_TIMEOUT = 100 * time.Second
+
 func (agg *Aggregator) handleBlsAggServiceResponse(blsAggServiceResp blsagg.BlsAggregationServiceResponse) {
 	agg.taskMutex.Lock()
 	agg.AggregatorConfig.BaseConfig.Logger.Info("- Locked Resources: Fetching task data")
@@ -275,6 +282,7 @@ func (agg *Aggregator) handleBlsAggServiceResponse(blsAggServiceResp blsagg.BlsA
 			agg.logger.Info("Aggregator successfully responded to task",
 				"taskIndex", blsAggServiceResp.TaskIndex,
 				"batchIdentifierHash", "0x"+hex.EncodeToString(batchIdentifierHash[:]))
+
 			return
 		}
 
@@ -361,12 +369,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)
@@ -377,3 +390,37 @@ 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 old Tasks from stored Maps
+func (agg *Aggregator) clearTasksFromMaps(period time.Duration, blocksOld uint64) {
+	agg.AggregatorConfig.BaseConfig.Logger.Info("- Removing finalized Task Infos from Maps every %d seconds", period)
+	lastIdxDeleted := uint32(0)
+
+	for {
+		time.Sleep(period)
+
+		agg.AggregatorConfig.BaseConfig.Logger.Info("Cleaning finalized tasks from maps")
+		oldTaskIdHash, err := agg.avsReader.GetOldTaskHash(blocksOld)
+		if err != nil {
+			agg.logger.Error("Error getting old task hash, skipping this garbage collect", "err", err)
+			continue // Retry in the next iteration
+		}
+
+		oldTaskIdx := agg.batchesIdxByIdentifierHash[*oldTaskIdHash]
+		agg.logger.Info("Old task found", "taskIndex", oldTaskIdx)
+		for i := lastIdxDeleted+1; i <= oldTaskIdx; i++ {
+			batchIdentifierHash, exists := agg.batchesIdentifierHashByIdx[i]
+			if exists {
+				agg.logger.Info("Cleaning up finalized task", "taskIndex", i)
+				delete(agg.batchesIdxByIdentifierHash, batchIdentifierHash)
+				delete(agg.batchCreatedBlockByIdx, i)
+				delete(agg.batchesIdentifierHashByIdx, i)
+				delete(agg.batchDataByIdentifierHash, batchIdentifierHash)
+			} else {
+				agg.logger.Warn("Task not found in maps", "taskIndex", i)
+			}
+		}
+		lastIdxDeleted = oldTaskIdx
+		agg.AggregatorConfig.BaseConfig.Logger.Info("Done cleaning finalized tasks from maps")
+	}
+}