Adds TreeStore module

Author: dylanlott

persistence/context.go

@@ -31,7 +31,7 @@ type PostgresContext struct {
 	//                 Need to simply access them via the bus.
 	blockStore blockstore.BlockStore
 	txIndexer  indexer.TxIndexer
-	stateTrees modules.TreeStore
+	stateTrees modules.TreeStoreModule
 
 	networkId string
 }
@@ -50,7 +50,7 @@ func (p *PostgresContext) RollbackToSavePoint(bytes []byte) error {
 // 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, p.txIndexer, uint64(p.Height))
+	stateHash, err := p.stateTrees.Update(p.tx, uint64(p.Height))
 	if err != nil {
 		return "", err
 	}

persistence/docs/CHANGELOG.md

@@ -7,7 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-## [0.0.0.58] - 2023-06-14
+## [0.0.0.59] - 2023-06-14
+
+- Refactors the persistence treeStore to be an IntegratableModule
+
+## [0.0.0.58] - 2023-06-13
 
 - Refactors the stateTrees implementation off of the PersistenceContext and into its in module
 - Implements the new Update and Commit pattern with the SMT trees in the trees module

persistence/module.go

@@ -18,7 +18,6 @@ import (
 
 var (
 	_ modules.PersistenceModule = &persistenceModule{}
-	_ modules.PersistenceModule = &persistenceModule{}
 
 	_ modules.PersistenceRWContext = &PostgresContext{}
 )
@@ -44,7 +43,7 @@ type persistenceModule struct {
 
 	// stateTrees manages all of the merkle trees maintained by the
 	// persistence module that roll up into the state commitment.
-	stateTrees modules.TreeStore
+	stateTrees modules.TreeStoreModule
 
 	// Only one write context is allowed at a time
 	writeContext *PostgresContext
@@ -104,8 +103,7 @@ func (*persistenceModule) Create(bus modules.Bus, options ...modules.ModuleOptio
 		return nil, err
 	}
 
-	// TECHDEBT (#808): Make TreeStore into a full Module
-	stateTrees, err := trees.NewStateTrees(persistenceCfg.TreesStoreDir)
+	treeModule, err := trees.Create(bus, trees.WithTreeStoreDirectory(persistenceCfg.TreesStoreDir))
 	if err != nil {
 		return nil, err
 	}
@@ -116,7 +114,7 @@ func (*persistenceModule) Create(bus modules.Bus, options ...modules.ModuleOptio
 
 	m.blockStore = blockStore
 	m.txIndexer = txIndexer
-	m.stateTrees = stateTrees
+	m.stateTrees = treeModule
 
 	// TECHDEBT: reconsider if this is the best place to call `populateGenesisState`. Note that
 	// 		     this forces the genesis state to be reloaded on every node startup until state
@@ -237,10 +235,6 @@ func (m *persistenceModule) GetTxIndexer() indexer.TxIndexer {
 	return m.txIndexer
 }
 
-func (m *persistenceModule) GetTreeStore() modules.TreeStore {
-	return m.stateTrees
-}
-
 func (m *persistenceModule) GetNetworkID() string {
 	return m.networkId
 }

persistence/trees/module.go

@@ -0,0 +1,73 @@
+package trees
+
+import (
+	"fmt"
+
+	"github.com/pokt-network/pocket/persistence/kvstore"
+	"github.com/pokt-network/pocket/shared/modules"
+	"github.com/pokt-network/smt"
+)
+
+func (*treeStore) Create(bus modules.Bus, options ...modules.TreeStoreOption) (modules.TreeStoreModule, error) {
+	m := &treeStore{}
+
+	for _, option := range options {
+		option(m)
+	}
+
+	m.SetBus(bus)
+
+	if err := m.setupTrees(); err != nil {
+		return nil, err
+	}
+
+	return m, nil
+}
+
+func Create(bus modules.Bus, options ...modules.TreeStoreOption) (modules.TreeStoreModule, error) {
+	return new(treeStore).Create(bus, options...)
+}
+
+// WithTreeStoreDirectory assigns the path where the tree store
+// saves its data.
+func WithTreeStoreDirectory(path string) modules.TreeStoreOption {
+	return func(m modules.TreeStoreModule) {
+		mod, ok := m.(*treeStore)
+		if ok {
+			mod.treeStoreDir = path
+		}
+	}
+}
+
+func (t *treeStore) setupTrees() error {
+	if t.treeStoreDir == ":memory:" {
+		return t.setupInMemory()
+	}
+
+	t.merkleTrees = make(map[merkleTree]*smt.SMT, int(numMerkleTrees))
+	t.nodeStores = make(map[merkleTree]kvstore.KVStore, int(numMerkleTrees))
+
+	for tree := merkleTree(0); tree < numMerkleTrees; tree++ {
+		nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", t.treeStoreDir, merkleTreeToString[tree]))
+		if err != nil {
+			return err
+		}
+		t.nodeStores[tree] = nodeStore
+		t.merkleTrees[tree] = smt.NewSparseMerkleTree(nodeStore, smtTreeHasher)
+	}
+
+	return nil
+}
+
+func (t *treeStore) setupInMemory() error {
+	t.merkleTrees = make(map[merkleTree]*smt.SMT, int(numMerkleTrees))
+	t.nodeStores = make(map[merkleTree]kvstore.KVStore, int(numMerkleTrees))
+
+	for tree := merkleTree(0); tree < numMerkleTrees; tree++ {
+		nodeStore := kvstore.NewMemKVStore() // For testing, `smt.NewSimpleMap()` can be used as well
+		t.nodeStores[tree] = nodeStore
+		t.merkleTrees[tree] = smt.NewSparseMerkleTree(nodeStore, smtTreeHasher)
+	}
+
+	return nil
+}

persistence/trees/trees.go

@@ -19,9 +19,13 @@ import (
 	"github.com/pokt-network/pocket/shared/codec"
 	coreTypes "github.com/pokt-network/pocket/shared/core/types"
 	"github.com/pokt-network/pocket/shared/crypto"
+	"github.com/pokt-network/pocket/shared/modules"
+	"github.com/pokt-network/pocket/shared/modules/base_modules"
 	"github.com/pokt-network/smt"
 )
 
+// smtTreeHasher sets the hasher used by the tree SMT trees
+// as a package level variable for visibility and internal use.
 var smtTreeHasher hash.Hash = sha256.New()
 
 var merkleTreeToString = map[merkleTree]string{
@@ -56,8 +60,7 @@ type merkleTree float64
 
 // A list of Merkle Trees used to maintain the state hash.
 const (
-	// IMPORTANT: The order in which these trees are defined is important and strict. It implicitly
-	// defines the index of the root hash each independent as they are concatenated together
+	// IMPORTANT: The order in which these trees are defined is important and strict. It implicitly // defines the index of the root hash each independent as they are concatenated together
 	// to generate the state hash.
 
 	// TECHDEBT(#834): Remove the need for enforced ordering
@@ -81,45 +84,29 @@ const (
 	numMerkleTrees
 )
 
+// Ensure treeStore implements TreeStore
+var _ modules.TreeStoreModule = &treeStore{}
+
 // treeStore stores a set of merkle trees that
 // it manages. It fulfills the modules.TreeStore interface.
 // * It is responsible for atomic commit or rollback behavior
 // of the underlying trees by utilizing the lazy loading
 // functionality provided by the underlying smt library.
 type treeStore struct {
+	base_modules.IntegratableModule
+
 	treeStoreDir string
 	merkleTrees  map[merkleTree]*smt.SMT
 	nodeStores   map[merkleTree]kvstore.KVStore
 }
 
-func (t *treeStore) Update(pgtx pgx.Tx, txi indexer.TxIndexer, height uint64) (string, error) {
+// Update takes a transaction and a height and updates
+// all of the trees in the treeStore for that height.
+func (t *treeStore) Update(pgtx pgx.Tx, height uint64) (string, error) {
+	txi := t.GetBus().GetPersistenceModule().GetTxIndexer()
 	return t.updateMerkleTrees(pgtx, txi, height)
 }
 
-// NewStateTrees is the constructor object for a treeStore and initializes and configures a new
-// tree for the appropriate type of store, i.e. in-memory vs file system storage.
-func NewStateTrees(treesStoreDir string) (*treeStore, error) {
-	if treesStoreDir == ":memory:" {
-		return newMemStateTrees()
-	}
-
-	stateTrees := &treeStore{
-		treeStoreDir: treesStoreDir,
-		merkleTrees:  make(map[merkleTree]*smt.SMT, int(numMerkleTrees)),
-		nodeStores:   make(map[merkleTree]kvstore.KVStore, int(numMerkleTrees)),
-	}
-
-	for tree := merkleTree(0); tree < numMerkleTrees; tree++ {
-		nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", treesStoreDir, merkleTreeToString[tree]))
-		if err != nil {
-			return nil, err
-		}
-		stateTrees.nodeStores[tree] = nodeStore
-		stateTrees.merkleTrees[tree] = smt.NewSparseMerkleTree(nodeStore, smtTreeHasher)
-	}
-	return stateTrees, nil
-}
-
 // DebugClearAll is used by the debug cli to completely reset all merkle trees.
 // This should only be called by the debug CLI.
 // TECHDEBT: Move this into a separate file with a debug build flag to avoid accidental usage in prod
@@ -134,20 +121,6 @@ func (t *treeStore) DebugClearAll() error {
 	return nil
 }
 
-// newMemStateTrees creates a new in-memory state tree
-func newMemStateTrees() (*treeStore, error) {
-	stateTrees := &treeStore{
-		merkleTrees: make(map[merkleTree]*smt.SMT, int(numMerkleTrees)),
-		nodeStores:  make(map[merkleTree]kvstore.KVStore, int(numMerkleTrees)),
-	}
-	for tree := merkleTree(0); tree < numMerkleTrees; tree++ {
-		nodeStore := kvstore.NewMemKVStore() // For testing, `smt.NewSimpleMap()` can be used as well
-		stateTrees.nodeStores[tree] = nodeStore
-		stateTrees.merkleTrees[tree] = smt.NewSparseMerkleTree(nodeStore, smtTreeHasher)
-	}
-	return stateTrees, nil
-}
-
 // updateMerkleTrees updates all of the merkle trees in order defined by `numMerkleTrees`
 // * it returns the new state hash capturing the state of all the trees or an error if one occurred
 func (t *treeStore) updateMerkleTrees(pgtx pgx.Tx, txi indexer.TxIndexer, height uint64) (string, error) {

shared/modules/persistence_module.go

@@ -3,7 +3,6 @@ package modules
 //go:generate mockgen -destination=./mocks/persistence_module_mock.go github.com/pokt-network/pocket/shared/modules PersistenceModule,PersistenceRWContext,PersistenceReadContext,PersistenceWriteContext
 
 import (
-	"github.com/jackc/pgx/v5"
 	"github.com/pokt-network/pocket/persistence/blockstore"
 	"github.com/pokt-network/pocket/persistence/indexer"
 	"github.com/pokt-network/pocket/runtime/genesis"
@@ -25,10 +24,6 @@ type PersistenceModule interface {
 	// BlockStore maps a block height to an *coreTypes.IndexedTransaction
 	GetBlockStore() blockstore.BlockStore
 
-	// TreeStore manages atomic access to a set of merkle trees
-	// that compose the state hash.
-	GetTreeStore() TreeStore
-
 	NewWriteContext() PersistenceRWContext
 
 	// Indexer operations
@@ -39,17 +34,6 @@ type PersistenceModule interface {
 	HandleDebugMessage(*messaging.DebugMessage) error
 }
 
-// TreeStore defines the interface for atomic updates and rollbacks to the internal
-// merkle trees that compose the state hash of pocket.
-type TreeStore interface {
-	// Update returns the new state hash for a given height.
-	// * Update inherits the pgx transaction's read view of the database and builds the trees according to that view.
-	// TODO(#808): Change interface to `Update(pgtx pgx.Tx, height uint64) (string, error)`
-	Update(pgtx pgx.Tx, txi indexer.TxIndexer, height uint64) (string, error)
-	// DebugClearAll completely clears the state of the trees. For debugging purposes only.
-	DebugClearAll() error
-}
-
 // Interface defining the context within which the node can operate with the persistence layer.
 // Operations in the context of a PersistenceContext are isolated from other operations and
 // other persistence contexts until committed, enabling parallelizability along other operations.

shared/modules/treestore_module.go

@@ -0,0 +1,30 @@
+package modules
+
+import (
+	"github.com/jackc/pgx/v5"
+)
+
+const (
+	TreeStoreModuleName = "tree_store"
+)
+
+type TreeStoreOption func(TreeStoreModule)
+
+type TreeStoreFactory = FactoryWithOptions[TreeStoreModule, TreeStoreOption]
+
+// TreeStoreModules defines the interface for atomic updates and rollbacks to the internal
+// merkle trees that compose the state hash of pocket.
+type TreeStoreModule interface {
+	IntegratableModule
+
+	// Update returns the new state hash for a given height.
+	// * Height is passed through to the Update function and is used to query the TxIndexer for transactions
+	// to update into the merkle tree set
+	// * Passing a higher height will cause a change but repeatedly calling the same or a lower height will
+	// not incur a change.
+	// * By nature of it taking a pgx transaction at runtime, Update inherits the pgx transaction's read view of the
+	// database.
+	Update(pgtx pgx.Tx, height uint64) (string, error)
+	// DebugClearAll completely clears the state of the trees. For debugging purposes only.
+	DebugClearAll() error
+}