memdb, badger: implement DBConnection.Revert

db/badgerdb/db.go

@@ -2,9 +2,9 @@ package badgerdb
 
 import (
 	"bytes"
+	"context"
 	"encoding/csv"
 	"errors"
-	"math"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -15,6 +15,8 @@ import (
 	dbutil "github.com/cosmos/cosmos-sdk/db/internal"
 
 	"github.com/dgraph-io/badger/v3"
+	bpb "github.com/dgraph-io/badger/v3/pb"
+	"github.com/dgraph-io/ristretto/z"
 )
 
 var (
@@ -112,6 +114,7 @@ func readVersionsFile(path string) (*versionManager, error) {
 		return nil, err
 	}
 	var versions []uint64
+	var lastTs uint64
 	vmap := map[uint64]uint64{}
 	for _, row := range rows {
 		version, err := strconv.ParseUint(row[0], 10, 64)
@@ -122,14 +125,17 @@ func readVersionsFile(path string) (*versionManager, error) {
 		if err != nil {
 			return nil, err
 		}
+		if version == 0 { // 0 maps to the latest TS
+			lastTs = ts
+		}
 		versions = append(versions, version)
 		vmap[version] = ts
 	}
 	vmgr := dbm.NewVersionManager(versions)
 	return &versionManager{
 		VersionManager: vmgr,
 		vmap:           vmap,
-		lastTs:         vmgr.Last(),
+		lastTs:         lastTs,
 	}, nil
 }
 
@@ -142,6 +148,10 @@ func writeVersionsFile(vm *versionManager, path string) error {
 	defer file.Close()
 	w := csv.NewWriter(file)
 	var rows [][]string
+	rows = append(rows, []string{
+		strconv.FormatUint(0, 10),
+		strconv.FormatUint(vm.lastTs, 10),
+	})
 	for it := vm.Iterator(); it.Next(); {
 		version := it.Value()
 		ts, ok := vm.vmap[version]
@@ -157,16 +167,20 @@ func writeVersionsFile(vm *versionManager, path string) error {
 }
 
 func (b *BadgerDB) Reader() dbm.DBReader {
-	return &badgerTxn{txn: b.db.NewTransactionAt(math.MaxUint64, false), db: b}
+	b.mtx.RLock()
+	ts := b.vmgr.lastTs
+	b.mtx.RUnlock()
+	return &badgerTxn{txn: b.db.NewTransactionAt(ts, false), db: b}
 }
 
 func (b *BadgerDB) ReaderAt(version uint64) (dbm.DBReader, error) {
 	b.mtx.RLock()
 	defer b.mtx.RUnlock()
-	if !b.vmgr.Exists(version) {
+	ts, has := b.vmgr.versionTs(version)
+	if !has {
 		return nil, dbm.ErrVersionDoesNotExist
 	}
-	return &badgerTxn{txn: b.db.NewTransactionAt(b.vmgr.versionTs(version), false), db: b}, nil
+	return &badgerTxn{txn: b.db.NewTransactionAt(ts, false), db: b}, nil
 }
 
 func (b *BadgerDB) ReadWriter() dbm.DBReadWriter {
@@ -232,6 +246,80 @@ func (b *BadgerDB) DeleteVersion(target uint64) error {
 	return nil
 }
 
+func (b *BadgerDB) Revert() error {
+	b.mtx.RLock()
+	defer b.mtx.RUnlock()
+	if b.openWriters > 0 {
+		return dbm.ErrOpenTransactions
+	}
+
+	last := b.vmgr.Last()
+	if last == 0 {
+		return dbm.ErrInvalidVersion
+	}
+	// Revert from latest commit TS to last "saved" TS
+	target, has := b.vmgr.versionTs(last)
+	if !has {
+		return errors.New("bad version history")
+	}
+	lastTs := b.vmgr.lastTs
+	if target == lastTs {
+		return nil
+	}
+
+	// Badger provides no way to rollback committed data, so we undo all changes
+	// since the target version using the Stream API
+	stream := b.db.NewStreamAt(lastTs)
+	// Skips unchanged keys
+	stream.ChooseKey = func(item *badger.Item) bool { return item.Version() > target }
+	// Scans for value at target version
+	stream.KeyToList = func(key []byte, itr *badger.Iterator) (*bpb.KVList, error) {
+		kv := bpb.KV{Key: key}
+		// advance down to <= target version
+		itr.Next() // we have at least one newer version
+		for itr.Valid() && bytes.Equal(key, itr.Item().Key()) && itr.Item().Version() > target {
+			itr.Next()
+		}
+		if itr.Valid() && bytes.Equal(key, itr.Item().Key()) && !itr.Item().IsDeletedOrExpired() {
+			var err error
+			kv.Value, err = itr.Item().ValueCopy(nil)
+			if err != nil {
+				return nil, err
+			}
+		}
+		return &bpb.KVList{Kv: []*bpb.KV{&kv}}, nil
+	}
+	txn := b.db.NewTransactionAt(lastTs, true)
+	defer txn.Discard()
+	stream.Send = func(buf *z.Buffer) error {
+		kvl, err := badger.BufferToKVList(buf)
+		if err != nil {
+			return err
+		}
+		// nil Value indicates a deleted entry
+		for _, kv := range kvl.Kv {
+			if kv.Value == nil {
+				err = txn.Delete(kv.Key)
+				if err != nil {
+					return err
+				}
+			} else {
+				err = txn.Set(kv.Key, kv.Value)
+				if err != nil {
+					return err
+				}
+			}
+		}
+		return nil
+	}
+
+	err := stream.Orchestrate(context.Background())
+	if err != nil {
+		return err
+	}
+	return txn.CommitAt(lastTs, nil)
+}
+
 func (b *BadgerDB) Stats() map[string]string { return nil }
 
 func (tx *badgerTxn) Get(key []byte) ([]byte, error) {
@@ -385,14 +473,23 @@ func (i *badgerIterator) Value() []byte {
 	return val
 }
 
-func (vm *versionManager) versionTs(ver uint64) uint64 {
-	return vm.vmap[ver]
+func (vm *versionManager) versionTs(ver uint64) (uint64, bool) {
+	ts, has := vm.vmap[ver]
+	return ts, has
+}
+
+// updateCommitTs increments the lastTs if equal to readts.
+func (vm *versionManager) updateCommitTs(readts uint64) {
+	if vm.lastTs == readts {
+		vm.lastTs += 1
+	}
 }
 
 // Atomically accesses the last commit timestamp used as a version marker.
 func (vm *versionManager) lastCommitTs() uint64 {
 	return atomic.LoadUint64(&vm.lastTs)
 }
+
 func (vm *versionManager) Copy() *versionManager {
 	vmap := map[uint64]uint64{}
 	for ver, ts := range vm.vmap {
@@ -405,12 +502,6 @@ func (vm *versionManager) Copy() *versionManager {
 	}
 }
 
-// updateCommitTs increments the lastTs if equal to readts.
-func (vm *versionManager) updateCommitTs(readts uint64) {
-	if vm.lastTs == readts {
-		vm.lastTs += 1
-	}
-}
 func (vm *versionManager) Save(target uint64) (uint64, error) {
 	id, err := vm.VersionManager.Save(target)
 	if err != nil {
@@ -419,3 +510,8 @@ func (vm *versionManager) Save(target uint64) (uint64, error) {
 	vm.vmap[id] = vm.lastTs // non-atomic, already guarded by the vmgr mutex
 	return id, nil
 }
+
+func (vm *versionManager) Delete(target uint64) {
+	vm.VersionManager.Delete(target)
+	delete(vm.vmap, target)
+}

db/badgerdb/db_test.go

@@ -31,6 +31,11 @@ func TestVersioning(t *testing.T) {
 	dbtest.DoTestVersioning(t, load)
 }
 
+func TestRevert(t *testing.T) {
+	dbtest.DoTestRevert(t, load, false)
+	dbtest.DoTestRevert(t, load, true)
+}
+
 func TestReloadDB(t *testing.T) {
 	dbtest.DoTestReloadDB(t, load)
 }

db/memdb/db.go

@@ -159,6 +159,30 @@ func (db *MemDB) DeleteVersion(target uint64) error {
 	return nil
 }
 
+func (db *MemDB) Revert() error {
+	db.mtx.RLock()
+	defer db.mtx.RUnlock()
+	if db.openWriters > 0 {
+		return dbm.ErrOpenTransactions
+	}
+
+	last := db.vmgr.Last()
+	if last == 0 {
+		return dbm.ErrInvalidVersion // revert needs at least one version
+	}
+	var has bool
+	db.btree, has = db.saved[last]
+	if !has {
+		return fmt.Errorf("bad version history: version %v not saved", last)
+	}
+	for ver, _ := range db.saved {
+		if ver > last {
+			delete(db.saved, ver)
+		}
+	}
+	return nil
+}
+
 // Get implements DBReader.
 func (tx *dbTxn) Get(key []byte) ([]byte, error) {
 	if tx.btree == nil {

db/memdb/db_test.go

@@ -44,6 +44,10 @@ func TestVersioning(t *testing.T) {
 	dbtest.DoTestVersioning(t, load)
 }
 
+func TestRevert(t *testing.T) {
+	dbtest.DoTestRevert(t, load, false)
+}
+
 func TestTransactions(t *testing.T) {
 	dbtest.DoTestTransactions(t, load, false)
 }