fix(server): Fix pruning height calculation (#24583)

Co-authored-by: Alex | Interchain Labs <alex@interchainlabs.io>
Co-authored-by: Avory <avorycorelli@gmail.com>

store/pruning/export_test.go

@@ -1,8 +0,0 @@
-package pruning
-
-var (
-	PruneSnapshotHeightsKey = pruneSnapshotHeightsKey
-
-	Int64SliceToBytes          = int64SliceToBytes
-	LoadPruningSnapshotHeights = loadPruningSnapshotHeights
-)

store/pruning/manager.go

@@ -3,6 +3,7 @@ package pruning
 import (
 	"encoding/binary"
 	"fmt"
+	"slices"
 	"sort"
 	"sync"
 
@@ -21,12 +22,14 @@ type Manager struct {
 	opts             types.PruningOptions
 	snapshotInterval uint64
 	// Snapshots are taken in a separate goroutine from the regular execution
-	// and can be delivered asynchrounously via HandleSnapshotHeight.
-	// Therefore, we sync access to pruneSnapshotHeights with this mutex.
+	// and can be delivered asynchronously via HandleSnapshotHeight.
+	// Therefore, we sync access to pruneSnapshotHeights, inflightSnapshotHeights and initFromStore with this mutex.
 	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
+	pruneSnapshotHeights    []int64
+	inflightSnapshotHeights []int64
+	initFromStore           bool
 }
 
 // NegativeHeightsError is returned when a negative height is provided to the manager.
@@ -51,7 +54,7 @@ func NewManager(db dbm.DB, logger log.Logger) *Manager {
 		db:                   db,
 		logger:               logger,
 		opts:                 types.NewPruningOptions(types.PruningNothing),
-		pruneSnapshotHeights: []int64{0},
+		pruneSnapshotHeights: []int64{0}, // init with 0 block height
 	}
 }
 
@@ -65,6 +68,17 @@ 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) {
+	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)
+}
+
 // 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.
@@ -74,63 +88,88 @@ func (m *Manager) HandleSnapshotHeight(height int64) {
 		return
 	}
 
+	m.logger.Debug("HandleSnapshotHeight", "height", height)
+
 	m.pruneSnapshotHeightsMx.Lock()
 	defer m.pruneSnapshotHeightsMx.Unlock()
 
-	m.logger.Debug("HandleSnapshotHeight", "height", 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:]...)
+	}
+
+	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.pruneSnapshotHeights = append(m.pruneSnapshotHeights, height)
 	sort.Slice(m.pruneSnapshotHeights, func(i, j int) bool { return m.pruneSnapshotHeights[i] < m.pruneSnapshotHeights[j] })
+
+	// 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
 		}
 	}
+	// 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 updates to disk so that they are not lost if crash happens.
-	if err := m.db.SetSync(pruneSnapshotHeightsKey, int64SliceToBytes(m.pruneSnapshotHeights)); err != nil {
+	// 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 {
 		panic(err)
 	}
 }
 
 // SetSnapshotInterval sets the interval at which the snapshots are taken.
+// This value should be set on startup and not exceed max int64 (2^63-1). Concurrent modifications are not supported.
 func (m *Manager) SetSnapshotInterval(snapshotInterval uint64) {
 	m.snapshotInterval = snapshotInterval
 }
 
 // GetPruningHeight returns the height which can prune upto if it is able to prune at the given height.
 func (m *Manager) GetPruningHeight(height int64) int64 {
-	if m.opts.GetPruningStrategy() == types.PruningNothing {
-		return 0
-	}
-	if m.opts.Interval <= 0 {
-		return 0
-	}
-
-	if height%int64(m.opts.Interval) != 0 || height <= int64(m.opts.KeepRecent) {
+	if m.opts.GetPruningStrategy() == types.PruningNothing ||
+		m.opts.Interval <= 0 ||
+		height <= int64(m.opts.KeepRecent) ||
+		height%int64(m.opts.Interval) != 0 {
 		return 0
 	}
 
 	// Consider the snapshot height
 	pruneHeight := height - 1 - int64(m.opts.KeepRecent) // we should keep the current height at least
 
-	m.pruneSnapshotHeightsMx.RLock()
-	defer m.pruneSnapshotHeightsMx.RUnlock()
-
 	// snapshotInterval is zero, indicating that all heights can be pruned
 	if m.snapshotInterval <= 0 {
 		return pruneHeight
 	}
 
-	if len(m.pruneSnapshotHeights) == 0 { // the length should be greater than zero
+	m.pruneSnapshotHeightsMx.RLock()
+	defer m.pruneSnapshotHeightsMx.RUnlock()
+
+	if len(m.pruneSnapshotHeights) == 0 { // do not prune before an initial snapshot
 		return 0
 	}
 
-	// the snapshot `m.pruneSnapshotHeights[0]` is already operated,
-	// so we can prune upto `m.pruneSnapshotHeights[0] + int64(m.snapshotInterval) - 1`
-	snHeight := m.pruneSnapshotHeights[0] + int64(m.snapshotInterval) - 1
-	return min(snHeight, pruneHeight)
+	// 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)
+	}
+	if len(m.inflightSnapshotHeights) == 0 {
+		return min(snHeight, pruneHeight)
+	}
+	// highest version based on started snapshots
+	inFlightHeight := m.inflightSnapshotHeights[0] - 1
+	return min(snHeight, pruneHeight, inFlightHeight)
 }
 
 // LoadSnapshotHeights loads the snapshot heights from the database as a crash recovery.
@@ -139,20 +178,27 @@ func (m *Manager) LoadSnapshotHeights(db dbm.DB) error {
 		return nil
 	}
 
+	// loading list for backwards compatibility
 	loadedPruneSnapshotHeights, err := loadPruningSnapshotHeights(db)
 	if err != nil {
 		return err
 	}
 
-	if len(loadedPruneSnapshotHeights) > 0 {
-		m.pruneSnapshotHeightsMx.Lock()
-		defer m.pruneSnapshotHeightsMx.Unlock()
-		m.pruneSnapshotHeights = loadedPruneSnapshotHeights
+	if len(loadedPruneSnapshotHeights) == 0 {
+		return nil
 	}
-
+	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
 	return nil
 }
 
+func storePruningSnapshotHeight(db dbm.DB, val int64) error {
+	return db.SetSync(pruneSnapshotHeightsKey, int64SliceToBytes(val))
+}
+
 func loadPruningSnapshotHeights(db dbm.DB) ([]int64, error) {
 	bz, err := db.Get(pruneSnapshotHeightsKey)
 	if err != nil {
@@ -177,7 +223,7 @@ func loadPruningSnapshotHeights(db dbm.DB) ([]int64, error) {
 	return pruneSnapshotHeights, nil
 }
 
-func int64SliceToBytes(slice []int64) []byte {
+func int64SliceToBytes(slice ...int64) []byte {
 	bz := make([]byte, 0, len(slice)*8)
 	for _, ph := range slice {
 		buf := make([]byte, 8)

store/pruning/manager_test.go

@@ -1,4 +1,4 @@
-package pruning_test
+package pruning
 
 import (
 	"errors"
@@ -6,19 +6,19 @@ import (
 	"testing"
 
 	db "github.com/cosmos/cosmos-db"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
 	"cosmossdk.io/log"
 	"cosmossdk.io/store/mock"
-	"cosmossdk.io/store/pruning"
 	"cosmossdk.io/store/pruning/types"
 )
 
 const dbErr = "db error"
 
 func TestNewManager(t *testing.T) {
-	manager := pruning.NewManager(db.NewMemDB(), log.NewNopLogger())
+	manager := NewManager(db.NewMemDB(), log.NewNopLogger())
 	require.NotNil(t, manager)
 	require.Equal(t, types.PruningNothing, manager.GetOptions().GetPruningStrategy())
 }
@@ -78,7 +78,7 @@ func TestStrategies(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
 
-			manager := pruning.NewManager(db.NewMemDB(), log.NewNopLogger())
+			manager := NewManager(db.NewMemDB(), log.NewNopLogger())
 			require.NotNil(t, manager)
 
 			curStrategy := tc.strategy
@@ -110,7 +110,9 @@ func TestStrategies(t *testing.T) {
 			for curHeight := int64(0); curHeight < 110000; curHeight++ {
 				if tc.snapshotInterval != 0 {
 					if curHeight > int64(tc.snapshotInterval) && curHeight%int64(tc.snapshotInterval) == int64(tc.snapshotInterval)-1 {
-						manager.HandleSnapshotHeight(curHeight - int64(tc.snapshotInterval) + 1)
+						snapHeight := curHeight - int64(tc.snapshotInterval) + 1
+						manager.AnnounceSnapshotHeight(snapHeight)
+						manager.HandleSnapshotHeight(snapHeight)
 						snHeight = curHeight
 					}
 				}
@@ -185,8 +187,7 @@ func TestPruningHeight_Inputs(t *testing.T) {
 
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
-			manager := pruning.NewManager(db.NewMemDB(), log.NewNopLogger())
-			require.NotNil(t, manager)
+			manager := NewManager(db.NewMemDB(), log.NewNopLogger())
 			manager.SetOptions(types.NewPruningOptions(tc.strategy))
 
 			pruningHeightActual := manager.GetPruningHeight(tc.height)
@@ -195,6 +196,158 @@ func TestPruningHeight_Inputs(t *testing.T) {
 	}
 }
 
+func TestGetPruningHeight(t *testing.T) {
+	specs := map[string]struct {
+		initDBState int64
+		opts        types.PruningOptions
+		setup       func(manager *Manager)
+		exp         map[int64]int64
+	}{
+		"init from store - no snap": {
+			initDBState: 10,
+			opts:        types.PruningOptions{KeepRecent: 5, Interval: 10, Strategy: types.PruningCustom},
+			setup: func(mgr *Manager) {
+				mgr.SetSnapshotInterval(15)
+			},
+			exp: map[int64]int64{
+				20: 9, // initDBState - 1
+				30: 9, // initDBState - 1
+				45: 0, // not a prune height
+			},
+		},
+		"init from store - snap landed": {
+			initDBState: 10,
+			opts:        types.PruningOptions{KeepRecent: 5, Interval: 10, Strategy: types.PruningCustom},
+			setup: func(mgr *Manager) {
+				mgr.SetSnapshotInterval(15)
+				mgr.AnnounceSnapshotHeight(15)
+				mgr.HandleSnapshotHeight(15)
+			},
+			exp: map[int64]int64{
+				10: 4,  // 10 - 5 (keep) - 1
+				15: 0,  // not on prune interval
+				20: 14, // 20 - 5 (keep) - 1
+				30: 24, // 30 - 5 (keep) - 1
+				40: 29, // 15 (last completed snap) + 15 (snap interval) - 1
+			},
+		},
+		"init from store - snap in-flight": {
+			initDBState: 10,
+			opts:        types.PruningOptions{KeepRecent: 5, Interval: 10, Strategy: types.PruningCustom},
+			setup: func(mgr *Manager) {
+				mgr.SetSnapshotInterval(15)
+				mgr.AnnounceSnapshotHeight(15)
+			},
+			exp: map[int64]int64{
+				10: 4, // 10 - 5 (keep) - 1
+				20: 9, // 10 - 5 (keep) - 1
+			},
+		},
+		"init from store - delayed in-flight snap": {
+			initDBState: 10,
+			opts:        types.PruningOptions{KeepRecent: 5, Interval: 10, Strategy: types.PruningCustom},
+			setup: func(mgr *Manager) {
+				mgr.SetSnapshotInterval(15)
+				mgr.AnnounceSnapshotHeight(15)
+				mgr.AnnounceSnapshotHeight(30)
+				mgr.HandleSnapshotHeight(30)
+			},
+			exp: map[int64]int64{
+				10: 4,  // 10 - 5 (keep) - 1
+				20: 14, // 15 (in-flight) - 1
+				30: 14, // 15 (in-flight) - 1
+				40: 14, // 15 (in-flight) - 1
+			},
+		},
+		"empty store": {
+			opts: types.PruningOptions{KeepRecent: 5, Interval: 10, Strategy: types.PruningCustom},
+			setup: func(mgr *Manager) {
+				mgr.SetSnapshotInterval(15)
+			},
+			exp: map[int64]int64{
+				10: 4,  // 10 -5 (keep) -1
+				20: 14, // 20 -5 (keep) -1
+			},
+		},
+		"empty snap interval set": {
+			initDBState: 10,
+			opts:        types.PruningOptions{KeepRecent: 5, Interval: 10, Strategy: types.PruningCustom},
+			setup:       func(mgr *Manager) {},
+			exp: map[int64]int64{
+				10: 4,  // 10 -5 (keep) -1
+				20: 14, // 20 -5 (keep) -1
+			},
+		},
+		"prune nothing set": {
+			initDBState: 10,
+			opts:        types.PruningOptions{Strategy: types.PruningNothing, Interval: 10, KeepRecent: 5},
+			setup: func(mgr *Manager) {
+				mgr.SetSnapshotInterval(15)
+			},
+			exp: map[int64]int64{
+				10: 0, // nothing
+				20: 0, // nothing
+				30: 0, // nothing
+			},
+		},
+		"empty prune interval": {
+			initDBState: 10,
+			opts:        types.PruningOptions{Strategy: types.PruningCustom, KeepRecent: 5},
+			setup: func(mgr *Manager) {
+				mgr.SetSnapshotInterval(15)
+			},
+			exp: map[int64]int64{
+				10: 0, // interval required
+				20: 0, // interval required
+				30: 0, // interval required
+			},
+		},
+		"height <= keep": {
+			initDBState: 10,
+			opts:        types.PruningOptions{Strategy: types.PruningCustom, Interval: 1, KeepRecent: 5},
+			setup: func(mgr *Manager) {
+				mgr.SetSnapshotInterval(15)
+			},
+			exp: map[int64]int64{
+				0: 0, // interval required
+				4: 0, // interval required
+				5: 0, // interval required
+			},
+		},
+		"height not on prune interval": {
+			initDBState: 10,
+			opts:        types.PruningOptions{Strategy: types.PruningCustom, Interval: 2},
+			setup: func(mgr *Manager) {
+				mgr.SetSnapshotInterval(15)
+			},
+			exp: map[int64]int64{
+				0: 0, // excluded
+				1: 0, // not on prune interval
+				2: 1, // 2 - 1
+				3: 0, // not on prune interval
+				4: 3, // 2 - 1
+			},
+		},
+	}
+	for name, spec := range specs {
+		t.Run(name, func(t *testing.T) {
+			memDB := db.NewMemDB()
+			if spec.initDBState != 0 {
+				require.NoError(t, storePruningSnapshotHeight(memDB, spec.initDBState))
+			}
+			mgr2 := NewManager(memDB, log.NewNopLogger())
+			mgr2.SetOptions(spec.opts)
+			require.NoError(t, mgr2.LoadSnapshotHeights(memDB))
+			spec.setup(mgr2)
+
+			for height, exp := range spec.exp {
+				gotHeight := mgr2.GetPruningHeight(height)
+				assert.Equal(t, exp, gotHeight, "height: %d", height)
+			}
+		})
+	}
+}
+
 func TestHandleSnapshotHeight_DbErr_Panic(t *testing.T) {
 	ctrl := gomock.NewController(t)
 
@@ -203,7 +356,8 @@ func TestHandleSnapshotHeight_DbErr_Panic(t *testing.T) {
 
 	dbMock.EXPECT().SetSync(gomock.Any(), gomock.Any()).Return(errors.New(dbErr)).Times(1)
 
-	manager := pruning.NewManager(dbMock, log.NewNopLogger())
+	manager := NewManager(dbMock, log.NewNopLogger())
+	manager.SetSnapshotInterval(1)
 	manager.SetOptions(types.NewPruningOptions(types.PruningEverything))
 	require.NotNil(t, manager)
 
@@ -221,7 +375,7 @@ func TestHandleSnapshotHeight_LoadFromDisk(t *testing.T) {
 
 	// Setup
 	db := db.NewMemDB()
-	manager := pruning.NewManager(db, log.NewNopLogger())
+	manager := NewManager(db, log.NewNopLogger())
 	require.NotNil(t, manager)
 
 	manager.SetOptions(types.NewPruningOptions(types.PruningEverything))
@@ -236,7 +390,7 @@ func TestHandleSnapshotHeight_LoadFromDisk(t *testing.T) {
 			expected = 1
 		}
 
-		loadedSnapshotHeights, err := pruning.LoadPruningSnapshotHeights(db)
+		loadedSnapshotHeights, err := loadPruningSnapshotHeights(db)
 		require.NoError(t, err)
 		require.Equal(t, expected, len(loadedSnapshotHeights), snapshotHeightStr)
 
@@ -244,7 +398,7 @@ func TestHandleSnapshotHeight_LoadFromDisk(t *testing.T) {
 		err = manager.LoadSnapshotHeights(db)
 		require.NoError(t, err)
 
-		loadedSnapshotHeights, err = pruning.LoadPruningSnapshotHeights(db)
+		loadedSnapshotHeights, err = loadPruningSnapshotHeights(db)
 		require.NoError(t, err)
 		require.Equal(t, expected, len(loadedSnapshotHeights), snapshotHeightStr)
 	}
@@ -252,7 +406,7 @@ func TestHandleSnapshotHeight_LoadFromDisk(t *testing.T) {
 
 func TestLoadPruningSnapshotHeights(t *testing.T) {
 	var (
-		manager = pruning.NewManager(db.NewMemDB(), log.NewNopLogger())
+		manager = NewManager(db.NewMemDB(), log.NewNopLogger())
 		err     error
 	)
 	require.NotNil(t, manager)
@@ -268,7 +422,7 @@ func TestLoadPruningSnapshotHeights(t *testing.T) {
 			getFlushedPruningSnapshotHeights: func() []int64 {
 				return []int64{5, -2, 3}
 			},
-			expectedResult: &pruning.NegativeHeightsError{Height: -2},
+			expectedResult: &NegativeHeightsError{Height: -2},
 		},
 		"non-negative - success": {
 			getFlushedPruningSnapshotHeights: func() []int64 {
@@ -282,7 +436,7 @@ func TestLoadPruningSnapshotHeights(t *testing.T) {
 			db := db.NewMemDB()
 
 			if tc.getFlushedPruningSnapshotHeights != nil {
-				err = db.Set(pruning.PruneSnapshotHeightsKey, pruning.Int64SliceToBytes(tc.getFlushedPruningSnapshotHeights()))
+				err = db.Set(pruneSnapshotHeightsKey, int64SliceToBytes(tc.getFlushedPruningSnapshotHeights()...))
 				require.NoError(t, err)
 			}
 
@@ -293,7 +447,7 @@ func TestLoadPruningSnapshotHeights(t *testing.T) {
 }
 
 func TestLoadSnapshotHeights_PruneNothing(t *testing.T) {
-	manager := pruning.NewManager(db.NewMemDB(), log.NewNopLogger())
+	manager := NewManager(db.NewMemDB(), log.NewNopLogger())
 	require.NotNil(t, manager)
 
 	manager.SetOptions(types.NewPruningOptions(types.PruningNothing))

store/pruning/types/options.go

@@ -3,6 +3,7 @@ package types
 import (
 	"errors"
 	"fmt"
+	"math"
 )
 
 // PruningOptions defines the pruning strategy used when determining which
@@ -56,6 +57,7 @@ var (
 	ErrPruningIntervalZero       = errors.New("'pruning-interval' must not be 0. If you want to disable pruning, select pruning = \"nothing\"")
 	ErrPruningIntervalTooSmall   = fmt.Errorf("'pruning-interval' must not be less than %d. For the most aggressive pruning, select pruning = \"everything\"", pruneEverythingInterval)
 	ErrPruningKeepRecentTooSmall = fmt.Errorf("'pruning-keep-recent' must not be less than %d. For the most aggressive pruning, select pruning = \"everything\"", pruneEverythingKeepRecent)
+	ErrPruningKeepRecentTooBig   = errors.New("'pruning-keep-recent' must not be greater than 2^63-1. Select pruning = \"nothing\"")
 )
 
 func NewPruningOptions(pruningStrategy PruningStrategy) PruningOptions {
@@ -110,6 +112,9 @@ func (po PruningOptions) Validate() error {
 	if po.KeepRecent < pruneEverythingKeepRecent {
 		return ErrPruningKeepRecentTooSmall
 	}
+	if po.KeepRecent > math.MaxInt64 {
+		return ErrPruningKeepRecentTooBig
+	}
 	return nil
 }
 

store/pruning/types/options_test.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"math"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -21,6 +22,7 @@ func TestPruningOptions_Validate(t *testing.T) {
 		{NewCustomPruningOptions(2, 9), ErrPruningIntervalTooSmall},
 		{NewCustomPruningOptions(2, 0), ErrPruningIntervalZero},
 		{NewCustomPruningOptions(2, 0), ErrPruningIntervalZero},
+		{NewCustomPruningOptions(math.MaxInt64+1, 10), ErrPruningKeepRecentTooBig},
 	}
 
 	for _, tc := range testCases {