fix: data race related to last commit info; flush pruning heights as soon as changed

pruning/manager.go

@@ -14,6 +14,7 @@ import (
 
 type Manager struct {
 	logger               log.Logger
+	db dbm.DB
 	opts                 *types.PruningOptions
 	snapshotInterval     uint64
 	pruneHeights         []int64
@@ -24,11 +25,14 @@ type Manager struct {
 const (
 	pruneHeightsKey         = "s/pruneheights"
 	pruneSnapshotHeightsKey = "s/pruneSnheights"
+
+	uint64Size = 8
 )
 
-func NewManager(logger log.Logger) *Manager {
+func NewManager(logger log.Logger, db dbm.DB) *Manager {
 	return &Manager{
 		logger:       logger,
+		db: db,
 		opts:         types.NewPruningOptions(types.PruningNothing),
 		pruneHeights: []int64{},
 		// These are the heights that are multiples of snapshotInterval and kept for state sync snapshots.
@@ -65,10 +69,22 @@ func (m *Manager) HandleHeight(previousHeight int64) int64 {
 		return 0
 	}
 
+	// Flag indicating whether should flush to disk.
+	// It is set to true when an update to one of
+	// pruneHeights or pruneSnapshotHeights is made
+	shouldFlush := false
+
 	defer func() {
 		// handle persisted snapshot heights
 		m.mx.Lock()
 		defer m.mx.Unlock()
+		defer func() {
+			if shouldFlush {
+				// Must be the unlocked implementation since we are under mutex
+				m.flushAllPruningHeightsUnlocked()
+			}
+		}()
+
 		var next *list.Element
 		for e := m.pruneSnapshotHeights.Front(); e != nil; e = next {
 			snHeight := e.Value.(int64)
@@ -78,6 +94,8 @@ func (m *Manager) HandleHeight(previousHeight int64) int64 {
 				// We must get next before removing to be able to continue iterating.
 				next = e.Next()
 				m.pruneSnapshotHeights.Remove(e)
+
+				shouldFlush = true
 			} else {
 				next = e.Next()
 			}
@@ -93,6 +111,7 @@ func (m *Manager) HandleHeight(previousHeight int64) int64 {
 		// a 'snapshot' height.
 		if m.snapshotInterval == 0 || pruneHeight%int64(m.snapshotInterval) != 0 {
 			m.pruneHeights = append(m.pruneHeights, pruneHeight)
+			shouldFlush = true
 			return pruneHeight
 		}
 	}
@@ -107,6 +126,7 @@ func (m *Manager) HandleHeightSnapshot(height int64) {
 	defer m.mx.Unlock()
 	m.logger.Debug("HandleHeightSnapshot", "height", height) // TODO: change log level to Debug
 	m.pruneSnapshotHeights.PushBack(height)
+	// m.flushPruningSnapshotHeightsUnlocked()
 }
 
 // SetSnapshotInterval sets the interval at which the snapshots are taken.
@@ -119,15 +139,6 @@ func (m *Manager) ShouldPruneAtHeight(height int64) bool {
 	return m.opts.GetPruningStrategy() != types.PruningNothing && m.opts.Interval > 0 && height%int64(m.opts.Interval) == 0
 }
 
-// FlushPruningHeights flushes the pruning heights to the database for crash recovery.
-func (m *Manager) FlushPruningHeights(batch dbm.Batch) {
-	if m.opts.GetPruningStrategy() == types.PruningNothing {
-		return
-	}
-	m.flushPruningHeights(batch)
-	m.flushPruningSnapshotHeights(batch)
-}
-
 // LoadPruningHeights loads the pruning heights from the database as a crash recovery.
 func (m *Manager) LoadPruningHeights(db dbm.DB) error {
 	if m.opts.GetPruningStrategy() == types.PruningNothing {
@@ -192,25 +203,80 @@ func (m *Manager) loadPruningSnapshotHeights(db dbm.DB) error {
 	return nil
 }
 
-func (m *Manager) flushPruningHeights(batch dbm.Batch) {
-	bz := make([]byte, 0)
+// FlushAllPruningHeights flushes the pruning heights to the database for crash recovery.
+// "All" refers to regular pruning heights and snapshot heights, if any.
+func (m *Manager) FlushAllPruningHeights() {
+	if m.opts.GetPruningStrategy() == types.PruningNothing {
+		return
+	}
+	batch := m.db.NewBatch()
+	defer batch.Close()
+	m.flushPruningHeightsBatch(batch)
+	m.flushPruningSnapshotHeightsBatch(batch)
+
+	if err := batch.Write(); err != nil {
+		panic(fmt.Errorf("error on batch write %w", err))
+	}
+}
+
+// flushAllPruningHeightsUnlocked flushes the pruning heights to the database for crash recovery.
+// "All" refers to regular pruning heights and snapshot heights, if any.
+// It serves the same function as exported FlushPruningHeights. However, it assummes that
+// mutex was acquired prior to calling this method.
+func (m *Manager) flushAllPruningHeightsUnlocked() {
+	if m.opts.GetPruningStrategy() == types.PruningNothing {
+		return
+	}
+	batch := m.db.NewBatch()
+	defer batch.Close()
+	m.flushPruningHeightsBatch(batch)
+	m.flushPruningSnapshotHeightsUnlockedBatch(batch)
+
+	if err := batch.Write(); err != nil {
+		panic(fmt.Errorf("error on batch write %w", err))
+	}
+}
+
+func (m *Manager) flushPruningHeightsBatch(batch dbm.Batch) {
+	bz := make([]byte, 0, uint64Size * len(m.pruneHeights))
 	for _, ph := range m.pruneHeights {
-		buf := make([]byte, 8)
+		buf := make([]byte, uint64Size)
 		binary.BigEndian.PutUint64(buf, uint64(ph))
 		bz = append(bz, buf...)
 	}
 
-	batch.Set([]byte(pruneHeightsKey), bz)
+	if err := batch.Set([]byte(pruneHeightsKey), bz); err != nil {
+		panic(err)
+	}
 }
 
-func (m *Manager) flushPruningSnapshotHeights(batch dbm.Batch) {
+func (m *Manager) flushPruningSnapshotHeightsBatch(batch dbm.Batch) {
 	m.mx.Lock()
 	defer m.mx.Unlock()
-	bz := make([]byte, 0)
-	for e := m.pruneSnapshotHeights.Front(); e != nil; e = e.Next() {
-		buf := make([]byte, 8)
+	m.flushPruningSnapshotHeightsUnlockedBatch(batch)
+}
+
+func (m *Manager) flushPruningSnapshotHeightsUnlockedBatch(batch dbm.Batch) {
+	bz := convertInt64ListToBytes(m.pruneSnapshotHeights)
+	if err := batch.Set([]byte(pruneSnapshotHeightsKey), bz); err != nil {
+		panic(err)
+	}
+}
+
+func (m *Manager) flushPruningSnapshotHeightsUnlocked() {
+	bz := convertInt64ListToBytes(m.pruneSnapshotHeights)
+	if err := m.db.Set([]byte(pruneSnapshotHeightsKey), bz); err != nil {
+		panic(err)
+	}
+}
+
+// INVARIANT: list contains ints
+func convertInt64ListToBytes(list *list.List) []byte {
+	bz := make([]byte, 0, uint64Size * list.Len())
+	for e := list.Front(); e != nil; e = e.Next() {
+		buf := make([]byte, uint64Size)
 		binary.BigEndian.PutUint64(buf, uint64(e.Value.(int64)))
 		bz = append(bz, buf...)
 	}
-	batch.Set([]byte(pruneSnapshotHeightsKey), bz)
+	return bz
 }