fix: smt fixes and memory efficiencies

Author: pablocampogo

controller/block.go

@@ -277,12 +277,15 @@ func (c *Controller) CommitCertificate(qc *lib.QuorumCertificate, block *lib.Blo
 		// exit with error
 		return
 	}
-	// delete each transaction from the mempool
-	c.Mempool.DeleteTransaction(block.Transactions...)
+	syncing := c.isSyncing.Load()
+	if !syncing {
+		// delete each transaction from the mempool
+		c.Mempool.DeleteTransaction(block.Transactions...)
+	}
 	// parse committed block for straw polls
 	c.FSM.ParsePollTransactions(blockResult)
 	// if self was the proposer
-	if bytes.Equal(qc.ProposerKey, c.PublicKey) && !c.isSyncing.Load() {
+	if bytes.Equal(qc.ProposerKey, c.PublicKey) && !syncing {
 		// send the certificate results transaction on behalf of the quorum
 		c.SendCertificateResultsTx(qc)
 	}
@@ -301,35 +304,38 @@ func (c *Controller) CommitCertificate(qc *lib.QuorumCertificate, block *lib.Blo
 		// exit with error
 		return err
 	}
-	// reset the current mempool store to prepare for the next height
-	c.Mempool.FSM.Discard()
-	// set up the mempool with the actual new FSM for the next height
-	// this makes c.Mempool.FSM.Reset() is unnecessary
-	if c.Mempool.FSM, err = c.FSM.Copy(); err != nil {
-		// exit with error
-		return err
+	if !syncing {
+		// reset the current mempool store to prepare for the next height
+		c.Mempool.FSM.Discard()
+		// set up the mempool with the actual new FSM for the next height
+		if c.Mempool.FSM, err = c.FSM.Copy(); err != nil {
+			// exit with error
+			return err
+		}
+		// check the mempool to cache a proposal block and validate the mempool itself
+		c.Mempool.CheckMempool()
+		// reset mempool FSM
+		c.Mempool.FSM.Reset()
 	}
-	// check the mempool to cache a proposal block and validate the mempool itself
-	c.Mempool.CheckMempool()
-	// reset mempool FSM
-	c.Mempool.FSM.Reset()
 	// update telemetry (using proper defer to ensure time.Since is evaluated at defer execution)
 	defer c.UpdateTelemetry(qc, block, time.Since(start))
-	// publish root chain information to all nested chain subscribers.
-	for _, id := range c.RCManager.ChainIds() {
-		// get the root chain info
-		info, e := c.FSM.LoadRootChainInfo(id, 0)
-		if e != nil {
-			// don't log 'no-validators' error as this is possible
-			if e.Error() != lib.ErrNoValidators().Error() {
-				c.log.Error(e.Error())
+	if !syncing {
+		// publish root chain information to all nested chain subscribers
+		for _, id := range c.RCManager.ChainIds() {
+			// get the root chain info
+			info, e := c.FSM.LoadRootChainInfo(id, 0)
+			if e != nil {
+				// don't log 'no-validators' error as this is possible
+				if e.Error() != lib.ErrNoValidators().Error() {
+					c.log.Error(e.Error())
+				}
+				continue
 			}
-			continue
+			// set the timestamp
+			info.Timestamp = ts
+			// publish root chain information
+			go c.RCManager.Publish(id, info)
 		}
-		// set the timestamp
-		info.Timestamp = ts
-		// publish root chain information
-		go c.RCManager.Publish(id, info)
 	}
 	// exit
 	return

controller/consensus.go

@@ -740,6 +740,16 @@ func (c *Controller) finishSyncing() {
 	c.Lock()
 	// when function completes, unlock
 	defer c.Unlock()
+	// reinitialize the mempool now that sync is complete
+	c.Mempool.L.Lock()
+	c.Mempool.Clear()
+	c.Mempool.FSM.Discard()
+	if mFSM, err := c.FSM.Copy(); err == nil {
+		c.Mempool.FSM = mFSM
+	}
+	c.Mempool.CheckMempool()
+	c.Mempool.FSM.Reset()
+	c.Mempool.L.Unlock()
 	// set the startup block metric (block height when first sync completed)
 	c.Metrics.SetStartupBlock(c.FSM.Height())
 	// signal a reset of bft for the chain

store/smt.go

@@ -190,79 +190,93 @@ func (s *SMT) Commit(unsortedOps map[uint64]valueOp) (err lib.ErrorI) {
 
 // CommitParallel() executes deferred operations in parallel by partitioning them into
 // 8 subtrees based on their 3-bit prefix (000-111), avoiding conflicts between operations
-// that would modify overlapping tree regions. Each subtree is processed independently,
-// then the results are merged back into the main tree.
+// that would modify overlapping tree regions. Each subtree is processed independently
+// with its own Txn copy to avoid lock contention, then the results are merged back into
+// the main tree.
 func (s *SMT) CommitParallel(unsortedOps map[uint64]valueOp) (err lib.ErrorI) {
 	// if there are too few operations, fall back to sequential processing
 	// if len(unsortedOps) < NumSubtrees*2 {
 	// 	return s.Commit(unsortedOps)
 	// }
 
-	// sort operations by their 3-bit prefix to avoid conflicts
+	parentTxn, ok := s.store.(*Txn)
+	if !ok {
+		return s.Commit(unsortedOps)
+	}
+
 	groups, err := s.sortOperationsByPrefix(unsortedOps)
 	if err != nil {
 		return err
 	}
 
-	// add synthetic borders to enable safe parallel processing
 	cleanup, err := s.addSyntheticBorders()
 	if err != nil {
 		return err
 	}
 	defer func() {
-		// cleanup synthetic borders regardless of success/failure
 		if cleanupErr := cleanup(); cleanupErr != nil && err == nil {
 			err = cleanupErr
 		}
 	}()
 
-	// get subtree roots for parallel processing
 	subtreeRoots, err := s.getSubtreeRoots()
 	if err != nil {
 		return err
 	}
 
-	// process subtrees in parallel using goroutines
 	type subtreeResult struct {
 		index int
+		store *subtreeStore
 		err   lib.ErrorI
 	}
 
 	resultChan := make(chan subtreeResult, NumSubtrees)
 	activeSubtrees := 0
 
-	// launch goroutines for each subtree that has operations
 	for i := 0; i < NumSubtrees; i++ {
 		if len(groups[i]) == 0 {
-			continue // skip empty groups
+			continue
 		}
-
 		activeSubtrees++
-		go func(idx int, ops []*node, root *node) {
-			// create an isolated subtree for this prefix
-			subtree := s.createSubtree(root, ops)
-			// reset subtree state for processing
+		st := parentTxn.newSubtreeStore()
+
+		go func(idx int, ops []*node, root *node, store *subtreeStore) {
+			subtree := &SMT{
+				store:        store,
+				root:         root,
+				keyBitLength: s.keyBitLength,
+				nodeCache:    make(map[string]*node),
+				operations:   ops,
+				minKey:       s.minKey,
+				maxKey:       s.maxKey,
+			}
 			subtree.reset()
-			// process operations in this subtree
-			err := subtree.commit(true)
-			// send result back
-			resultChan <- subtreeResult{index: idx, err: err}
-		}(i, groups[i], subtreeRoots[i])
+			commitErr := subtree.commit(true)
+			resultChan <- subtreeResult{
+				index: idx,
+				store: store,
+				err:   commitErr,
+			}
+		}(i, groups[i], subtreeRoots[i], st)
 	}
 
-	// collect results from all active subtrees
+	results := make([]subtreeResult, 0, activeSubtrees)
 	for completed := 0; completed < activeSubtrees; completed++ {
 		result := <-resultChan
 		if result.err != nil {
-			// if any subtree fails, we need to return the error
-			// the cleanup function will handle synthetic border removal
+			// Drain remaining results before returning so goroutines
+			// finish and don't race with the deferred cleanup.
+			for completed++; completed < activeSubtrees; completed++ {
+				<-resultChan
+			}
 			return result.err
 		}
+		results = append(results, result)
+	}
+	for _, result := range results {
+		parentTxn.mergeSubtreeOps(result.store)
 	}
 
-	// after all subtrees are processed, the tree state is already consistent
-	// because each subtree operation updated the shared store
-	// we just need to refresh our in-memory view of the root
 	s.root, err = s.getNode(s.root.Key.bytes())
 	if err != nil {
 		return err
@@ -547,19 +561,6 @@ func (s *SMT) getSubtreeRoots() (roots []*node, err lib.ErrorI) {
 	return
 }
 
-// createSubtree() initializes the subtree structure
-func (s *SMT) createSubtree(root *node, operations []*node) *SMT {
-	return &SMT{
-		store:        s.store,
-		root:         root,
-		keyBitLength: s.keyBitLength,
-		nodeCache:    make(map[string]*node),
-		operations:   operations,
-		minKey:       s.minKey,
-		maxKey:       s.maxKey,
-	}
-}
-
 // sortOperationsByPrefix returns 8 sorted slices grouped by 3-bit prefix: 000 to 111
 func (s *SMT) sortOperationsByPrefix(unsortedOps map[uint64]valueOp) (groups [8][]*node, err lib.ErrorI) {
 	// for each unsorted operation

store/txn.go

@@ -2,6 +2,7 @@ package store
 
 import (
 	"bytes"
+	"fmt"
 	"sync"
 
 	"github.com/canopy-network/canopy/lib"
@@ -274,6 +275,100 @@ func (t *Txn) NewIterator(prefix []byte, reverse bool, seek bool) (lib.IteratorI
 	return newTxnIterator(parentIterator, t.txn, prefix, t.prefix, reverse), nil
 }
 
+// newSubtreeStore creates a lightweight store for a parallel SMT subtree.
+// Reads check: (1) the subtree's own ops, (2) the parent Txn's ops, then
+// (3) an independent PebbleDB reader. Writes go only to the subtree's ops.
+// After the goroutine completes, call mergeSubtreeOps to fold the subtree
+// writes back into the parent Txn.
+func (t *Txn) newSubtreeStore() *subtreeStore {
+	var reader *VersionedStore
+	if vs, ok := t.reader.(*VersionedStore); ok {
+		reader = vs.NewParallelReader()
+	}
+	t.txn.l.Lock()
+	parentOps := make(map[uint64]valueOp, len(t.txn.ops))
+	for k, v := range t.txn.ops {
+		parentOps[k] = v
+	}
+	t.txn.l.Unlock()
+	return &subtreeStore{
+		ops:       make(map[uint64]valueOp),
+		parentOps: parentOps,
+		reader:    reader,
+		prefix:    t.prefix,
+		version:   t.writeVersion,
+	}
+}
+
+// mergeSubtreeOps copies all operations from a subtreeStore into the Txn.
+// Must only be called when no goroutines are accessing the subtreeStore.
+func (t *Txn) mergeSubtreeOps(s *subtreeStore) {
+	for k, v := range s.ops {
+		if _, exists := t.txn.ops[k]; !exists && t.sort {
+			t.addToSorted(v.key, k)
+		}
+		t.txn.ops[k] = v
+	}
+}
+
+// subtreeStore implements lib.RWStoreI for a parallel SMT subtree. It
+// provides three-layer reads (own ops → parent ops snapshot → PebbleDB) with
+// independent snapshot access to avoid lock contention on the DB path.
+// The parent ops are snapshotted at creation time so reads are lock-free.
+type subtreeStore struct {
+	l         sync.Mutex
+	ops       map[uint64]valueOp
+	parentOps map[uint64]valueOp
+	reader    *VersionedStore
+	prefix    []byte
+	version   uint64
+}
+
+func (s *subtreeStore) Get(key []byte) ([]byte, lib.ErrorI) {
+	h := lib.MemHash(key)
+	s.l.Lock()
+	if v, found := s.ops[h]; found {
+		s.l.Unlock()
+		if v.op == opDelete {
+			return nil, nil
+		}
+		return v.value, nil
+	}
+	s.l.Unlock()
+	if v, found := s.parentOps[h]; found {
+		if v.op == opDelete {
+			return nil, nil
+		}
+		return v.value, nil
+	}
+	if s.reader == nil {
+		return nil, nil
+	}
+	return s.reader.Get(lib.Append(s.prefix, key))
+}
+
+func (s *subtreeStore) Set(key, value []byte) lib.ErrorI {
+	s.l.Lock()
+	s.ops[lib.MemHash(key)] = valueOp{key: key, value: value, version: s.version, op: opSet}
+	s.l.Unlock()
+	return nil
+}
+
+func (s *subtreeStore) Delete(key []byte) lib.ErrorI {
+	s.l.Lock()
+	s.ops[lib.MemHash(key)] = valueOp{key: key, version: s.version, op: opDelete}
+	s.l.Unlock()
+	return nil
+}
+
+func (s *subtreeStore) Iterator([]byte) (lib.IteratorI, lib.ErrorI) {
+	return nil, ErrStoreGet(fmt.Errorf("iterator not supported on subtree store"))
+}
+
+func (s *subtreeStore) RevIterator([]byte) (lib.IteratorI, lib.ErrorI) {
+	return nil, ErrStoreGet(fmt.Errorf("reverse iterator not supported on subtree store"))
+}
+
 // Copy creates a new Txn with the same configuration and txn as the original
 func (t *Txn) Copy(reader TxnReaderI, writer TxnWriterI) *Txn {
 	return &Txn{

store/versioned_store.go

@@ -77,6 +77,18 @@ func NewVersionedStore(db pebble.Reader, batch *pebble.Batch, version uint64) *V
 	}
 }
 
+// NewParallelReader creates a read-only VersionedStore sharing the same
+// underlying database (snapshot) but with its own buffers, safe for concurrent
+// use from a separate goroutine. The returned store must NOT be closed, as it
+// does not own the underlying reader.
+func (vs *VersionedStore) NewParallelReader() *VersionedStore {
+	return &VersionedStore{
+		db:           vs.db,
+		version:      vs.version,
+		decodeBuffer: make([][]byte, 0, 5),
+	}
+}
+
 // Set() stores a key-value pair at the current version
 func (vs *VersionedStore) Set(key, value []byte) (err lib.ErrorI) {
 	return vs.SetAt(key, value, vs.version)