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..1ca98e3e22 --- /dev/null +++ b/sei-db/db_engine/litt/disktable/segment/rollback.go @@ -0,0 +1,101 @@ +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), 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 { + 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) + } + + survivingKeys := keys[:survivingKeyCount] + + // 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 + } + + // 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..44101bd8c0 --- /dev/null +++ b/sei-db/db_engine/litt/rollback/rollback.go @@ -0,0 +1,320 @@ +// 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" + "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. 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 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). 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 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() + + 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, tableName, lowestSegmentIndex, highestSegmentIndex, rollbackFilter) + if err != nil { + return err + } + if pivot == nil { + logger.Warn("no rollback point found, leaving table unchanged", "table", tableName) + 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, + "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, + tableName string, + 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(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) + } + 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 +} + +// 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) { + 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..929593c2a8 --- /dev/null +++ b/sei-db/db_engine/litt/rollback/rollback_test.go @@ -0,0 +1,312 @@ +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/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" + "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(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) + + // 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(tableName string, 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(tableName string, 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(tableName string, 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(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-")) + 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()) +} + +// 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()) +}