[Persistence] Adds atomic Update for TreeStore (#861)

Co-authored-by: Daniel Olshansky <[email protected]>

persistence/context.go

@@ -12,6 +12,7 @@ import (
 	"github.com/pokt-network/pocket/persistence/indexer"
 	coreTypes "github.com/pokt-network/pocket/shared/core/types"
 	"github.com/pokt-network/pocket/shared/modules"
+	"go.uber.org/multierr"
 )
 
 var _ modules.PersistenceRWContext = &PostgresContext{}
@@ -36,29 +37,35 @@ type PostgresContext struct {
 	networkId string
 }
 
-func (p *PostgresContext) NewSavePoint(bytes []byte) error {
-	p.logger.Info().Bool("TODO", true).Msg("NewSavePoint not implemented")
+// SetSavePoint generates a new Savepoint for this context.
+func (p *PostgresContext) SetSavePoint() error {
+	if err := p.stateTrees.Savepoint(); err != nil {
+		return err
+	}
 	return nil
 }
 
-// TECHDEBT(#327): Guarantee atomicity betweens `prepareBlock`, `insertBlock` and `storeBlock` for save points & rollbacks.
-func (p *PostgresContext) RollbackToSavePoint(bytes []byte) error {
-	p.logger.Info().Bool("TODO", true).Msg("RollbackToSavePoint not fully implemented")
-	return p.tx.Rollback(context.TODO())
+// RollbackToSavepoint triggers a rollback for the current pgx transaction and the underylying submodule stores.
+func (p *PostgresContext) RollbackToSavePoint() error {
+	ctx, _ := p.getCtxAndTx()
+	pgErr := p.tx.Rollback(ctx)
+	treesErr := p.stateTrees.Rollback()
+	return multierr.Combine(pgErr, treesErr)
 }
 
-// IMPROVE(#361): Guarantee the integrity of the state
 // Full details in the thread from the PR review: https://github.com/pokt-network/pocket/pull/285#discussion_r1018471719
 func (p *PostgresContext) ComputeStateHash() (string, error) {
 	stateHash, err := p.stateTrees.Update(p.tx, uint64(p.Height))
 	if err != nil {
 		return "", err
 	}
+	if err := p.stateTrees.Commit(); err != nil {
+		return "", err
+	}
 	p.stateHash = stateHash
 	return p.stateHash, nil
 }
 
-// TECHDEBT(#327): Make sure these operations are atomic
 func (p *PostgresContext) Commit(proposerAddr, quorumCert []byte) error {
 	p.logger.Info().Int64("height", p.Height).Msg("About to commit block & context")
 

persistence/db.go

@@ -37,6 +37,7 @@ var protocolActorSchemas = []types.ProtocolActorSchema{
 	types.ValidatorActor,
 }
 
+// TECHDEBT(#595): Properly handle context threading and passing for the entire persistence module
 func (pg *PostgresContext) getCtxAndTx() (context.Context, pgx.Tx) {
 	return context.TODO(), pg.tx
 }

persistence/docs/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.0.0.60] - 2023-07-11
+
+- Adds savepoints and rollbacks implementation to TreeStore
+
 ## [0.0.0.60] - 2023-06-26
 
 - Add place-holder for local context and servicer token usage support methods

persistence/trees/atomic_test.go

@@ -0,0 +1,92 @@
+package trees
+
+import (
+	"encoding/hex"
+	"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"
+
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	// the root hash of a tree store where each tree is empty but present and initialized
+	h0 = "302f2956c084cc3e0e760cf1b8c2da5de79c45fa542f68a660a5fc494b486972"
+	// the root hash of a tree store where each tree has has key foo value bar added to it
+	h1 = "7d5712ea1507915c40e295845fa58773baa405b24b87e9d99761125d826ff915"
+)
+
+func TestTreeStore_AtomicUpdatesWithSuccessfulRollback(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
+	mockTxIndexer := mock_types.NewMockTxIndexer(ctrl)
+	mockBus := mockModules.NewMockBus(ctrl)
+	mockPersistenceMod := mockModules.NewMockPersistenceModule(ctrl)
+
+	mockBus.EXPECT().GetPersistenceModule().AnyTimes().Return(mockPersistenceMod)
+	mockPersistenceMod.EXPECT().GetTxIndexer().AnyTimes().Return(mockTxIndexer)
+
+	ts := &treeStore{
+		logger:       logger.Global.CreateLoggerForModule(modules.TreeStoreSubmoduleName),
+		treeStoreDir: ":memory:",
+	}
+	require.NoError(t, ts.setupTrees())
+	require.NotEmpty(t, ts.merkleTrees[TransactionsTreeName])
+
+	hash0 := ts.getStateHash()
+	require.NotEmpty(t, hash0)
+	require.Equal(t, hash0, h0)
+
+	require.NoError(t, ts.Savepoint())
+
+	// insert test data into every tree
+	for _, treeName := range stateTreeNames {
+		err := ts.merkleTrees[treeName].tree.Update([]byte("foo"), []byte("bar"))
+		require.NoError(t, err)
+	}
+
+	// commit the above changes
+	require.NoError(t, ts.Commit())
+
+	// assert state hash is changed
+	hash1 := ts.getStateHash()
+	require.NotEmpty(t, hash1)
+	require.NotEqual(t, hash0, hash1)
+	require.Equal(t, hash1, h1)
+
+	// set a new savepoint
+	require.NoError(t, ts.Savepoint())
+	require.NotEmpty(t, ts.prevState.merkleTrees)
+	require.NotEmpty(t, ts.prevState.rootTree)
+	// assert that savepoint creation doesn't mutate state hash
+	require.Equal(t, hash1, hex.EncodeToString(ts.prevState.rootTree.tree.Root()))
+
+	// verify that creating a savepoint does not change state hash
+	hash2 := ts.getStateHash()
+	require.Equal(t, hash2, hash1)
+	require.Equal(t, hash2, h1)
+
+	// validate that state tree was updated and a previous savepoint is created
+	for _, treeName := range stateTreeNames {
+		require.NotEmpty(t, ts.merkleTrees[treeName])
+		require.NotEmpty(t, ts.prevState.merkleTrees[treeName])
+	}
+
+	// insert additional test data into all of the trees
+	for _, treeName := range stateTreeNames {
+		require.NoError(t, ts.merkleTrees[treeName].tree.Update([]byte("fiz"), []byte("buz")))
+	}
+
+	// rollback the changes made to the trees above BEFORE anything was committed
+	err := ts.Rollback()
+	require.NoError(t, err)
+
+	// validate that the state hash is unchanged after new data was inserted but rolled back before commitment
+	hash3 := ts.getStateHash()
+	require.Equal(t, hash3, hash2)
+	require.Equal(t, hash3, h1)
+}

persistence/trees/main_test.go

@@ -0,0 +1,12 @@
+//go:build test
+
+package trees
+
+import (
+	"crypto/sha256"
+	"hash"
+)
+
+type TreeStore = treeStore
+
+var SMTTreeHasher hash.Hash = sha256.New()

persistence/trees/module_test.go

@@ -48,7 +48,6 @@ func TestTreeStore_Create(t *testing.T) {
 
 	treemod, err := trees.Create(mockBus, trees.WithTreeStoreDirectory(":memory:"))
 	assert.NoError(t, err)
-
 	got := treemod.GetBus()
 	assert.Equal(t, got, mockBus)
 

persistence/trees/prove_test.go

@@ -0,0 +1,90 @@
+package trees
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/pokt-network/pocket/persistence/kvstore"
+	"github.com/pokt-network/smt"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTreeStore_Prove(t *testing.T) {
+	nodeStore := kvstore.NewMemKVStore()
+	tree := smt.NewSparseMerkleTree(nodeStore, smtTreeHasher)
+	testTree := &stateTree{
+		name:      "test",
+		tree:      tree,
+		nodeStore: nodeStore,
+	}
+
+	require.NoError(t, testTree.tree.Update([]byte("key"), []byte("value")))
+	require.NoError(t, testTree.tree.Commit())
+
+	treeStore := &treeStore{
+		merkleTrees: make(map[string]*stateTree, 1),
+	}
+	treeStore.merkleTrees["test"] = testTree
+
+	testCases := []struct {
+		name        string
+		treeName    string
+		key         []byte
+		value       []byte
+		valid       bool
+		expectedErr error
+	}{
+		{
+			name:        "valid inclusion proof: key and value in tree",
+			treeName:    "test",
+			key:         []byte("key"),
+			value:       []byte("value"),
+			valid:       true,
+			expectedErr: nil,
+		},
+		{
+			name:        "valid exclusion proof: key not in tree",
+			treeName:    "test",
+			key:         []byte("key2"),
+			value:       nil,
+			valid:       true,
+			expectedErr: nil,
+		},
+		{
+			name:        "invalid proof: tree not in store",
+			treeName:    "unstored tree",
+			key:         []byte("key"),
+			value:       []byte("value"),
+			valid:       false,
+			expectedErr: fmt.Errorf("tree not found: %s", "unstored tree"),
+		},
+		{
+			name:        "invalid inclusion proof: key in tree, wrong value",
+			treeName:    "test",
+			key:         []byte("key"),
+			value:       []byte("wrong value"),
+			valid:       false,
+			expectedErr: nil,
+		},
+		{
+			name:        "invalid exclusion proof: key in tree",
+			treeName:    "test",
+			key:         []byte("key"),
+			value:       nil,
+			valid:       false,
+			expectedErr: nil,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			valid, err := treeStore.Prove(tc.treeName, tc.key, tc.value)
+			require.Equal(t, valid, tc.valid)
+			if tc.expectedErr == nil {
+				require.NoError(t, err)
+				return
+			}
+			require.ErrorAs(t, err, &tc.expectedErr)
+		})
+	}
+}