Prevent subtree data race

Author: oskarszoon

services/blockassembly/subtreeprocessor/SubtreeProcessor.go

@@ -207,7 +207,8 @@ type SubtreeProcessor struct {
 	chainedSubtreeCount atomic.Int32
 
 	// currentSubtree represents the subtree currently being built
-	currentSubtree *subtreepkg.Subtree
+	// Uses atomic.Pointer for safe concurrent access from external callers (e.g., gRPC handlers)
+	currentSubtree atomic.Pointer[subtreepkg.Subtree]
 
 	// currentBlockHeader stores the current block header being processed
 	currentBlockHeader *model.BlockHeader
@@ -384,7 +385,6 @@ func NewSubtreeProcessor(ctx context.Context, logger ulogger.Logger, tSettings *
 		newSubtreeChan:           newSubtreeChan,
 		chainedSubtrees:          make([]*subtreepkg.Subtree, 0, ExpectedNumberOfSubtrees),
 		chainedSubtreeCount:      atomic.Int32{},
-		currentSubtree:           firstSubtree,
 		batcher:                  NewTxIDAndFeeBatch(tSettings.BlockAssembly.SubtreeProcessorBatcherSize),
 		queue:                    queue,
 		currentTxMap:             txmap.NewSyncedMap[chainhash.Hash, subtreepkg.TxInpoints](),
@@ -397,6 +397,7 @@ func NewSubtreeProcessor(ctx context.Context, logger ulogger.Logger, tSettings *
 		currentRunningState:      atomic.Value{},
 		announcementTicker:       time.NewTicker(tSettings.BlockAssembly.SubtreeAnnouncementInterval),
 	}
+	stp.currentSubtree.Store(firstSubtree)
 	stp.setCurrentRunningState(StateStarting)
 
 	// need to make sure first coinbase tx is counted when we start
@@ -457,7 +458,7 @@ func (stp *SubtreeProcessor) Start(ctx context.Context) {
 					completeSubtrees = append(completeSubtrees, stp.chainedSubtrees...)
 
 					// incomplete subtrees ?
-					if chainedCount == 0 && stp.currentSubtree.Length() > 1 {
+					if chainedCount == 0 && stp.currentSubtree.Load().Length() > 1 {
 						incompleteSubtree, err := stp.createIncompleteSubtreeCopy()
 						if err != nil {
 							logger.Errorf("[SubtreeProcessor] error creating incomplete subtree: %s", err.Error())
@@ -513,8 +514,8 @@ func (stp *SubtreeProcessor) Start(ctx context.Context) {
 						subtreeHashes = append(subtreeHashes, *subtree.RootHash())
 					}
 
-					if stp.currentSubtree.Length() > 0 {
-						subtreeHashes = append(subtreeHashes, *stp.currentSubtree.RootHash())
+					if stp.currentSubtree.Load().Length() > 0 {
+						subtreeHashes = append(subtreeHashes, *stp.currentSubtree.Load().RootHash())
 					}
 
 					getSubtreeHashesChan <- subtreeHashes
@@ -530,8 +531,8 @@ func (stp *SubtreeProcessor) Start(ctx context.Context) {
 							transactionHashes = append(transactionHashes, node.Hash)
 						}
 					}
-					if stp.currentSubtree.Length() > 0 {
-						for _, node := range stp.currentSubtree.Nodes {
+					if stp.currentSubtree.Load().Length() > 0 {
+						for _, node := range stp.currentSubtree.Load().Nodes {
 							transactionHashes = append(transactionHashes, node.Hash)
 						}
 					}
@@ -560,14 +561,14 @@ func (stp *SubtreeProcessor) Start(ctx context.Context) {
 
 					// store current state before attempting to move forward the block
 					originalChainedSubtrees := stp.chainedSubtrees
-					originalCurrentSubtree := stp.currentSubtree
+					originalCurrentSubtree := stp.currentSubtree.Load()
 					originalCurrentTxMap := stp.currentTxMap
 					currentBlockHeader := stp.currentBlockHeader
 
 					if _, err = stp.moveForwardBlock(processorCtx, moveForwardReq.block, false, processedConflictingHashesMap, false, true); err != nil {
 						// rollback to previous state
 						stp.chainedSubtrees = originalChainedSubtrees
-						stp.currentSubtree = originalCurrentSubtree
+						stp.currentSubtree.Store(originalCurrentSubtree)
 						stp.currentTxMap = originalCurrentTxMap
 						stp.currentBlockHeader = currentBlockHeader
 
@@ -608,7 +609,7 @@ func (stp *SubtreeProcessor) Start(ctx context.Context) {
 
 				case lengthCh := <-stp.lengthCh:
 					// return the length of the current subtree
-					lengthCh <- stp.currentSubtree.Length()
+					lengthCh <- stp.currentSubtree.Load().Length()
 
 				case errCh := <-stp.checkSubtreeProcessorCh:
 					stp.setCurrentRunningState(StateCheckSubtreeProcessor)
@@ -619,8 +620,8 @@ func (stp *SubtreeProcessor) Start(ctx context.Context) {
 
 				case <-stp.announcementTicker.C:
 					// Periodically announce the current subtree if it has transactions
-					if stp.currentSubtree.Length() > 1 {
-						logger.Debugf("[SubtreeProcessor] periodic announcement of current subtree with %d transactions", stp.currentSubtree.Length()-1)
+					if stp.currentSubtree.Load().Length() > 1 {
+						logger.Debugf("[SubtreeProcessor] periodic announcement of current subtree with %d transactions", stp.currentSubtree.Load().Length()-1)
 
 						incompleteSubtree, err := stp.createIncompleteSubtreeCopy()
 						if err != nil {
@@ -757,14 +758,14 @@ func (stp *SubtreeProcessor) createIncompleteSubtreeCopy() (*subtreepkg.Subtree,
 	}
 
 	// Copy all nodes from current subtree (skipping the coinbase placeholder at index 0)
-	for _, node := range stp.currentSubtree.Nodes[1:] {
+	for _, node := range stp.currentSubtree.Load().Nodes[1:] {
 		if err = incompleteSubtree.AddSubtreeNode(node); err != nil {
 			return nil, err
 		}
 	}
 
 	// Copy fees
-	incompleteSubtree.Fees = stp.currentSubtree.Fees
+	incompleteSubtree.Fees = stp.currentSubtree.Load().Fees
 
 	return incompleteSubtree, nil
 }
@@ -827,8 +828,9 @@ func (stp *SubtreeProcessor) reset(blockHeader *model.BlockHeader, moveBackBlock
 	stp.chainedSubtrees = make([]*subtreepkg.Subtree, 0, ExpectedNumberOfSubtrees)
 	stp.chainedSubtreeCount.Store(0)
 
-	stp.currentSubtree, _ = subtreepkg.NewTreeByLeafCount(stp.currentItemsPerFile)
-	if err := stp.currentSubtree.AddCoinbaseNode(); err != nil {
+	newSubtree, _ := subtreepkg.NewTreeByLeafCount(stp.currentItemsPerFile)
+	stp.currentSubtree.Store(newSubtree)
+	if err := stp.currentSubtree.Load().AddCoinbaseNode(); err != nil {
 		return errors.NewProcessingError("[SubtreeProcessor][Reset] error adding coinbase placeholder to new current subtree", err)
 	}
 
@@ -1042,11 +1044,12 @@ func (stp *SubtreeProcessor) SetCurrentBlockHeader(blockHeader *model.BlockHeade
 }
 
 // GetCurrentSubtree returns the subtree currently being built.
+// This method is safe for concurrent access.
 //
 // Returns:
 //   - *util.Subtree: Current subtree
 func (stp *SubtreeProcessor) GetCurrentSubtree() *subtreepkg.Subtree {
-	return stp.currentSubtree
+	return stp.currentSubtree.Load()
 }
 
 // GetCurrentTxMap returns the map of transactions currently held in the subtree processor.
@@ -1348,27 +1351,28 @@ func (stp *SubtreeProcessor) addNode(node subtreepkg.Node, parents *subtreepkg.T
 		}
 	}
 
-	if stp.currentSubtree == nil {
-		stp.currentSubtree, err = subtreepkg.NewTreeByLeafCount(stp.currentItemsPerFile)
+	if stp.currentSubtree.Load() == nil {
+		newSubtree, err := subtreepkg.NewTreeByLeafCount(stp.currentItemsPerFile)
 		if err != nil {
 			return err
 		}
+		stp.currentSubtree.Store(newSubtree)
 
 		// This is the first subtree for this block - we need a coinbase placeholder
-		err = stp.currentSubtree.AddCoinbaseNode()
+		err = stp.currentSubtree.Load().AddCoinbaseNode()
 		if err != nil {
 			return err
 		}
 
 		stp.txCount.Add(1)
 	}
 
-	err = stp.currentSubtree.AddSubtreeNode(node)
+	err = stp.currentSubtree.Load().AddSubtreeNode(node)
 	if err != nil {
 		return errors.NewProcessingError("error adding node to subtree", err)
 	}
 
-	if stp.currentSubtree.IsComplete() {
+	if stp.currentSubtree.Load().IsComplete() {
 		if err = stp.processCompleteSubtree(skipNotification); err != nil {
 			return err
 		}
@@ -1378,36 +1382,38 @@ func (stp *SubtreeProcessor) addNode(node subtreepkg.Node, parents *subtreepkg.T
 }
 
 func (stp *SubtreeProcessor) processCompleteSubtree(skipNotification bool) (err error) {
+	currentSubtree := stp.currentSubtree.Load()
 	if !skipNotification {
-		stp.logger.Debugf("[%s] append subtree", stp.currentSubtree.RootHash().String())
+		stp.logger.Debugf("[%s] append subtree", currentSubtree.RootHash().String())
 	}
 
 	// Track the actual number of nodes in this subtree
 	// We don't exclude coinbase because:
 	// 1. Only the first subtree in a block has a coinbase
 	// 2. The coinbase is still a transaction that takes space
 	// 3. For sizing decisions, we care about total throughput
-	actualNodeCount := len(stp.currentSubtree.Nodes)
+	actualNodeCount := len(currentSubtree.Nodes)
 	if actualNodeCount > 0 {
 		// Add to ring buffer (overwrites oldest value automatically)
 		stp.subtreeNodeCounts.Value = actualNodeCount
 		stp.subtreeNodeCounts = stp.subtreeNodeCounts.Next()
 	}
 
 	// Add the subtree to the chain
-	stp.chainedSubtrees = append(stp.chainedSubtrees, stp.currentSubtree)
+	stp.chainedSubtrees = append(stp.chainedSubtrees, currentSubtree)
 	stp.chainedSubtreeCount.Add(1)
 
 	stp.subtreesInBlock++ // Track number of subtrees in current block
 
-	oldSubtree := stp.currentSubtree
+	oldSubtree := currentSubtree
 	oldSubtreeHash := oldSubtree.RootHash()
 
 	// create a new subtree with the same height as the previous subtree
-	stp.currentSubtree, err = subtreepkg.NewTree(stp.currentSubtree.Height)
+	newSubtree, err := subtreepkg.NewTree(oldSubtree.Height)
 	if err != nil {
 		return errors.NewProcessingError("[%s] error creating new subtree", oldSubtreeHash.String(), err)
 	}
+	stp.currentSubtree.Store(newSubtree)
 
 	// Send the subtree to the newSubtreeChan, including a reference to the parent transactions map
 	errCh := make(chan error)
@@ -1506,7 +1512,7 @@ func (stp *SubtreeProcessor) removeTxFromSubtrees(ctx context.Context, hash chai
 	defer deferFn()
 
 	// find the transaction in the current and all chained subtrees
-	foundIndex := stp.currentSubtree.NodeIndex(hash)
+	foundIndex := stp.currentSubtree.Load().NodeIndex(hash)
 	foundSubtreeIndex := -1
 
 	if foundIndex == -1 {
@@ -1528,7 +1534,7 @@ func (stp *SubtreeProcessor) removeTxFromSubtrees(ctx context.Context, hash chai
 		if foundSubtreeIndex == -1 {
 			// it was found in the current tree, remove it from there
 			// further processing is not needed, as the subtrees in the chainedSubtrees are older than the current subtree
-			return stp.currentSubtree.RemoveNodeAtIndex(foundIndex)
+			return stp.currentSubtree.Load().RemoveNodeAtIndex(foundIndex)
 		}
 
 		// it was found in a chained subtree, remove it from there and chain the subtrees again from the point it was removed
@@ -1570,7 +1576,7 @@ func (stp *SubtreeProcessor) removeTxsFromSubtrees(ctx context.Context, hashes [
 
 	for _, hash := range hashes {
 		// find the transaction in the current and all chained subtrees
-		foundIndex := stp.currentSubtree.NodeIndex(hash)
+		foundIndex := stp.currentSubtree.Load().NodeIndex(hash)
 		foundSubtreeIndex := -1
 
 		if foundIndex == -1 {
@@ -1592,7 +1598,7 @@ func (stp *SubtreeProcessor) removeTxsFromSubtrees(ctx context.Context, hashes [
 			if foundSubtreeIndex == -1 {
 				// it was found in the current tree, remove it from there
 				// further processing is not needed, as the subtrees in the chainedSubtrees are older than the current subtree
-				return stp.currentSubtree.RemoveNodeAtIndex(foundIndex)
+				return stp.currentSubtree.Load().RemoveNodeAtIndex(foundIndex)
 			}
 
 			// it was found in a chained subtree, remove it from there and chain the subtrees again from the point it was removed
@@ -1625,7 +1631,7 @@ func (stp *SubtreeProcessor) removeTxsFromSubtrees(ctx context.Context, hashes [
 func (stp *SubtreeProcessor) reChainSubtrees(fromIndex int) error {
 	// copy the original subtrees from the given index into a new structure
 	originalSubtrees := stp.chainedSubtrees[fromIndex:]
-	originalSubtrees = append(originalSubtrees, stp.currentSubtree)
+	originalSubtrees = append(originalSubtrees, stp.currentSubtree.Load())
 
 	// reset the chained subtrees and the current subtree
 	stp.chainedSubtrees = stp.chainedSubtrees[:fromIndex]
@@ -1637,14 +1643,15 @@ func (stp *SubtreeProcessor) reChainSubtrees(fromIndex int) error {
 
 	stp.chainedSubtreeCount.Store(fromIndexInt32)
 
-	stp.currentSubtree, err = subtreepkg.NewTreeByLeafCount(stp.currentItemsPerFile)
+	newSubtree, err := subtreepkg.NewTreeByLeafCount(stp.currentItemsPerFile)
 	if err != nil {
 		return errors.NewProcessingError("error creating new current subtree", err)
 	}
+	stp.currentSubtree.Store(newSubtree)
 
 	if len(originalSubtrees) == 0 {
 		// we must add the coinbase tx if we have no original subtrees
-		if err = stp.currentSubtree.AddCoinbaseNode(); err != nil {
+		if err = stp.currentSubtree.Load().AddCoinbaseNode(); err != nil {
 			return errors.NewProcessingError("error adding coinbase node to new current subtree", err)
 		}
 	}
@@ -1724,7 +1731,7 @@ func (stp *SubtreeProcessor) checkSubtreeProcessor(errCh chan error) {
 	}
 
 	// check all the transactions in the subtrees are in the currentTxMap
-	for nodeIdx, node := range stp.currentSubtree.Nodes {
+	for nodeIdx, node := range stp.currentSubtree.Load().Nodes {
 		if _, ok := stp.currentTxMap.Get(node.Hash); !ok {
 			if node.Hash.Equal(subtreepkg.CoinbasePlaceholderHashValue) {
 				// check that the coinbase placeholder is in the first subtree
@@ -1880,15 +1887,15 @@ func (stp *SubtreeProcessor) reorgBlocks(ctx context.Context, moveBackBlocks []*
 
 	// store current state for restore in case of error
 	originalChainedSubtrees := stp.chainedSubtrees
-	originalCurrentSubtree := stp.currentSubtree
+	originalCurrentSubtree := stp.currentSubtree.Load()
 	originalCurrentTxMap := stp.currentTxMap
 	currentBlockHeader := stp.currentBlockHeader
 
 	defer func() {
 		if err != nil {
 			// restore the original state
 			stp.chainedSubtrees = originalChainedSubtrees
-			stp.currentSubtree = originalCurrentSubtree
+			stp.currentSubtree.Store(originalCurrentSubtree)
 			stp.currentTxMap = originalCurrentTxMap
 			stp.currentBlockHeader = currentBlockHeader
 
@@ -2007,9 +2014,10 @@ func (stp *SubtreeProcessor) reorgBlocks(ctx context.Context, moveBackBlocks []*
 	}
 
 	// also for the current subtree
-	notOnLongestChain := make([]chainhash.Hash, 0, len(stp.currentSubtree.Nodes))
+	currentSubtree := stp.currentSubtree.Load()
+	notOnLongestChain := make([]chainhash.Hash, 0, len(currentSubtree.Nodes))
 
-	for _, node := range stp.currentSubtree.Nodes {
+	for _, node := range currentSubtree.Nodes {
 		if node.Hash.Equal(subtreepkg.CoinbasePlaceholderHashValue) {
 			// skip coinbase placeholder
 			continue
@@ -2182,7 +2190,7 @@ func (stp *SubtreeProcessor) setTxCountFromSubtrees() {
 		stp.txCount.Add(subtreeLen)
 	}
 
-	currSubtreeLenUint64, err := safeconversion.IntToUint64(stp.currentSubtree.Length())
+	currSubtreeLenUint64, err := safeconversion.IntToUint64(stp.currentSubtree.Load().Length())
 	if err != nil {
 		stp.logger.Errorf("error converting current subtree length: %s", err)
 		return
@@ -2220,7 +2228,7 @@ func (stp *SubtreeProcessor) moveBackBlock(ctx context.Context, block *model.Blo
 		deferFn()
 	}()
 
-	lastIncompleteSubtree := stp.currentSubtree
+	lastIncompleteSubtree := stp.currentSubtree.Load()
 	chainedSubtrees := stp.chainedSubtrees
 
 	// process coinbase utxos
@@ -2308,16 +2316,17 @@ func (stp *SubtreeProcessor) moveBackBlockCreateNewSubtrees(ctx context.Context,
 		// we create as few subtrees as possible when moving back, to avoid fragmentation and lots of small writes to disk
 		subtreeSize = 1024 * 1024
 	}
-	stp.currentSubtree, err = subtreepkg.NewTreeByLeafCount(subtreeSize)
+	newSubtree, err := subtreepkg.NewTreeByLeafCount(subtreeSize)
 	if err != nil {
 		return nil, nil, errors.NewProcessingError("[moveBackBlock:CreateNewSubtrees][%s] error creating new subtree", block.String(), err)
 	}
+	stp.currentSubtree.Store(newSubtree)
 
 	stp.chainedSubtrees = make([]*subtreepkg.Subtree, 0, ExpectedNumberOfSubtrees)
 	stp.chainedSubtreeCount.Store(0)
 
 	// add first coinbase placeholder transaction
-	_ = stp.currentSubtree.AddCoinbaseNode()
+	_ = stp.currentSubtree.Load().AddCoinbaseNode()
 
 	// run through the nodes of the subtrees in order and add to the new subtrees
 	if len(subtreesNodes) > 0 {
@@ -2576,16 +2585,17 @@ func (stp *SubtreeProcessor) resetSubtreeState(createProperlySizedSubtrees bool)
 		subtreeSize = 1024 * 1024
 	}
 
-	stp.currentSubtree, err = subtreepkg.NewTreeByLeafCount(subtreeSize)
+	newSubtree, err := subtreepkg.NewTreeByLeafCount(subtreeSize)
 	if err != nil {
 		return err
 	}
+	stp.currentSubtree.Store(newSubtree)
 
 	stp.chainedSubtrees = make([]*subtreepkg.Subtree, 0, ExpectedNumberOfSubtrees)
 	stp.chainedSubtreeCount.Store(0)
 
 	// Add first coinbase placeholder transaction
-	_ = stp.currentSubtree.AddCoinbaseNode()
+	_ = stp.currentSubtree.Load().AddCoinbaseNode()
 
 	return nil
 }
@@ -2762,7 +2772,7 @@ func (stp *SubtreeProcessor) moveForwardBlock(ctx context.Context, block *model.
 		return nil, err
 	}
 
-	originalCurrentSubtree := stp.currentSubtree
+	originalCurrentSubtree := stp.currentSubtree.Load()
 	originalCurrentTxMap := stp.currentTxMap
 
 	// Reset subtree state