From c242ef648da5545c039ae518c07f46b1ed44dc93 Mon Sep 17 00:00:00 2001 From: kbhat1 Date: Fri, 26 Jun 2026 10:56:31 -0400 Subject: [PATCH 1/2] Add offline rollback utility for LittDB MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds rollback.RollbackLittDB(dataDirs, rollbackFilter): a static, offline utility that rewinds a LittDB instance to a chosen point — e.g. to roll a node's state back to a specific block height while the database is stopped. For each table it walks the key files newest-to-oldest, invoking the filter once per record (isPrimary is derived from the record's KeyKind). The first record the filter accepts is the rollback point: that record's whole key group and everything written before it are kept; everything written after the group is permanently removed from the segment files. - segment.Segment.RollbackToKeyCount truncates a sealed segment in place: atomic key-file swap, value-file truncation, metadata key-count update. - The orchestrator locks the data dirs (so it won't touch a live DB), then discards the table's keymap and snapshot, deletes whole newer segments (highest index first), and truncates the rollback segment. The keymap and snapshot are derived state and are rebuilt from the truncated segments on the next start, which keeps them consistent with the rolled-back data — the same approach cli/prune.go uses to finalize an offline mutation. Steps are ordered for crash-safety: discarding the keymap first ensures the DB always rebuilds it from whatever segment state exists rather than trusting a keymap that points into truncated data. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../litt/disktable/segment/rollback.go | 94 ++++++ sei-db/db_engine/litt/rollback/rollback.go | 282 ++++++++++++++++++ .../db_engine/litt/rollback/rollback_test.go | 253 ++++++++++++++++ 3 files changed, 629 insertions(+) create mode 100644 sei-db/db_engine/litt/disktable/segment/rollback.go create mode 100644 sei-db/db_engine/litt/rollback/rollback.go create mode 100644 sei-db/db_engine/litt/rollback/rollback_test.go diff --git a/sei-db/db_engine/litt/disktable/segment/rollback.go b/sei-db/db_engine/litt/disktable/segment/rollback.go new file mode 100644 index 0000000000..01a05acd4b --- /dev/null +++ b/sei-db/db_engine/litt/disktable/segment/rollback.go @@ -0,0 +1,94 @@ +package segment + +import ( + "fmt" + "os" +) + +// RollbackToKeyCount truncates a sealed segment so that it retains only its first survivingKeyCount key-file +// records (in the order they were written), discarding every key and value written afterwards. +// +// This is a destructive, offline operation. The caller must guarantee that the database is not running and +// that nothing else is touching the segment's files. +// +// survivingKeyCount counts individual key-file records; a primary key and each of its secondary keys count +// separately. It must not exceed the number of records currently in the segment. To keep the segment +// internally consistent, callers should pass a count that lands on a group boundary (a primary plus all of +// its secondaries are either all kept or all discarded); RollbackToKeyCount itself does not enforce this. +// +// The steps are ordered so that an interruption never leaves a torn record: +// 1. the key file is rewritten via an atomic swap (the commit point), +// 2. each shard's value file is truncated, and +// 3. the segment's key count is recorded in the metadata file. +// +// The surviving records always occupy a contiguous prefix of the key file and of each shard's value file, +// so the addresses of the kept records are never disturbed. +func (s *Segment) RollbackToKeyCount(survivingKeyCount uint32) error { + if !s.IsSealed() { + return fmt.Errorf("segment %d is not sealed, cannot roll back", s.index) + } + + keys, err := s.keys.readKeys() + if err != nil { + return fmt.Errorf("failed to read keys for segment %d: %w", s.index, err) + } + + if int(survivingKeyCount) > len(keys) { + return fmt.Errorf("surviving key count %d exceeds the %d records in segment %d", + survivingKeyCount, len(keys), s.index) + } + if int(survivingKeyCount) == len(keys) { + // Nothing was written after the boundary; the segment is already in its target state. + return nil + } + + survivingKeys := keys[:survivingKeyCount] + + // 1. Rewrite the key file with only the surviving records. The atomic rename of the swap file over the + // original key file is the commit point for this segment's rollback. + swapFile, err := createKeyFile(s.logger, s.index, s.keys.segmentPath, true) + if err != nil { + return fmt.Errorf("failed to create swap key file for segment %d: %w", s.index, err) + } + for _, key := range survivingKeys { + if err = swapFile.write(key); err != nil { + return fmt.Errorf("failed to write key to swap file for segment %d: %w", s.index, err) + } + } + if err = swapFile.seal(); err != nil { + return fmt.Errorf("failed to seal swap key file for segment %d: %w", s.index, err) + } + if err = swapFile.atomicSwap(s.fsync); err != nil { + return fmt.Errorf("failed to swap key file for segment %d: %w", s.index, err) + } + s.keys = swapFile + + // 2. Truncate each shard's value file to the end of its last surviving value. Values carry no length + // prefix on disk, so a value occupies exactly [offset, offset+valueSize), and the surviving values form + // a prefix of each shard because values are appended in write order. + shardEnds := make([]uint64, len(s.shards)) + for _, key := range survivingKeys { + shardID := key.Address.ShardID() + if int(shardID) >= len(s.shards) { + return fmt.Errorf("segment %d has a key with shard ID %d outside its sharding factor %d", + s.index, shardID, len(s.shards)) + } + end := uint64(key.Address.Offset()) + uint64(key.Address.ValueSize()) + if end > shardEnds[shardID] { + shardEnds[shardID] = end + } + } + for shardID, valueFile := range s.shards { + if err = os.Truncate(valueFile.path(), int64(shardEnds[shardID])); err != nil { //nolint:gosec // value offsets are bounded to 2^32 + return fmt.Errorf("failed to truncate value file for segment %d shard %d: %w", s.index, shardID, err) + } + } + + // 3. Record the new key count. The seal time is preserved so the segment's TTL/expiry is unaffected. + s.metadata.keyCount = survivingKeyCount + if err = s.metadata.write(); err != nil { + return fmt.Errorf("failed to update metadata for segment %d: %w", s.index, err) + } + + return nil +} diff --git a/sei-db/db_engine/litt/rollback/rollback.go b/sei-db/db_engine/litt/rollback/rollback.go new file mode 100644 index 0000000000..66c364c0db --- /dev/null +++ b/sei-db/db_engine/litt/rollback/rollback.go @@ -0,0 +1,282 @@ +// Package rollback implements an offline rollback utility for LittDB. It rewinds a database to a chosen +// point by discarding the most recently written keys, and is intended for operational use (for example, +// rolling a node's state back to a specific block height) while the database is not running. +package rollback + +import ( + "context" + "fmt" + "log/slog" + "os" + "path/filepath" + "sort" + "time" + + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable/keymap" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable/segment" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/util" +) + +// RollbackFilter decides where a rollback stops. It is invoked once per key-file record, walking each +// table from the most recently written key to the oldest. isPrimary is true for primary keys (standalone +// primaries and primaries that own secondary keys) and false for secondary keys. The first record for +// which the filter returns true is the rollback point: that record's group is kept along with everything +// written before it, and everything written after the group is discarded. +type RollbackFilter func(key []byte, isPrimary bool) (bool, error) + +// RollbackLittDB performs an offline rollback of the LittDB instance stored across the given data +// directories (the same paths passed to the database as its storage roots). +// +// For every table found under dataDirs, RollbackLittDB walks the key files from newest to oldest and +// invokes rollbackFilter for each key. The first key for which the filter returns true marks the rollback +// point: that key's group (a primary plus any secondary keys written with it) and everything written +// before it are retained; everything written after the group is permanently deleted from the segment +// files. A table for which the filter never returns true is left unchanged. +// +// The keymap and any snapshot are discarded rather than edited: the database rebuilds both from the +// truncated segment files the next time it starts, which keeps them exactly consistent with the +// rolled-back data (the same approach cli/prune.go uses after an offline mutation). +// +// The database must NOT be running while this is called. RollbackLittDB takes the same directory locks the +// database uses, so it will fail rather than corrupt a live database, and it assumes nothing else mutates +// the files while it works. +// +// The operation is idempotent: if it is interrupted, re-running it with the same filter completes the +// rollback. (An interrupted run may briefly leave a segment's recorded key count stale, but that value only +// feeds metrics and self-corrects on the next run.) +func RollbackLittDB(dataDirs []string, rollbackFilter RollbackFilter) error { + logger := slog.Default() + + if len(dataDirs) == 0 { + return fmt.Errorf("no data directories provided") + } + if rollbackFilter == nil { + return fmt.Errorf("rollback filter must not be nil") + } + + roots := make([]string, len(dataDirs)) + for i, dir := range dataDirs { + sanitized, err := util.SanitizePath(dir) + if err != nil { + return fmt.Errorf("invalid data directory %q: %w", dir, err) + } + roots[i] = sanitized + } + + // Refuse to operate on a database that is in active use. The DB holds these same locks while running. + releaseLocks, err := util.LockDirectories(logger, roots, util.LockfileName, true) + if err != nil { + return fmt.Errorf("failed to lock data directories %v: %w", roots, err) + } + defer releaseLocks() + + tables, err := findTables(roots) + if err != nil { + return fmt.Errorf("failed to enumerate tables under %v: %w", roots, err) + } + + for _, table := range tables { + if err := rollbackTable(logger, roots, table, rollbackFilter); err != nil { + return fmt.Errorf("failed to roll back table %q: %w", table, err) + } + } + + return nil +} + +// rollbackPoint identifies where a table's rollback boundary falls: the segment that contains the matched +// key and the number of key-file records to retain in that segment. Everything after that prefix in the +// rollback segment, and every newer segment, is discarded. +type rollbackPoint struct { + segmentIndex uint32 + survivingKeyCount uint32 +} + +// rollbackTable rolls back a single table. +func rollbackTable( + logger *slog.Logger, + roots []string, + tableName string, + rollbackFilter RollbackFilter, +) error { + errorMonitor := util.NewErrorMonitor(context.Background(), logger, nil) + + segmentPaths, err := segment.BuildSegmentPaths(roots, "", tableName) + if err != nil { + return fmt.Errorf("failed to build segment paths: %w", err) + } + + lowestSegmentIndex, highestSegmentIndex, segments, err := segment.GatherSegmentFiles( + logger, errorMonitor, segmentPaths, false /* snapshottingEnabled */, time.Now(), + true /* cleanOrphans */, true /* fsync */) + if err != nil { + return fmt.Errorf("failed to gather segment files: %w", err) + } + if len(segments) == 0 { + logger.Info("table has no segments, nothing to roll back", "table", tableName) + return nil + } + + // Refuse to operate on a symlinked snapshot directory: truncating symlinked value files would corrupt + // the real segment data they point at. Rollback must run against the database's storage roots, not a + // snapshot. (cli/prune.go makes the same check.) + isSnapshot, err := segments[lowestSegmentIndex].IsSnapshot() + if err != nil { + return fmt.Errorf("failed to determine whether table %q is a snapshot: %w", tableName, err) + } + if isSnapshot { + return fmt.Errorf("table %q is a symlinked snapshot; refusing to roll back "+ + "(point the tool at the database's storage roots, not a snapshot)", tableName) + } + + pivot, err := findRollbackPoint(segments, lowestSegmentIndex, highestSegmentIndex, rollbackFilter) + if err != nil { + return err + } + if pivot == nil { + logger.Warn("no rollback point found, leaving table unchanged", "table", tableName) + return nil + } + + logger.Info("rolling back table", + "table", tableName, + "rollbackSegment", pivot.segmentIndex, + "survivingRecordsInRollbackSegment", pivot.survivingKeyCount, + "deletedSegments", highestSegmentIndex-pivot.segmentIndex, + ) + + // 1. Discard the derived keymap and snapshot first. They are rebuilt from the segment files on the next + // start, so doing this before touching the segments guarantees that however the steps below are + // interrupted, the database rebuilds the keymap from whatever segment state exists rather than trusting + // a keymap that points into truncated or deleted data. + if err = discardDerivedState(roots, tableName); err != nil { + return fmt.Errorf("failed to discard derived state: %w", err) + } + + // 2. Delete whole segments newer than the rollback segment, highest index first so that an interruption + // never leaves a gap in the middle of the segment sequence. + for segmentIndex := highestSegmentIndex; segmentIndex > pivot.segmentIndex; segmentIndex-- { + for _, filePath := range segments[segmentIndex].GetFilePaths() { + if err = util.DeepDelete(filePath); err != nil { + return fmt.Errorf("failed to delete %s: %w", filePath, err) + } + } + } + + // 3. Truncate the rollback segment down to the surviving records. + if err = segments[pivot.segmentIndex].RollbackToKeyCount(pivot.survivingKeyCount); err != nil { + return fmt.Errorf("failed to truncate segment %d: %w", pivot.segmentIndex, err) + } + + return nil +} + +// findRollbackPoint walks the table's key files from the newest record to the oldest, invoking the filter +// on each. The first record for which the filter returns true is the rollback point. The whole group that +// contains that record (a standalone primary, or a primary together with the secondaries that follow it) +// is retained, so the surviving boundary is set to the end of that group. Returns nil if no record matches. +func findRollbackPoint( + segments map[uint32]*segment.Segment, + lowestSegmentIndex uint32, + highestSegmentIndex uint32, + rollbackFilter RollbackFilter, +) (*rollbackPoint, error) { + + for segmentIndex := highestSegmentIndex; ; segmentIndex-- { + keys, err := segments[segmentIndex].GetKeys() + if err != nil { + return nil, fmt.Errorf("failed to read keys from segment %d: %w", segmentIndex, err) + } + + for i := len(keys) - 1; i >= 0; i-- { + match, err := rollbackFilter(keys[i].Key, keys[i].Kind.IsPrimary()) + if err != nil { + return nil, fmt.Errorf("rollback filter returned an error in segment %d: %w", segmentIndex, err) + } + if match { + groupEnd, err := groupEndIndex(keys, i) + if err != nil { + return nil, fmt.Errorf("segment %d: %w", segmentIndex, err) + } + return &rollbackPoint{ + segmentIndex: segmentIndex, + survivingKeyCount: uint32(groupEnd + 1), //nolint:gosec // bounded by the segment's key count + }, nil + } + } + + if segmentIndex == lowestSegmentIndex { + break + } + } + + return nil, nil +} + +// groupEndIndex returns the index of the last record in the group that contains record i. A group is +// either a single standalone primary, or a primary followed by one or more secondaries terminated by a +// KeyKindFinalSecondary record. +func groupEndIndex(keys []*types.ScopedKey, i int) (int, error) { + if keys[i].Kind == types.KeyKindStandalone { + return i, nil + } + for j := i; j < len(keys); j++ { + if keys[j].Kind == types.KeyKindFinalSecondary { + return j, nil + } + } + return 0, fmt.Errorf("key group starting at record %d has no terminating final-secondary record", i) +} + +// discardDerivedState removes the keymap and snapshot directories for a table from every root. Both are +// derived entirely from the segment files: the database rebuilds the keymap (via reloadKeymap) and the +// snapshot on its next start, so deleting them forces both back into sync with the truncated segments. +// Leaving them would let a stale keymap reference discarded keys, or let a snapshot's hard links pin the +// rolled-back data on disk. Removing a directory that does not exist is a no-op. +func discardDerivedState(roots []string, tableName string) error { + for _, root := range roots { + dirs := []string{ + filepath.Join(root, tableName, keymap.KeymapDirectoryName), + filepath.Join(root, tableName, segment.HardLinkDirectory), + } + for _, dir := range dirs { + if err := os.RemoveAll(dir); err != nil { + return fmt.Errorf("failed to remove %s: %w", dir, err) + } + } + } + return nil +} + +// findTables returns the names of all LittDB tables found under the given roots. A table is any directory +// that contains a "segments" sub-directory. +func findTables(roots []string) ([]string, error) { + tableSet := make(map[string]struct{}) + for _, root := range roots { + entries, err := os.ReadDir(root) + if err != nil { + return nil, fmt.Errorf("failed to read directory %s: %w", root, err) + } + for _, entry := range entries { + if !entry.IsDir() { + continue + } + segmentsDir := filepath.Join(root, entry.Name(), segment.SegmentDirectory) + isDir, err := util.IsDirectory(segmentsDir) + if err != nil { + return nil, fmt.Errorf("failed to check directory %s: %w", segmentsDir, err) + } + if isDir { + tableSet[entry.Name()] = struct{}{} + } + } + } + + tables := make([]string, 0, len(tableSet)) + for table := range tableSet { + tables = append(tables, table) + } + sort.Strings(tables) + return tables, nil +} diff --git a/sei-db/db_engine/litt/rollback/rollback_test.go b/sei-db/db_engine/litt/rollback/rollback_test.go new file mode 100644 index 0000000000..5f38215de6 --- /dev/null +++ b/sei-db/db_engine/litt/rollback/rollback_test.go @@ -0,0 +1,253 @@ +package rollback + +import ( + "fmt" + "path/filepath" + "strconv" + "strings" + "testing" + + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/littbuilder" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/util" + "github.com/stretchr/testify/require" +) + +const rollbackTestTable = "rollback-test" + +// newRollbackTestDB returns a littDB config (with a handful of storage roots) and an open database, sized +// so that even a modest number of writes produces many sealed segments. The caller must close the DB. +func newRollbackTestDB(t *testing.T) (*litt.Config, []string) { + t.Helper() + rand := util.NewTestRandom() + testDirectory := t.TempDir() + + rootPathCount := rand.Uint64Range(1, 4) + rootPaths := make([]string, rootPathCount) + for i := uint64(0); i < rootPathCount; i++ { + rootPaths[i] = filepath.Join(testDirectory, fmt.Sprintf("root-%d", i)) + } + + config, err := litt.DefaultConfig(rootPaths...) + require.NoError(t, err) + config.Fsync = false + config.DoubleWriteProtection = true + // A tiny target file size forces the data to be spread over many sealed segments, exercising both + // whole-segment deletion and partial (truncating) rollback of a single segment. + config.TargetSegmentFileSize = 100 + + return config, rootPaths +} + +// writeSequentialKeys writes count primary keys named "key-NNNNN" in index order, each with value +// "value-NNNNN", returning the value for every key index. The DB is flushed and closed before returning so +// the data is sealed on disk and ready for an offline rollback. +func writeSequentialKeys(t *testing.T, config *litt.Config, count int) map[int][]byte { + t.Helper() + + db, err := littbuilder.NewDB(config) + require.NoError(t, err) + + tableConfig := litt.DefaultTableConfig(rollbackTestTable) + tableConfig.ShardingFactor = 3 + table, err := db.BuildTable(tableConfig) + require.NoError(t, err) + + values := make(map[int][]byte, count) + for i := 0; i < count; i++ { + value := []byte(fmt.Sprintf("value-%05d", i)) + require.NoError(t, table.Put(keyForIndex(i), value)) + values[i] = value + } + + require.NoError(t, table.Flush()) + require.NoError(t, db.Close()) + return values +} + +func keyForIndex(i int) []byte { + return []byte(fmt.Sprintf("key-%05d", i)) +} + +// indexFromKey parses the integer index encoded in a "key-NNNNN" key. +func indexFromKey(t *testing.T, key []byte) int { + t.Helper() + idx, err := strconv.Atoi(strings.TrimPrefix(string(key), "key-")) + require.NoError(t, err) + return idx +} + +// openTable reopens the test table. +func openTable(t *testing.T, config *litt.Config) (litt.DB, litt.Table) { + t.Helper() + db, err := littbuilder.NewDB(config) + require.NoError(t, err) + tableConfig := litt.DefaultTableConfig(rollbackTestTable) + tableConfig.ShardingFactor = 3 + table, err := db.BuildTable(tableConfig) + require.NoError(t, err) + return db, table +} + +// assertSequentialState verifies that exactly indices [0, keepThrough] are present (with their original +// values) and all higher indices are absent. +func assertSequentialState(t *testing.T, table litt.Table, count int, keepThrough int, values map[int][]byte) { + t.Helper() + for i := 0; i < count; i++ { + got, ok, err := table.Get(keyForIndex(i)) + require.NoError(t, err) + if i <= keepThrough { + require.Truef(t, ok, "key %d should survive rollback", i) + require.Equalf(t, values[i], got, "value mismatch for surviving key %d", i) + } else { + require.Falsef(t, ok, "key %d should have been rolled back", i) + } + } + require.Equal(t, uint64(keepThrough+1), table.KeyCount()) +} + +// TestRollbackLittDB rolls back to a key in the middle of the write history and verifies that the surviving +// keys keep their values and the newer keys are gone. +func TestRollbackLittDB(t *testing.T) { + t.Parallel() + + const count = 200 + const keepThrough = 137 + + config, roots := newRollbackTestDB(t) + values := writeSequentialKeys(t, config, count) + + err := RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { + require.True(t, isPrimary) // these are all standalone primary keys + return indexFromKey(t, key) <= keepThrough, nil + }) + require.NoError(t, err) + + // The rollback discards the keymap, so reopening rebuilds it from the truncated segment files. A correct + // read of every surviving key therefore also confirms the segments were truncated consistently. + db, table := openTable(t, config) + assertSequentialState(t, table, count, keepThrough, values) + require.NoError(t, db.Close()) +} + +// TestRollbackNoMatch verifies that a table for which the filter never returns true is left untouched. +func TestRollbackNoMatch(t *testing.T) { + t.Parallel() + + const count = 50 + + config, roots := newRollbackTestDB(t) + values := writeSequentialKeys(t, config, count) + + err := RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { + return false, nil + }) + require.NoError(t, err) + + db, table := openTable(t, config) + assertSequentialState(t, table, count, count-1, values) // everything survives + require.NoError(t, db.Close()) +} + +// TestRollbackKeepsEverything verifies that when the newest key matches, nothing is deleted. +func TestRollbackKeepsEverything(t *testing.T) { + t.Parallel() + + const count = 50 + + config, roots := newRollbackTestDB(t) + values := writeSequentialKeys(t, config, count) + + err := RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { + return true, nil // the very first key visited (the newest) matches + }) + require.NoError(t, err) + + db, table := openTable(t, config) + assertSequentialState(t, table, count, count-1, values) + require.NoError(t, db.Close()) +} + +// TestRollbackPropagatesFilterError verifies that an error from the filter aborts the rollback. +func TestRollbackPropagatesFilterError(t *testing.T) { + t.Parallel() + + config, roots := newRollbackTestDB(t) + writeSequentialKeys(t, config, 20) + + wantErr := fmt.Errorf("boom") + err := RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { + return false, wantErr + }) + require.ErrorIs(t, err, wantErr) +} + +// TestRollbackWithSecondaryKeys verifies that secondary keys are handled correctly: the rollback point's +// whole group (its primary plus the secondaries written after it) is retained, isPrimary is reported +// correctly to the filter, and discarded groups lose both their primary and secondary keys. +func TestRollbackWithSecondaryKeys(t *testing.T) { + t.Parallel() + + const count = 60 + const keepThrough = 28 + + config, roots := newRollbackTestDB(t) + + db, err := littbuilder.NewDB(config) + require.NoError(t, err) + tableConfig := litt.DefaultTableConfig(rollbackTestTable) + tableConfig.ShardingFactor = 2 + table, err := db.BuildTable(tableConfig) + require.NoError(t, err) + + primaryKey := func(i int) []byte { return []byte(fmt.Sprintf("pk-%05d", i)) } + secondaryKey := func(i int) []byte { return []byte(fmt.Sprintf("sk-%05d", i)) } + + values := make(map[int][]byte, count) + for i := 0; i < count; i++ { + value := []byte(fmt.Sprintf("value-%05d", i)) + // One secondary key aliasing the entire value. + secondary := &types.SecondaryKey{Key: secondaryKey(i), Offset: 0, Length: uint32(len(value))} + require.NoError(t, table.Put(primaryKey(i), value, secondary)) + values[i] = value + } + require.NoError(t, table.Flush()) + require.NoError(t, db.Close()) + + err = RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { + // Only primary keys carry an index we want to stop on; secondaries are reported with isPrimary=false. + if !isPrimary { + require.True(t, strings.HasPrefix(string(key), "sk-")) + return false, nil + } + require.True(t, strings.HasPrefix(string(key), "pk-")) + idx, err := strconv.Atoi(strings.TrimPrefix(string(key), "pk-")) + require.NoError(t, err) + return idx <= keepThrough, nil + }) + require.NoError(t, err) + + db, err = littbuilder.NewDB(config) + require.NoError(t, err) + table, err = db.BuildTable(tableConfig) + require.NoError(t, err) + + for i := 0; i < count; i++ { + gotPrimary, okPrimary, err := table.Get(primaryKey(i)) + require.NoError(t, err) + gotSecondary, okSecondary, err := table.Get(secondaryKey(i)) + require.NoError(t, err) + + if i <= keepThrough { + require.Truef(t, okPrimary, "primary %d should survive", i) + require.Equal(t, values[i], gotPrimary) + require.Truef(t, okSecondary, "secondary %d should survive (same group as its primary)", i) + require.Equal(t, values[i], gotSecondary) + } else { + require.Falsef(t, okPrimary, "primary %d should be rolled back", i) + require.Falsef(t, okSecondary, "secondary %d should be rolled back", i) + } + } + require.NoError(t, db.Close()) +} From 9e2d6c84584b04e21eeca6cd30ce74b3adcf8d18 Mon Sep 17 00:00:00 2001 From: kbhat1 Date: Mon, 6 Jul 2026 14:02:12 -0400 Subject: [PATCH 2/2] Address PR review: table-scoped filter, gc-watermark guard, idempotent truncation - RollbackFilter now receives the table name (per @cody-littley's review): key schemas differ across tables, so the filter must decode keys table-aware. - Refuse to roll a table back below its durable gc-watermark. Segments below it were logically garbage collected; a target there cannot be faithfully reconstructed and would leave lowestReadableSegment above the highest surviving segment, which the DB rejects at startup. The watermark is left in place (not deleted) so a keymap rebuild does not resurrect collected keys. - Segment.RollbackToKeyCount no longer early-returns when the key file already holds exactly the surviving records: it still truncates the value files and rewrites the metadata key count, so a run interrupted after the atomic key-file swap is repaired on re-invocation. This makes the documented idempotency guarantee actually hold. - Tests: per-table filter (table-name assertion), gc-watermark refusal leaves the table untouched, and idempotent re-run. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../litt/disktable/segment/rollback.go | 49 +++++++------ sei-db/db_engine/litt/rollback/rollback.go | 72 ++++++++++++++----- .../db_engine/litt/rollback/rollback_test.go | 71 ++++++++++++++++-- 3 files changed, 148 insertions(+), 44 deletions(-) diff --git a/sei-db/db_engine/litt/disktable/segment/rollback.go b/sei-db/db_engine/litt/disktable/segment/rollback.go index 01a05acd4b..1ca98e3e22 100644 --- a/sei-db/db_engine/litt/disktable/segment/rollback.go +++ b/sei-db/db_engine/litt/disktable/segment/rollback.go @@ -17,10 +17,14 @@ import ( // its secondaries are either all kept or all discarded); RollbackToKeyCount itself does not enforce this. // // The steps are ordered so that an interruption never leaves a torn record: -// 1. the key file is rewritten via an atomic swap (the commit point), +// 1. the key file is rewritten via an atomic swap (the commit point), skipped when it already holds exactly +// the surviving records, // 2. each shard's value file is truncated, and // 3. the segment's key count is recorded in the metadata file. // +// It is idempotent: re-invoking with the same survivingKeyCount is a no-op on an already-rolled-back segment +// and repairs one whose earlier run was interrupted after step 1 (finishing steps 2 and 3). +// // The surviving records always occupy a contiguous prefix of the key file and of each shard's value file, // so the addresses of the kept records are never disturbed. func (s *Segment) RollbackToKeyCount(survivingKeyCount uint32) error { @@ -37,31 +41,34 @@ func (s *Segment) RollbackToKeyCount(survivingKeyCount uint32) error { return fmt.Errorf("surviving key count %d exceeds the %d records in segment %d", survivingKeyCount, len(keys), s.index) } - if int(survivingKeyCount) == len(keys) { - // Nothing was written after the boundary; the segment is already in its target state. - return nil - } survivingKeys := keys[:survivingKeyCount] - // 1. Rewrite the key file with only the surviving records. The atomic rename of the swap file over the - // original key file is the commit point for this segment's rollback. - swapFile, err := createKeyFile(s.logger, s.index, s.keys.segmentPath, true) - if err != nil { - return fmt.Errorf("failed to create swap key file for segment %d: %w", s.index, err) - } - for _, key := range survivingKeys { - if err = swapFile.write(key); err != nil { - return fmt.Errorf("failed to write key to swap file for segment %d: %w", s.index, err) + // 1. Rewrite the key file to contain only the surviving records. This is skipped when the key file + // already holds exactly those records — either nothing was written after the rollback boundary, or a + // prior run was interrupted right after this step. The atomic rename of the swap file over the original + // key file is the commit point for dropping records; steps 2 and 3 below always run, so a rollback + // interrupted after this swap is still repaired (over-long value files truncated, stale metadata key + // count corrected) when it is re-invoked. + if int(survivingKeyCount) < len(keys) { + var swapFile *keyFile + swapFile, err = createKeyFile(s.logger, s.index, s.keys.segmentPath, true) + if err != nil { + return fmt.Errorf("failed to create swap key file for segment %d: %w", s.index, err) } + for _, key := range survivingKeys { + if err = swapFile.write(key); err != nil { + return fmt.Errorf("failed to write key to swap file for segment %d: %w", s.index, err) + } + } + if err = swapFile.seal(); err != nil { + return fmt.Errorf("failed to seal swap key file for segment %d: %w", s.index, err) + } + if err = swapFile.atomicSwap(s.fsync); err != nil { + return fmt.Errorf("failed to swap key file for segment %d: %w", s.index, err) + } + s.keys = swapFile } - if err = swapFile.seal(); err != nil { - return fmt.Errorf("failed to seal swap key file for segment %d: %w", s.index, err) - } - if err = swapFile.atomicSwap(s.fsync); err != nil { - return fmt.Errorf("failed to swap key file for segment %d: %w", s.index, err) - } - s.keys = swapFile // 2. Truncate each shard's value file to the end of its last surviving value. Values carry no length // prefix on disk, so a value occupies exactly [offset, offset+valueSize), and the surviving values form diff --git a/sei-db/db_engine/litt/rollback/rollback.go b/sei-db/db_engine/litt/rollback/rollback.go index 66c364c0db..44101bd8c0 100644 --- a/sei-db/db_engine/litt/rollback/rollback.go +++ b/sei-db/db_engine/litt/rollback/rollback.go @@ -12,39 +12,42 @@ import ( "sort" "time" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable/keymap" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable/segment" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/util" ) -// RollbackFilter decides where a rollback stops. It is invoked once per key-file record, walking each -// table from the most recently written key to the oldest. isPrimary is true for primary keys (standalone -// primaries and primaries that own secondary keys) and false for secondary keys. The first record for -// which the filter returns true is the rollback point: that record's group is kept along with everything -// written before it, and everything written after the group is discarded. -type RollbackFilter func(key []byte, isPrimary bool) (bool, error) +// RollbackFilter decides where a rollback stops. It is invoked once per key-file record, walking each table +// from the most recently written key to the oldest. tableName identifies the table the record belongs to: +// key schemas differ across tables, so the filter must decode the key in a table-aware way. isPrimary is +// true for primary keys (standalone primaries and primaries that own secondary keys) and false for +// secondary keys. The first record for which the filter returns true is the rollback point: that record's +// group is kept along with everything written before it, and everything written after the group is discarded. +type RollbackFilter func(tableName string, key []byte, isPrimary bool) (bool, error) // RollbackLittDB performs an offline rollback of the LittDB instance stored across the given data // directories (the same paths passed to the database as its storage roots). // -// For every table found under dataDirs, RollbackLittDB walks the key files from newest to oldest and -// invokes rollbackFilter for each key. The first key for which the filter returns true marks the rollback -// point: that key's group (a primary plus any secondary keys written with it) and everything written -// before it are retained; everything written after the group is permanently deleted from the segment -// files. A table for which the filter never returns true is left unchanged. +// For every table found under dataDirs, RollbackLittDB walks that table's key files from newest to oldest +// and invokes rollbackFilter(tableName, key, isPrimary) for each record. The first key for which the filter +// returns true marks the rollback point: that key's group (a primary plus any secondary keys written with +// it) and everything written before it are retained; everything written after the group is permanently +// deleted from the segment files. A table for which the filter never returns true is left unchanged. // // The keymap and any snapshot are discarded rather than edited: the database rebuilds both from the // truncated segment files the next time it starts, which keeps them exactly consistent with the -// rolled-back data (the same approach cli/prune.go uses after an offline mutation). +// rolled-back data (the same approach cli/prune.go uses after an offline mutation). The durable gc-watermark +// is deliberately left in place; rolling a table back below its gc-watermark — into data that garbage +// collection has already reclaimed — is refused, because that state cannot be faithfully reconstructed. // // The database must NOT be running while this is called. RollbackLittDB takes the same directory locks the // database uses, so it will fail rather than corrupt a live database, and it assumes nothing else mutates // the files while it works. // -// The operation is idempotent: if it is interrupted, re-running it with the same filter completes the -// rollback. (An interrupted run may briefly leave a segment's recorded key count stale, but that value only -// feeds metrics and self-corrects on the next run.) +// The operation is idempotent and safe to re-run: re-running with the same filter completes (and repairs) +// a rollback that was interrupted partway through. func RollbackLittDB(dataDirs []string, rollbackFilter RollbackFilter) error { logger := slog.Default() @@ -130,7 +133,7 @@ func rollbackTable( "(point the tool at the database's storage roots, not a snapshot)", tableName) } - pivot, err := findRollbackPoint(segments, lowestSegmentIndex, highestSegmentIndex, rollbackFilter) + pivot, err := findRollbackPoint(segments, tableName, lowestSegmentIndex, highestSegmentIndex, rollbackFilter) if err != nil { return err } @@ -139,6 +142,21 @@ func rollbackTable( return nil } + // Refuse to roll back below the durable gc-watermark. Segments below it are logically garbage collected + // (their keys were durably removed from the keymap), so a rollback target there cannot be faithfully + // reconstructed — and it would leave lowestReadableSegment above the highest surviving segment, which the + // database rejects at startup. We must not simply delete the watermark either: keymap reload honors it so + // a rebuild does not resurrect collected keys. + watermark, defined, err := highestGCWatermark(roots, tableName) + if err != nil { + return err + } + if defined && watermark > pivot.segmentIndex { + return fmt.Errorf("cannot roll back table %q to segment %d: it is below the gc-watermark "+ + "(lowest readable segment %d); that data has already been garbage collected", + tableName, pivot.segmentIndex, watermark) + } + logger.Info("rolling back table", "table", tableName, "rollbackSegment", pivot.segmentIndex, @@ -178,6 +196,7 @@ func rollbackTable( // is retained, so the surviving boundary is set to the end of that group. Returns nil if no record matches. func findRollbackPoint( segments map[uint32]*segment.Segment, + tableName string, lowestSegmentIndex uint32, highestSegmentIndex uint32, rollbackFilter RollbackFilter, @@ -190,7 +209,7 @@ func findRollbackPoint( } for i := len(keys) - 1; i >= 0; i-- { - match, err := rollbackFilter(keys[i].Key, keys[i].Kind.IsPrimary()) + match, err := rollbackFilter(tableName, keys[i].Key, keys[i].Kind.IsPrimary()) if err != nil { return nil, fmt.Errorf("rollback filter returned an error in segment %d: %w", segmentIndex, err) } @@ -249,6 +268,25 @@ func discardDerivedState(roots []string, tableName string) error { return nil } +// highestGCWatermark returns the highest gc-watermark (lowest readable segment index) recorded for a table +// across all roots, and whether any root defined one. Segments below this index have been logically garbage +// collected. The watermark lives at the table root and survives a keymap rebuild, matching how the database +// loads it at startup. +func highestGCWatermark(roots []string, tableName string) (watermark uint32, defined bool, err error) { + for _, root := range roots { + f, err := disktable.LoadGCWatermarkFile(filepath.Join(root, tableName)) + if err != nil { + return 0, false, fmt.Errorf("failed to load gc-watermark for table %q under %s: %w", + tableName, root, err) + } + if f.IsDefined() && (!defined || f.LowestReadableSegment() > watermark) { + watermark = f.LowestReadableSegment() + defined = true + } + } + return watermark, defined, nil +} + // findTables returns the names of all LittDB tables found under the given roots. A table is any directory // that contains a "segments" sub-directory. func findTables(roots []string) ([]string, error) { diff --git a/sei-db/db_engine/litt/rollback/rollback_test.go b/sei-db/db_engine/litt/rollback/rollback_test.go index 5f38215de6..929593c2a8 100644 --- a/sei-db/db_engine/litt/rollback/rollback_test.go +++ b/sei-db/db_engine/litt/rollback/rollback_test.go @@ -8,6 +8,7 @@ import ( "testing" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/littbuilder" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/util" @@ -118,8 +119,9 @@ func TestRollbackLittDB(t *testing.T) { config, roots := newRollbackTestDB(t) values := writeSequentialKeys(t, config, count) - err := RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { - require.True(t, isPrimary) // these are all standalone primary keys + err := RollbackLittDB(roots, func(tableName string, key []byte, isPrimary bool) (bool, error) { + require.Equal(t, rollbackTestTable, tableName) // filter is invoked per table + require.True(t, isPrimary) // these are all standalone primary keys return indexFromKey(t, key) <= keepThrough, nil }) require.NoError(t, err) @@ -140,7 +142,7 @@ func TestRollbackNoMatch(t *testing.T) { config, roots := newRollbackTestDB(t) values := writeSequentialKeys(t, config, count) - err := RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { + err := RollbackLittDB(roots, func(tableName string, key []byte, isPrimary bool) (bool, error) { return false, nil }) require.NoError(t, err) @@ -159,7 +161,7 @@ func TestRollbackKeepsEverything(t *testing.T) { config, roots := newRollbackTestDB(t) values := writeSequentialKeys(t, config, count) - err := RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { + err := RollbackLittDB(roots, func(tableName string, key []byte, isPrimary bool) (bool, error) { return true, nil // the very first key visited (the newest) matches }) require.NoError(t, err) @@ -177,7 +179,7 @@ func TestRollbackPropagatesFilterError(t *testing.T) { writeSequentialKeys(t, config, 20) wantErr := fmt.Errorf("boom") - err := RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { + err := RollbackLittDB(roots, func(tableName string, key []byte, isPrimary bool) (bool, error) { return false, wantErr }) require.ErrorIs(t, err, wantErr) @@ -215,7 +217,7 @@ func TestRollbackWithSecondaryKeys(t *testing.T) { require.NoError(t, table.Flush()) require.NoError(t, db.Close()) - err = RollbackLittDB(roots, func(key []byte, isPrimary bool) (bool, error) { + err = RollbackLittDB(roots, func(tableName string, key []byte, isPrimary bool) (bool, error) { // Only primary keys carry an index we want to stop on; secondaries are reported with isPrimary=false. if !isPrimary { require.True(t, strings.HasPrefix(string(key), "sk-")) @@ -251,3 +253,60 @@ func TestRollbackWithSecondaryKeys(t *testing.T) { } require.NoError(t, db.Close()) } + +// TestRollbackRefusesBelowGCWatermark verifies that rolling a table back to a point below its durable +// gc-watermark (data garbage collection has already reclaimed) is refused before any files are mutated. +func TestRollbackRefusesBelowGCWatermark(t *testing.T) { + t.Parallel() + + const count = 50 + + config, roots := newRollbackTestDB(t) + writeSequentialKeys(t, config, count) + + // Record a gc-watermark far above any surviving segment so the rollback point is guaranteed to fall + // below it. The watermark lives at the table root on one of the storage roots. + tableDir := filepath.Join(roots[0], rollbackTestTable) + watermark, err := disktable.LoadGCWatermarkFile(tableDir) + require.NoError(t, err) + require.NoError(t, watermark.Update(1_000_000)) + + keymapDir := filepath.Join(tableDir, "keymap") + existsBefore, err := util.Exists(keymapDir) + require.NoError(t, err) + require.True(t, existsBefore) + + err = RollbackLittDB(roots, func(tableName string, key []byte, isPrimary bool) (bool, error) { + return indexFromKey(t, key) <= 10, nil + }) + require.Error(t, err) + require.Contains(t, err.Error(), "gc-watermark") + + // The guard fires before any destructive step, so the derived state is left untouched. + existsAfter, err := util.Exists(keymapDir) + require.NoError(t, err) + require.True(t, existsAfter, "rollback must not mutate anything when it refuses") +} + +// TestRollbackIsIdempotent verifies that re-running the same rollback is a safe no-op that leaves the table +// in the same correct state (exercising RollbackToKeyCount's survivingKeyCount == len(keys) path). +func TestRollbackIsIdempotent(t *testing.T) { + t.Parallel() + + const count = 120 + const keepThrough = 71 + + config, roots := newRollbackTestDB(t) + values := writeSequentialKeys(t, config, count) + + filter := func(tableName string, key []byte, isPrimary bool) (bool, error) { + return indexFromKey(t, key) <= keepThrough, nil + } + + require.NoError(t, RollbackLittDB(roots, filter)) + require.NoError(t, RollbackLittDB(roots, filter)) // second run must be a safe no-op + + db, table := openTable(t, config) + assertSequentialState(t, table, count, keepThrough, values) + require.NoError(t, db.Close()) +}