Prevent subtree data race

Author: oskarszoon

services/blockassembly/subtreeprocessor/SubtreeProcessor.go

@@ -209,6 +209,9 @@ type SubtreeProcessor struct {
 	// currentSubtree represents the subtree currently being built
 	currentSubtree *subtreepkg.Subtree
 
+	// currentSubtreeMu protects currentSubtree pointer from concurrent read/write access
+	currentSubtreeMu sync.RWMutex
+
 	// currentBlockHeader stores the current block header being processed
 	currentBlockHeader *model.BlockHeader
 
@@ -567,7 +570,7 @@ func (stp *SubtreeProcessor) Start(ctx context.Context) {
 					if _, err = stp.moveForwardBlock(processorCtx, moveForwardReq.block, false, processedConflictingHashesMap, false, true); err != nil {
 						// rollback to previous state
 						stp.chainedSubtrees = originalChainedSubtrees
-						stp.currentSubtree = originalCurrentSubtree
+						stp.setCurrentSubtree(originalCurrentSubtree)
 						stp.currentTxMap = originalCurrentTxMap
 						stp.currentBlockHeader = currentBlockHeader
 
@@ -827,7 +830,8 @@ 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)
+	newSubtree, _ := subtreepkg.NewTreeByLeafCount(stp.currentItemsPerFile)
+	stp.setCurrentSubtree(newSubtree)
 	if err := stp.currentSubtree.AddCoinbaseNode(); err != nil {
 		return errors.NewProcessingError("[SubtreeProcessor][Reset] error adding coinbase placeholder to new current subtree", err)
 	}
@@ -1041,11 +1045,22 @@ func (stp *SubtreeProcessor) SetCurrentBlockHeader(blockHeader *model.BlockHeade
 	stp.currentBlockHeader = blockHeader
 }
 
+// setCurrentSubtree safely updates the currentSubtree pointer with write lock protection.
+// This is used internally to prevent data races when external callers read via GetCurrentSubtree().
+func (stp *SubtreeProcessor) setCurrentSubtree(st *subtreepkg.Subtree) {
+	stp.currentSubtreeMu.Lock()
+	stp.currentSubtree = st
+	stp.currentSubtreeMu.Unlock()
+}
+
 // GetCurrentSubtree returns the subtree currently being built.
+// This method is safe for concurrent access from external callers.
 //
 // Returns:
 //   - *util.Subtree: Current subtree
 func (stp *SubtreeProcessor) GetCurrentSubtree() *subtreepkg.Subtree {
+	stp.currentSubtreeMu.RLock()
+	defer stp.currentSubtreeMu.RUnlock()
 	return stp.currentSubtree
 }
 
@@ -1349,10 +1364,11 @@ func (stp *SubtreeProcessor) addNode(node subtreepkg.Node, parents *subtreepkg.T
 	}
 
 	if stp.currentSubtree == nil {
-		stp.currentSubtree, err = subtreepkg.NewTreeByLeafCount(stp.currentItemsPerFile)
+		newSubtree, err := subtreepkg.NewTreeByLeafCount(stp.currentItemsPerFile)
 		if err != nil {
 			return err
 		}
+		stp.setCurrentSubtree(newSubtree)
 
 		// This is the first subtree for this block - we need a coinbase placeholder
 		err = stp.currentSubtree.AddCoinbaseNode()
@@ -1404,10 +1420,11 @@ func (stp *SubtreeProcessor) processCompleteSubtree(skipNotification bool) (err
 	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.setCurrentSubtree(newSubtree)
 
 	// Send the subtree to the newSubtreeChan, including a reference to the parent transactions map
 	errCh := make(chan error)
@@ -1637,10 +1654,11 @@ 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.setCurrentSubtree(newSubtree)
 
 	if len(originalSubtrees) == 0 {
 		// we must add the coinbase tx if we have no original subtrees
@@ -1888,7 +1906,7 @@ func (stp *SubtreeProcessor) reorgBlocks(ctx context.Context, moveBackBlocks []*
 		if err != nil {
 			// restore the original state
 			stp.chainedSubtrees = originalChainedSubtrees
-			stp.currentSubtree = originalCurrentSubtree
+			stp.setCurrentSubtree(originalCurrentSubtree)
 			stp.currentTxMap = originalCurrentTxMap
 			stp.currentBlockHeader = currentBlockHeader
 
@@ -2308,10 +2326,11 @@ 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.setCurrentSubtree(newSubtree)
 
 	stp.chainedSubtrees = make([]*subtreepkg.Subtree, 0, ExpectedNumberOfSubtrees)
 	stp.chainedSubtreeCount.Store(0)
@@ -2576,10 +2595,11 @@ 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.setCurrentSubtree(newSubtree)
 
 	stp.chainedSubtrees = make([]*subtreepkg.Subtree, 0, ExpectedNumberOfSubtrees)
 	stp.chainedSubtreeCount.Store(0)