Fix batch DB garbage collection. (#507)

* Fix batch DB garbage collection.
* Make sure batches are garbage-collected only once

Signed-off-by: Matej Pavlovic <matopavlovic@gmail.com>
Co-authored-by: Alejandro Ranchal-Pedrosa <alejandro@protocol.ai>

pkg/availability/batchdb/fakebatchdb/fakebatchdb.go

@@ -17,13 +17,32 @@ type ModuleConfig struct {
 }
 
 type moduleState struct {
-	batchStore      map[msctypes.BatchID]*batch
-	batchesByRetIdx map[tt.RetentionIndex][]msctypes.BatchID
-	retIdx          tt.RetentionIndex
+
+	// batchStore stores the actual batch data.
+	batchStore map[msctypes.BatchID]*batch
+
+	// batchesByRetIdx is an index that stores a set of batches associated with each particular retention index
+	// that has not yet been garbage-collected.
+	// Upon garbage-collection of a particular retention index, all the batches stored here under that retention index
+	// are deleted, unless they are also associated with a higher retention index.
+	// Using a set of batches rather than a list prevents duplicate entries.
+	batchesByRetIdx map[tt.RetentionIndex]map[msctypes.BatchID]struct{}
+
+	// retIdx is the lowest retention index that has not yet been garbage-collected.
+	retIdx tt.RetentionIndex
 }
 
 type batch struct {
+
+	// Transactions in the batch.
 	txs []*trantorpbtypes.Transaction
+
+	// The maximal retention index with which the batch was stored.
+	// Note that the same batch can be stored in the batchDB (under the same batch ID)
+	// multiple times if, e.g., multiple nodes propose the same transactions in different epoch in Trantor.
+	// Thus, when garbage-collecting, we must not delete a batch that still needs to be retained
+	// for a higher retention index.
+	maxRetIdx tt.RetentionIndex
 }
 
 // NewModule returns a new module for a fake batch database.
@@ -33,7 +52,7 @@ func NewModule(mc ModuleConfig) modules.Module {
 
 	state := moduleState{
 		batchStore:      make(map[msctypes.BatchID]*batch),
-		batchesByRetIdx: make(map[tt.RetentionIndex][]msctypes.BatchID),
+		batchesByRetIdx: make(map[tt.RetentionIndex]map[msctypes.BatchID]struct{}),
 		retIdx:          0,
 	}
 
@@ -47,9 +66,18 @@ func NewModule(mc ModuleConfig) modules.Module {
 
 		// Only save the batch if its retention index has not yet been garbage-collected.
 		if retIdx >= state.retIdx {
-			b := batch{txs}
-			state.batchStore[batchID] = &b
-			state.batchesByRetIdx[retIdx] = append(state.batchesByRetIdx[retIdx], batchID)
+			// Check if we already have the batch.
+			b, ok := state.batchStore[batchID]
+			if !ok || b.maxRetIdx < retIdx {
+				// If we do not, or if the stored batch's retention index is lower,
+				// store the received batch with the up-to-date retention index
+				state.batchStore[batchID] = &batch{txs, retIdx}
+			}
+
+			if _, ok := state.batchesByRetIdx[retIdx]; !ok {
+				state.batchesByRetIdx[retIdx] = make(map[msctypes.BatchID]struct{})
+			}
+			state.batchesByRetIdx[retIdx][batchID] = struct{}{}
 		}
 
 		// Note that we emit a BatchStored event even if the batch's retention index was too low
@@ -77,8 +105,10 @@ func NewModule(mc ModuleConfig) modules.Module {
 
 	batchdbpbdsl.UponGarbageCollect(m, func(retentionIndex tt.RetentionIndex) error {
 		for ; state.retIdx < retentionIndex; state.retIdx++ {
-			for _, batchID := range state.batchesByRetIdx[state.retIdx] {
-				delete(state.batchStore, batchID)
+			for batchID := range state.batchesByRetIdx[state.retIdx] {
+				if state.batchStore[batchID].maxRetIdx <= state.retIdx {
+					delete(state.batchStore, batchID)
+				}
 			}
 			delete(state.batchesByRetIdx, state.retIdx)
 		}