fix(pruning): clear in-flight snapshot state when creation fails

CHANGELOG.md

@@ -55,6 +55,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Bug Fixes
 
+* (store) [#26787](https://github.com/cosmos/cosmos-sdk/pull/26787) clear in-flight snapshot state when creation fails.
 * (client/tx) [#26759](https://github.com/cosmos/cosmos-sdk/issues/26759) Populate the multisig bit array in simulation txs so `--gas auto` works for multisig senders.
 * (blockstm) [#26772](https://github.com/cosmos/cosmos-sdk/pull/26772) Panic with a descriptive error when accessing an unregistered store instead of silently using store index zero.
 * (x/genutil) [#26741](https://github.com/cosmos/cosmos-sdk/issues/26741) Preserve vote extension enable height when exporting genesis state.

store/pruning/manager.go

@@ -4,7 +4,6 @@ import (
 	"encoding/binary"
 	"fmt"
 	"slices"
-	"sort"
 	"sync"
 
 	dbm "github.com/cosmos/cosmos-db"
@@ -22,15 +21,13 @@ type Manager struct {
 	logger           log.Logger
 	opts             types.PruningOptions
 	snapshotInterval uint64
-	// Snapshots are taken in a separate goroutine from the regular execution
-	// and can be delivered asynchronously via HandleSnapshotHeight.
-	// Therefore, we sync access to pruneSnapshotHeights, inflightSnapshotHeights and initFromStore with this mutex.
+	// Snapshots are taken in a separate goroutine from regular execution.
 	pruneSnapshotHeightsMx sync.RWMutex
-	// These are the heights that are multiples of snapshotInterval and kept for state sync snapshots.
-	// The heights are added to be pruned when a snapshot is complete.
-	pruneSnapshotHeights    []int64
-	inflightSnapshotHeights []int64
-	initFromStore           bool
+	// completedSnapshotHeight is the highest durably completed snapshot.
+	completedSnapshotHeight int64
+	// inflightSnapshotHeights is memory-only and contains snapshots without a terminal transition.
+	inflightSnapshotHeights map[int64]struct{}
+	loadedFromDisk          bool
 }
 
 // NegativeHeightsError is returned when a negative height is provided to the manager.
@@ -52,10 +49,10 @@ var pruneSnapshotHeightsKey = []byte("s/prunesnapshotheights")
 // by calling SetOptions.
 func NewManager(db dbm.DB, logger log.Logger) *Manager {
 	return &Manager{
-		db:                   db,
-		logger:               logger,
-		opts:                 types.NewPruningOptions(types.PruningNothing),
-		pruneSnapshotHeights: []int64{0}, // init with 0 block height
+		db:                      db,
+		logger:                  logger,
+		opts:                    types.NewPruningOptions(types.PruningNothing),
+		inflightSnapshotHeights: make(map[int64]struct{}),
 	}
 }
 
@@ -69,64 +66,64 @@ func (m *Manager) GetOptions() types.PruningOptions {
 	return m.opts
 }
 
-// AnnounceSnapshotHeight announces a new snapshot height for tracking and pruning.
-func (m *Manager) AnnounceSnapshotHeight(height int64) {
+// StartSnapshot tracks a snapshot while it is being created.
+func (m *Manager) StartSnapshot(height int64) {
 	if m.opts.GetPruningStrategy() == types.PruningNothing || height <= 0 {
 		return
 	}
 	m.pruneSnapshotHeightsMx.Lock()
 	defer m.pruneSnapshotHeightsMx.Unlock()
-	// called in ascending order so no sorting required
-	m.inflightSnapshotHeights = append(m.inflightSnapshotHeights, height)
+	if height <= m.completedSnapshotHeight {
+		return
+	}
+	m.inflightSnapshotHeights[height] = struct{}{}
 }
 
-// HandleSnapshotHeight persists the snapshot height to be pruned at the next appropriate
-// height defined by the pruning strategy. It flushes the update to disk and panics if the flush fails.
-// The input height must be greater than 0, and the pruning strategy must not be set to pruning nothing.
-// If either of these conditions is not met, this function does nothing.
-func (m *Manager) HandleSnapshotHeight(height int64) {
+// AnnounceSnapshotHeight announces a new snapshot height for tracking and pruning.
+func (m *Manager) AnnounceSnapshotHeight(height int64) {
+	m.StartSnapshot(height)
+}
+
+// FailSnapshot removes a failed snapshot height from in-flight tracking.
+func (m *Manager) FailSnapshot(height int64) {
 	if m.opts.GetPruningStrategy() == types.PruningNothing || height <= 0 {
 		return
 	}
-
-	m.logger.Debug("HandleSnapshotHeight", "height", height)
-
 	m.pruneSnapshotHeightsMx.Lock()
 	defer m.pruneSnapshotHeightsMx.Unlock()
+	delete(m.inflightSnapshotHeights, height)
+}
 
-	// remove from the in-flight list
-	if position := slices.Index(m.inflightSnapshotHeights, height); position != -1 {
-		m.inflightSnapshotHeights = append(m.inflightSnapshotHeights[:position], m.inflightSnapshotHeights[position+1:]...)
+// CompleteSnapshot records a successfully completed snapshot height.
+func (m *Manager) CompleteSnapshot(height int64) {
+	if m.opts.GetPruningStrategy() == types.PruningNothing || height <= 0 {
+		return
 	}
 
-	if m.initFromStore {
-		// drop the legacy state as it may belong to a different interval or an outdated snapshot
-		// that is not in sequence with the current one
-		m.pruneSnapshotHeights = m.pruneSnapshotHeights[1:]
-		m.initFromStore = false
-	}
+	m.logger.Debug("CompleteSnapshot", "height", height)
 
-	m.pruneSnapshotHeights = append(m.pruneSnapshotHeights, height)
-	sort.Slice(m.pruneSnapshotHeights, func(i, j int) bool { return m.pruneSnapshotHeights[i] < m.pruneSnapshotHeights[j] })
+	m.pruneSnapshotHeightsMx.Lock()
+	defer m.pruneSnapshotHeightsMx.Unlock()
 
-	// in-flight snapshots may land out of order due to the concurrent nature of the snapshotter.
-	// we need to detect them to prevent pruning their heights while the snapshots are still in progress.
-	k := 1
-	for ; k < len(m.pruneSnapshotHeights); k++ {
-		if m.pruneSnapshotHeights[k] != m.pruneSnapshotHeights[k-1]+int64(m.snapshotInterval) {
-			// gap detected, snapshot is in-flight
-			break
-		}
+	if height <= m.completedSnapshotHeight {
+		delete(m.inflightSnapshotHeights, height)
+		return
 	}
-	// compact the height list for the snapshots in sequence
-	// the last snapshot height is used to allow pruning up to the next interval height
-	m.pruneSnapshotHeights = m.pruneSnapshotHeights[k-1:]
 
 	// flush the max height to store so that they are not lost if a crash happens.
 	// only the max height matters as there are no in-flight snapshots after a restart
-	if err := storePruningSnapshotHeight(m.db, slices.Max(m.pruneSnapshotHeights)); err != nil {
+	if err := storePruningSnapshotHeight(m.db, height); err != nil {
 		panic(err)
 	}
+
+	m.completedSnapshotHeight = height
+	m.loadedFromDisk = false
+	delete(m.inflightSnapshotHeights, height)
+}
+
+// HandleSnapshotHeight is kept for compatibility with legacy callers.
+func (m *Manager) HandleSnapshotHeight(height int64) {
+	m.CompleteSnapshot(height)
 }
 
 // SetSnapshotInterval sets the interval at which the snapshots are taken.
@@ -155,22 +152,18 @@ func (m *Manager) GetPruningHeight(height int64) int64 {
 	m.pruneSnapshotHeightsMx.RLock()
 	defer m.pruneSnapshotHeightsMx.RUnlock()
 
-	if len(m.pruneSnapshotHeights) == 0 { // do not prune before an initial snapshot
-		return 0
-	}
-
-	// highest version based on completed snapshots
-	snHeight := m.pruneSnapshotHeights[0] - 1
-	if !m.initFromStore { // ensure non-legacy data
-		// with no inflight snapshots, we may prune up to the next snap interval -1
-		snHeight += int64(m.snapshotInterval)
+	completedLimit := m.completedSnapshotHeight - 1
+	if !m.loadedFromDisk {
+		completedLimit += int64(m.snapshotInterval)
 	}
 	if len(m.inflightSnapshotHeights) == 0 {
-		return min(snHeight, pruneHeight)
+		return min(completedLimit, pruneHeight)
+	}
+	inFlightHeight := int64(^uint64(0) >> 1)
+	for snapshotHeight := range m.inflightSnapshotHeights {
+		inFlightHeight = min(inFlightHeight, snapshotHeight-1)
 	}
-	// highest version based on started snapshots
-	inFlightHeight := m.inflightSnapshotHeights[0] - 1
-	return min(snHeight, pruneHeight, inFlightHeight)
+	return min(completedLimit, pruneHeight, inFlightHeight)
 }
 
 // LoadSnapshotHeights loads the snapshot heights from the database as a crash recovery.
@@ -191,8 +184,8 @@ func (m *Manager) LoadSnapshotHeights(db dbm.DB) error {
 	m.pruneSnapshotHeightsMx.Lock()
 	defer m.pruneSnapshotHeightsMx.Unlock()
 	// restore max only as there are no in-flight snapshots after a restart
-	m.pruneSnapshotHeights = []int64{slices.Max(loadedPruneSnapshotHeights)}
-	m.initFromStore = true
+	m.completedSnapshotHeight = slices.Max(loadedPruneSnapshotHeights)
+	m.loadedFromDisk = true
 	return nil
 }
 

store/pruning/manager_test.go

@@ -25,6 +25,122 @@ func TestNewManager(t *testing.T) {
 	require.Equal(t, types.PruningNothing, manager.GetOptions().GetPruningStrategy())
 }
 
+func TestFailSnapshot(t *testing.T) {
+	manager := NewManager(db.NewMemDB(), log.NewNopLogger())
+	manager.SetOptions(types.NewCustomPruningOptions(10, 1))
+	manager.SetSnapshotInterval(10)
+	manager.StartSnapshot(10)
+	manager.FailSnapshot(10)
+	manager.FailSnapshot(10)
+
+	assert.Empty(t, manager.inflightSnapshotHeights)
+	assert.Equal(t, int64(9), manager.GetPruningHeight(20))
+}
+
+func TestSnapshotLifecycleTransitions(t *testing.T) {
+	newManager := func() *Manager {
+		manager := NewManager(db.NewMemDB(), log.NewNopLogger())
+		manager.SetOptions(types.NewCustomPruningOptions(1, 10))
+		manager.SetSnapshotInterval(10)
+		return manager
+	}
+
+	manager := newManager()
+	manager.StartSnapshot(10)
+	manager.FailSnapshot(10)
+	manager.StartSnapshot(20)
+	manager.CompleteSnapshot(20)
+	assert.Equal(t, int64(28), manager.GetPruningHeight(30))
+
+	manager = newManager()
+	manager.StartSnapshot(10)
+	manager.StartSnapshot(20)
+	manager.CompleteSnapshot(20)
+	assert.Equal(t, int64(9), manager.GetPruningHeight(30))
+	manager.FailSnapshot(10)
+	assert.Equal(t, int64(28), manager.GetPruningHeight(30))
+}
+
+func TestSnapshotLifecycleRestartAndDuplicateCalls(t *testing.T) {
+	db := db.NewMemDB()
+	manager := NewManager(db, log.NewNopLogger())
+	manager.SetOptions(types.NewCustomPruningOptions(1, 10))
+	manager.SetSnapshotInterval(10)
+	manager.StartSnapshot(10)
+	manager.StartSnapshot(10)
+	assert.Len(t, manager.inflightSnapshotHeights, 1)
+	manager.CompleteSnapshot(10)
+	manager.CompleteSnapshot(10)
+
+	restarted := NewManager(db, log.NewNopLogger())
+	restarted.SetOptions(types.NewCustomPruningOptions(1, 10))
+	restarted.SetSnapshotInterval(10)
+	require.NoError(t, restarted.LoadSnapshotHeights(db))
+	assert.Empty(t, restarted.inflightSnapshotHeights)
+	assert.Equal(t, int64(9), restarted.GetPruningHeight(20))
+	restarted.CompleteSnapshot(10)
+	restarted.CompleteSnapshot(5)
+	assert.True(t, restarted.loadedFromDisk)
+	assert.Equal(t, int64(9), restarted.GetPruningHeight(20))
+	restarted.StartSnapshot(20)
+	restarted.CompleteSnapshot(20)
+	assert.False(t, restarted.loadedFromDisk)
+	assert.Equal(t, int64(28), restarted.GetPruningHeight(30))
+}
+
+func TestCompleteSnapshotDoesNotPersistDuplicateOrLowerHeight(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	manager := NewManager(mock.NewMockDB(ctrl), log.NewNopLogger())
+	manager.SetOptions(types.NewPruningOptions(types.PruningEverything))
+	manager.completedSnapshotHeight = 10
+
+	manager.StartSnapshot(10)
+	manager.CompleteSnapshot(10)
+	manager.StartSnapshot(5)
+	manager.CompleteSnapshot(5)
+
+	assert.Empty(t, manager.inflightSnapshotHeights)
+	assert.Equal(t, int64(10), manager.completedSnapshotHeight)
+}
+
+func TestStartSnapshotDoesNotTrackCompletedHeight(t *testing.T) {
+	manager := NewManager(db.NewMemDB(), log.NewNopLogger())
+	manager.SetOptions(types.NewPruningOptions(types.PruningEverything))
+	manager.completedSnapshotHeight = 10
+
+	manager.StartSnapshot(10)
+	manager.StartSnapshot(5)
+
+	assert.Empty(t, manager.inflightSnapshotHeights)
+}
+
+func TestCompleteSnapshotPersistenceFailureIsRetryable(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	dbMock := mock.NewMockDB(ctrl)
+	dbMock.EXPECT().SetSync(pruneSnapshotHeightsKey, int64SliceToBytes(10)).Return(errors.New(dbErr))
+	dbMock.EXPECT().SetSync(pruneSnapshotHeightsKey, int64SliceToBytes(10)).Return(nil)
+
+	manager := NewManager(dbMock, log.NewNopLogger())
+	manager.SetOptions(types.NewPruningOptions(types.PruningEverything))
+	manager.SetSnapshotInterval(10)
+	manager.completedSnapshotHeight = 5
+	manager.loadedFromDisk = true
+	manager.StartSnapshot(10)
+
+	require.Panics(t, func() {
+		manager.CompleteSnapshot(10)
+	})
+	assert.Equal(t, int64(5), manager.completedSnapshotHeight)
+	assert.True(t, manager.loadedFromDisk)
+	assert.Contains(t, manager.inflightSnapshotHeights, int64(10))
+	assert.Equal(t, int64(4), manager.GetPruningHeight(20))
+
+	manager.CompleteSnapshot(10)
+	assert.Equal(t, int64(10), manager.completedSnapshotHeight)
+	assert.False(t, manager.loadedFromDisk)
+	assert.Empty(t, manager.inflightSnapshotHeights)
+}
+
 func TestStrategies(t *testing.T) {
 	testcases := map[string]struct {
 		strategy         types.PruningOptions
@@ -113,7 +229,7 @@ func TestStrategies(t *testing.T) {
 				if tc.snapshotInterval != 0 {
 					if curHeight > int64(tc.snapshotInterval) && curHeight%int64(tc.snapshotInterval) == int64(tc.snapshotInterval)-1 {
 						snapHeight := curHeight - int64(tc.snapshotInterval) + 1
-						manager.AnnounceSnapshotHeight(snapHeight)
+						manager.StartSnapshot(snapHeight)
 						manager.HandleSnapshotHeight(snapHeight)
 						snHeight = curHeight
 					}
@@ -222,7 +338,7 @@ func TestGetPruningHeight(t *testing.T) {
 			opts:        types.PruningOptions{KeepRecent: 5, Interval: 10, Strategy: types.PruningCustom},
 			setup: func(mgr *Manager) {
 				mgr.SetSnapshotInterval(15)
-				mgr.AnnounceSnapshotHeight(15)
+				mgr.StartSnapshot(15)
 				mgr.HandleSnapshotHeight(15)
 			},
 			exp: map[int64]int64{
@@ -238,7 +354,7 @@ func TestGetPruningHeight(t *testing.T) {
 			opts:        types.PruningOptions{KeepRecent: 5, Interval: 10, Strategy: types.PruningCustom},
 			setup: func(mgr *Manager) {
 				mgr.SetSnapshotInterval(15)
-				mgr.AnnounceSnapshotHeight(15)
+				mgr.StartSnapshot(15)
 			},
 			exp: map[int64]int64{
 				10: 4, // 10 - 5 (keep) - 1
@@ -250,8 +366,8 @@ func TestGetPruningHeight(t *testing.T) {
 			opts:        types.PruningOptions{KeepRecent: 5, Interval: 10, Strategy: types.PruningCustom},
 			setup: func(mgr *Manager) {
 				mgr.SetSnapshotInterval(15)
-				mgr.AnnounceSnapshotHeight(15)
-				mgr.AnnounceSnapshotHeight(30)
+				mgr.StartSnapshot(15)
+				mgr.StartSnapshot(30)
 				mgr.HandleSnapshotHeight(30)
 			},
 			exp: map[int64]int64{

store/rootmulti/store.go

@@ -81,7 +81,7 @@ type Store struct {
 var (
 	_ types.CommitMultiStore          = (*Store)(nil)
 	_ types.Queryable                 = (*Store)(nil)
-	_ snapshottypes.SnapshotAnnouncer = (*Store)(nil)
+	_ snapshottypes.SnapshotLifecycle = (*Store)(nil)
 )
 
 // NewStore returns a reference to a new Store object with the provided DB. The
@@ -365,6 +365,21 @@ func (rs *Store) AnnounceSnapshotHeight(height int64) {
 	rs.pruningManager.AnnounceSnapshotHeight(height)
 }
 
+// StartSnapshot tracks a snapshot while it is being created.
+func (rs *Store) StartSnapshot(height int64) {
+	rs.pruningManager.StartSnapshot(height)
+}
+
+// CompleteSnapshot records a successfully completed snapshot.
+func (rs *Store) CompleteSnapshot(height int64) {
+	rs.pruningManager.CompleteSnapshot(height)
+}
+
+// FailSnapshot removes a failed snapshot from in-flight tracking.
+func (rs *Store) FailSnapshot(height int64) {
+	rs.pruningManager.FailSnapshot(height)
+}
+
 // SetInterBlockCache sets the Store's internal inter-block (persistent) cache.
 // When this is defined, all CommitKVStores will be wrapped with their respective
 // inter-block cache.