[Persistence] Adds Save and Load functionality to TreeStore

persistence/blockstore/block_store.go

@@ -93,6 +93,10 @@ func (bs *blockStore) Stop() error {
 	return bs.kv.Stop()
 }
 
+func (bs *blockStore) Backup(path string) error {
+	return bs.kv.Backup(path)
+}
+
 func (bs *blockStore) Delete(key []byte) error         { return bs.kv.Delete(key) }
 func (bs *blockStore) Exists(key []byte) (bool, error) { return bs.kv.Exists(key) }
 func (bs *blockStore) GetAll(prefixKey []byte, descending bool) (keys, values [][]byte, err error) {

persistence/kvstore/kvstore.go

@@ -4,7 +4,10 @@ package kvstore
 
 import (
 	"errors"
+	"fmt"
 	"log"
+	"os"
+	"path/filepath"
 
 	badger "github.com/dgraph-io/badger/v3"
 	"github.com/pokt-network/smt"
@@ -22,6 +25,8 @@ type KVStore interface {
 	GetAll(prefixKey []byte, descending bool) (keys, values [][]byte, err error)
 	Exists(key []byte) (bool, error)
 	ClearAll() error
+
+	Backup(filepath string) error
 }
 
 const (
@@ -141,6 +146,30 @@ func (store *badgerKVStore) Stop() error {
 	return store.db.Close()
 }
 
+// Backup creates a backup for the badgerDB at the provided path.
+// It creates a file
+func (store *badgerKVStore) Backup(backupPath string) error {
+	// create backup directory if it doesn't exist
+	if err := os.MkdirAll(filepath.Dir(backupPath), os.ModePerm); err != nil {
+		return fmt.Errorf("failed to create backup directory: %v", err)
+	}
+
+	// create the backup file itself
+	backupFile, err := os.Create(backupPath)
+	if err != nil {
+		return fmt.Errorf("failed to create backup file: %v", err)
+	}
+	defer backupFile.Close()
+
+	// dump the database to the backup file
+	_, err = store.db.Backup(backupFile, 0)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // PrefixEndBytes returns the end byteslice for a noninclusive range
 // that would include all byte slices for which the input is the prefix
 func prefixEndBytes(prefix []byte) []byte {

persistence/kvstore/kvstore_test.go

@@ -2,6 +2,9 @@ package kvstore
 
 import (
 	"encoding/hex"
+	"io"
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 
@@ -330,10 +333,74 @@ func TestKVStore_ClearAll(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestKVStore_Backup(t *testing.T) {
+	t.Run("should backup an in-memory database", func(t *testing.T) {
+		store := NewMemKVStore()
+		require.NotNil(t, store)
+
+		tmpdir := t.TempDir()
+		path := filepath.Join(tmpdir, "TestKVStore_Backup_InMemory.bak")
+		require.NoError(t, store.Backup(path))
+
+		empty, err := isEmpty(t, tmpdir)
+		require.NoError(t, err)
+		require.False(t, empty)
+
+		// open the directory and assert on individual files
+		dir, err := os.Open(tmpdir)
+		require.NoError(t, err)
+		defer dir.Close()
+
+		files, err := dir.Readdir(0) // 0 means read all directory entries
+		require.NoError(t, err)
+		require.Equal(t, len(files), 1)
+	})
+	t.Run("should backup an on-disk store database", func(t *testing.T) {
+		tmpdir := t.TempDir()
+		kvpath := filepath.Join(tmpdir, "TestKVStore_Backup_OnDisk_Source.bak")
+		store, err := NewKVStore(kvpath)
+		require.NoError(t, err)
+		require.NotNil(t, store)
+
+		backupDir := t.TempDir()
+		path := filepath.Join(backupDir, "TestKVStore_Backup_OnDisk_Destination.bak")
+		require.NoError(t, store.Backup(path))
+
+		empty, err := isEmpty(t, backupDir)
+		require.NoError(t, err)
+		require.False(t, empty)
+
+		// open the directory and assert on individual files
+		dir, err := os.Open(backupDir)
+		require.NoError(t, err)
+		defer dir.Close()
+
+		files, err := dir.Readdir(0) // 0 means read all directory entries
+		require.NoError(t, err)
+		require.NoError(t, err)
+		require.Equal(t, len(files), 1)
+	})
+}
+
 func setupStore(t *testing.T, store KVStore) {
 	t.Helper()
 	err := store.Set([]byte("foo"), []byte("bar"))
 	require.NoError(t, err)
 	err = store.Set([]byte("baz"), []byte("bin"))
 	require.NoError(t, err)
 }
+
+func isEmpty(t *testing.T, dir string) (bool, error) {
+	t.Helper()
+	f, err := os.Open(dir)
+	if err != nil {
+		return false, err
+	}
+	defer f.Close()
+
+	_, err = f.Readdirnames(1)
+	if err == io.EOF {
+		return true, nil
+	}
+	return false, err
+}

persistence/module.go

@@ -99,12 +99,15 @@ func (*persistenceModule) Create(bus modules.Bus, options ...modules.ModuleOptio
 		return nil, err
 	}
 
-	_, err = trees.Create(
+	treemod, err := trees.Create(
 		bus,
 		trees.WithTreeStoreDirectory(persistenceCfg.TreesStoreDir),
 		trees.WithLogger(m.logger))
 	if err != nil {
-		return nil, fmt.Errorf("failed to create TreeStoreModule: %w", err)
+		return nil, fmt.Errorf("failed to create %s: %w", modules.TreeStoreSubmoduleName, err)
+	}
+	if err := treemod.Start(); err != nil {
+		return nil, fmt.Errorf("failed to start %s: %w", modules.TreeStoreSubmoduleName, err)
 	}
 
 	m.config = persistenceCfg

persistence/trees/atomic_test.go

@@ -2,14 +2,16 @@ package trees
 
 import (
 	"encoding/hex"
+	"io"
+	"os"
 	"testing"
 
-	"github.com/golang/mock/gomock"
 	"github.com/pokt-network/pocket/logger"
 	mock_types "github.com/pokt-network/pocket/persistence/types/mocks"
 	"github.com/pokt-network/pocket/shared/modules"
-	mockModules "github.com/pokt-network/pocket/shared/modules/mocks"
+	mock_modules "github.com/pokt-network/pocket/shared/modules/mocks"
 
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -24,8 +26,8 @@ func TestTreeStore_AtomicUpdatesWithSuccessfulRollback(t *testing.T) {
 	ctrl := gomock.NewController(t)
 
 	mockTxIndexer := mock_types.NewMockTxIndexer(ctrl)
-	mockBus := mockModules.NewMockBus(ctrl)
-	mockPersistenceMod := mockModules.NewMockPersistenceModule(ctrl)
+	mockBus := mock_modules.NewMockBus(ctrl)
+	mockPersistenceMod := mock_modules.NewMockPersistenceModule(ctrl)
 
 	mockBus.EXPECT().GetPersistenceModule().AnyTimes().Return(mockPersistenceMod)
 	mockPersistenceMod.EXPECT().GetTxIndexer().AnyTimes().Return(mockTxIndexer)
@@ -89,4 +91,132 @@ func TestTreeStore_AtomicUpdatesWithSuccessfulRollback(t *testing.T) {
 	hash3 := ts.getStateHash()
 	require.Equal(t, hash3, hash2)
 	require.Equal(t, hash3, h1)
+	require.NoError(t, ts.Rollback())
+
+	// confirm it's not in the tree
+	v, err := ts.merkleTrees[TransactionsTreeName].tree.Get([]byte("fiz"))
+	require.NoError(t, err)
+	require.Nil(t, v)
+}
+
+func TestTreeStore_SaveAndLoad(t *testing.T) {
+	t.Parallel()
+	t.Run("should save a backup in a directory", func(t *testing.T) {
+		ts := newTestTreeStore(t)
+		tmpdir := t.TempDir()
+		// assert that the directory is empty before backup
+		ok, err := isEmpty(tmpdir)
+		require.NoError(t, err)
+		require.True(t, ok)
+
+		// Trigger a backup
+		require.NoError(t, ts.Backup(tmpdir))
+
+		// assert that the directory is not empty after Backup has returned
+		ok, err = isEmpty(tmpdir)
+		require.NoError(t, err)
+		require.False(t, ok)
+	})
+	t.Run("should load a backup and maintain TreeStore hash integrity", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		tmpDir := t.TempDir()
+
+		mockTxIndexer := mock_types.NewMockTxIndexer(ctrl)
+		mockBus := mock_modules.NewMockBus(ctrl)
+		mockPersistenceMod := mock_modules.NewMockPersistenceModule(ctrl)
+
+		mockBus.EXPECT().GetPersistenceModule().AnyTimes().Return(mockPersistenceMod)
+		mockPersistenceMod.EXPECT().GetTxIndexer().AnyTimes().Return(mockTxIndexer)
+
+		ts := &treeStore{
+			logger:       logger.Global.CreateLoggerForModule(modules.TreeStoreSubmoduleName),
+			treeStoreDir: tmpDir,
+		}
+		require.NoError(t, ts.Start())
+		require.NotNil(t, ts.rootTree.tree)
+
+		for _, treeName := range stateTreeNames {
+			err := ts.merkleTrees[treeName].tree.Update([]byte("foo"), []byte("bar"))
+			require.NoError(t, err)
+		}
+
+		err := ts.Commit()
+		require.NoError(t, err)
+
+		hash1 := ts.getStateHash()
+		require.NotEmpty(t, hash1)
+
+		w, err := ts.save()
+		require.NoError(t, err)
+		require.NotNil(t, w)
+		require.NotNil(t, w.rootHash)
+		require.NotNil(t, w.merkleRoots)
+
+		// Stop the first tree store so that it's databases are no longer used
+		require.NoError(t, ts.Stop())
+
+		// declare a second TreeStore with no trees then load the first worldstate into it
+		ts2 := &treeStore{
+			logger:       logger.Global.CreateLoggerForModule(modules.TreeStoreSubmoduleName),
+			treeStoreDir: tmpDir,
+		}
+
+		// Load sets a tree store to the provided worldstate
+		err = ts2.Load(w)
+		require.NoError(t, err)
+
+		hash2 := ts2.getStateHash()
+
+		// Assert that hash is unchanged from save and load
+		require.Equal(t, hash1, hash2)
+	})
+}
+
+// creates a new tree store with a tmp directory for nodestore persistence
+// and then starts the tree store and returns its pointer.
+func newTestTreeStore(t *testing.T) *treeStore {
+	t.Helper()
+	ctrl := gomock.NewController(t)
+	tmpDir := t.TempDir()
+
+	mockTxIndexer := mock_types.NewMockTxIndexer(ctrl)
+	mockBus := mock_modules.NewMockBus(ctrl)
+	mockPersistenceMod := mock_modules.NewMockPersistenceModule(ctrl)
+
+	mockBus.EXPECT().GetPersistenceModule().AnyTimes().Return(mockPersistenceMod)
+	mockPersistenceMod.EXPECT().GetTxIndexer().AnyTimes().Return(mockTxIndexer)
+
+	ts := &treeStore{
+		logger:       logger.Global.CreateLoggerForModule(modules.TreeStoreSubmoduleName),
+		treeStoreDir: tmpDir,
+	}
+	require.NoError(t, ts.Start())
+	require.NotNil(t, ts.rootTree.tree)
+
+	for _, treeName := range stateTreeNames {
+		err := ts.merkleTrees[treeName].tree.Update([]byte("foo"), []byte("bar"))
+		require.NoError(t, err)
+	}
+
+	err := ts.Commit()
+	require.NoError(t, err)
+
+	hash1 := ts.getStateHash()
+	require.NotEmpty(t, hash1)
+
+	return ts
+}
+
+func isEmpty(dir string) (bool, error) {
+	f, err := os.Open(dir)
+	if err != nil {
+		return false, err
+	}
+	defer f.Close()
+
+	_, err = f.Readdirnames(1) // Or f.Readdir(1)
+	if err == io.EOF {
+		return true, nil
+	}
+	return false, err // Either not empty or error, suits both cases
 }

persistence/trees/module.go

@@ -1,6 +1,8 @@
 package trees
 
 import (
+	"encoding/hex"
+	"errors"
 	"fmt"
 
 	"github.com/pokt-network/pocket/persistence/kvstore"
@@ -10,6 +12,8 @@ import (
 
 var _ modules.TreeStoreModule = &treeStore{}
 
+// Create returns a TreeStoreSubmodule that has been setup with the provided TreeStoreOptions, started,
+// and then registered to the bus.
 func (*treeStore) Create(bus modules.Bus, options ...modules.TreeStoreOption) (modules.TreeStoreModule, error) {
 	m := &treeStore{}
 
@@ -19,10 +23,6 @@ func (*treeStore) Create(bus modules.Bus, options ...modules.TreeStoreOption) (m
 
 	bus.RegisterModule(m)
 
-	if err := m.setupTrees(); err != nil {
-		return nil, err
-	}
-
 	return m, nil
 }
 
@@ -50,8 +50,30 @@ func WithTreeStoreDirectory(path string) modules.TreeStoreOption {
 	}
 }
 
+// Start loads up the trees from the configured tree store directory.
+func (t *treeStore) Start() error {
+	return t.setupTrees()
+}
+
+// Stop shuts down the database connection to the nodestore for the root tree and then for each merkle tree.
+// If Commit has not been called before Stop is called, data will be lost.
+func (t *treeStore) Stop() error {
+	t.logger.Debug().Msgf("🛑 tree store stop initiated at %s 🛑", hex.EncodeToString(t.rootTree.tree.Root()))
+	errs := []error{}
+	if err := t.rootTree.nodeStore.Stop(); err != nil {
+		errs = append(errs, err)
+	}
+	for _, st := range t.merkleTrees {
+		if err := st.nodeStore.Stop(); err != nil {
+			errs = append(errs, err)
+		}
+	}
+	return errors.Join(errs...)
+}
+
 func (t *treeStore) GetModuleName() string { return modules.TreeStoreSubmoduleName }
 
+// setupTrees is called by Start and it loads the treestore at the given directory
 func (t *treeStore) setupTrees() error {
 	if t.treeStoreDir == ":memory:" {
 		return t.setupInMemory()