Skip to content
102 changes: 56 additions & 46 deletions services/blockassembly/subtreeprocessor/SubtreeProcessor.go

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -395,11 +395,12 @@ func TestSubtreeProcessor_CompleteSubtreeTracking(t *testing.T) {

t.Run("tracks node counts correctly", func(t *testing.T) {
// Create a subtree with known node count
stp.currentSubtree, err = subtreepkg.NewTreeByLeafCount(8)
newSubtree, err := subtreepkg.NewTreeByLeafCount(8)
require.NoError(t, err)
stp.currentSubtree.Store(newSubtree)

// Add some nodes (including coinbase)
err = stp.currentSubtree.AddCoinbaseNode()
err = stp.currentSubtree.Load().AddCoinbaseNode()
require.NoError(t, err)

// Add 4 more transaction nodes
Expand All @@ -411,7 +412,7 @@ func TestSubtreeProcessor_CompleteSubtreeTracking(t *testing.T) {
Fee: 100,
SizeInBytes: 250,
}
err = stp.currentSubtree.AddSubtreeNode(node)
err = stp.currentSubtree.Load().AddSubtreeNode(node)
require.NoError(t, err)
}

Expand Down Expand Up @@ -444,9 +445,10 @@ func TestSubtreeProcessor_CompleteSubtreeTracking(t *testing.T) {
}

// Create another subtree that should trigger the limit
stp.currentSubtree, err = subtreepkg.NewTreeByLeafCount(8)
newSubtree2, err := subtreepkg.NewTreeByLeafCount(8)
require.NoError(t, err)
err = stp.currentSubtree.AddCoinbaseNode()
stp.currentSubtree.Store(newSubtree2)
err = stp.currentSubtree.Load().AddCoinbaseNode()
require.NoError(t, err)

for i := 0; i < 3; i++ {
Expand All @@ -457,7 +459,7 @@ func TestSubtreeProcessor_CompleteSubtreeTracking(t *testing.T) {
Fee: 100,
SizeInBytes: 250,
}
err = stp.currentSubtree.AddSubtreeNode(node)
err = stp.currentSubtree.Load().AddSubtreeNode(node)
require.NoError(t, err)
}

Expand All @@ -475,9 +477,10 @@ func TestSubtreeProcessor_CompleteSubtreeTracking(t *testing.T) {
require.Equal(t, 18, count, "Should have 18 samples (full buffer)")

// Add one more to test that it maintains the limit
stp.currentSubtree, err = subtreepkg.NewTreeByLeafCount(8)
newSubtree3, err := subtreepkg.NewTreeByLeafCount(8)
require.NoError(t, err)
err = stp.currentSubtree.AddCoinbaseNode()
stp.currentSubtree.Store(newSubtree3)
err = stp.currentSubtree.Load().AddCoinbaseNode()
require.NoError(t, err)

hash := chainhash.Hash{}
Expand All @@ -487,7 +490,7 @@ func TestSubtreeProcessor_CompleteSubtreeTracking(t *testing.T) {
Fee: 100,
SizeInBytes: 250,
}
err = stp.currentSubtree.AddSubtreeNode(node)
err = stp.currentSubtree.Load().AddSubtreeNode(node)
require.NoError(t, err)

err = stp.processCompleteSubtree(false)
Expand Down
74 changes: 38 additions & 36 deletions services/blockassembly/subtreeprocessor/SubtreeProcessor_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ type SubtreeProcessorState struct {
func captureSubtreeProcessorState(stp *SubtreeProcessor) SubtreeProcessorState {
return SubtreeProcessorState{
ChainedSubtreesCount: len(stp.chainedSubtrees),
CurrentSubtreeLength: stp.currentSubtree.Length(),
CurrentSubtreeLength: stp.currentSubtree.Load().Length(),
TxCount: stp.TxCount(),
CurrentTxMapLength: stp.currentTxMap.Length(),
}
Expand Down Expand Up @@ -263,12 +263,12 @@ func Test_RemoveTxFromSubtrees(t *testing.T) {
// Subtrees 2-10: 4 txs each = 36 txs
// Remaining in current: 42 - 39 = 3 txs
t.Logf("Number of chained subtrees: %d", len(stp.chainedSubtrees))
t.Logf("Current subtree nodes: %d", len(stp.currentSubtree.Nodes))
t.Logf("Current subtree nodes: %d", len(stp.currentSubtree.Load().Nodes))
if len(stp.chainedSubtrees) > 5 {
t.Logf("Subtree 5 has %d nodes", len(stp.chainedSubtrees[5].Nodes))
}
assert.Len(t, stp.chainedSubtrees, 10)
assert.Len(t, stp.currentSubtree.Nodes, 3)
assert.Len(t, stp.currentSubtree.Load().Nodes, 3)

// get the middle transaction from the middle subtree
txHash := stp.chainedSubtrees[5].Nodes[2].Hash
Expand Down Expand Up @@ -299,11 +299,11 @@ func Test_RemoveTxFromSubtrees(t *testing.T) {
// After removing and rechaining, we may have fewer subtrees due to proper duplicate handling
// The rechaining process rebuilds from the removal point, properly detecting duplicates
t.Logf("After rechaining - Number of chained subtrees: %d", len(stp.chainedSubtrees))
t.Logf("After rechaining - Current subtree nodes: %d", len(stp.currentSubtree.Nodes))
t.Logf("After rechaining - Current subtree nodes: %d", len(stp.currentSubtree.Load().Nodes))
// We should have at least the subtrees before the removal point
assert.GreaterOrEqual(t, len(stp.chainedSubtrees), 5)
// Current subtree should have some nodes but may vary due to rechaining
assert.GreaterOrEqual(t, len(stp.currentSubtree.Nodes), 0)
assert.GreaterOrEqual(t, len(stp.currentSubtree.Load().Nodes), 0)

// check that the txHash node has been removed from the currentTxMap
_, ok = stp.currentTxMap.Get(txHash)
Expand Down Expand Up @@ -360,7 +360,7 @@ func TestReChainSubtrees(t *testing.T) {
// With 42 unique transactions and 4 items per subtree:
// 10 complete subtrees + 3 remaining in current
assert.Len(t, stp.chainedSubtrees, 10)
assert.Len(t, stp.currentSubtree.Nodes, 3)
assert.Len(t, stp.currentSubtree.Load().Nodes, 3)

// check the fee in the middle node the middle subtree
assert.Len(t, stp.chainedSubtrees[5].Nodes, 4)
Expand Down Expand Up @@ -394,7 +394,7 @@ func TestReChainSubtrees(t *testing.T) {
assert.Len(t, stp.chainedSubtrees[5].Nodes, 4)

// currentSubtree should have 2 nodes
assert.Len(t, stp.currentSubtree.Nodes, 2)
assert.Len(t, stp.currentSubtree.Load().Nodes, 2)

// all chainedSubtrees should have 4 nodes
for i := 0; i < 10; i++ {
Expand Down Expand Up @@ -457,7 +457,7 @@ func TestGetMerkleProofForCoinbase(t *testing.T) {
require.NoError(t, err)

if i == 0 {
stp.currentSubtree.ReplaceRootNode(hash, 0, 0)
stp.currentSubtree.Load().ReplaceRootNode(hash, 0, 0)
} else {
stp.Add(subtreepkg.Node{Hash: *hash, Fee: 1}, subtreepkg.TxInpoints{ParentTxHashes: []chainhash.Hash{*hash}})
}
Expand Down Expand Up @@ -499,7 +499,7 @@ func TestGetMerkleProofForCoinbase(t *testing.T) {
require.NoError(t, err)

if i == 0 {
stp.currentSubtree.ReplaceRootNode(hash, 0, 0)
stp.currentSubtree.Load().ReplaceRootNode(hash, 0, 0)
} else {
stp.Add(subtreepkg.Node{Hash: *hash, Fee: 1}, subtreepkg.TxInpoints{ParentTxHashes: []chainhash.Hash{*hash}})
}
Expand Down Expand Up @@ -592,7 +592,7 @@ func TestMoveForwardBlock(t *testing.T) {
require.NoError(t, err)

if i == 0 {
stp.currentSubtree.ReplaceRootNode(hash, 0, 0)
stp.currentSubtree.Load().ReplaceRootNode(hash, 0, 0)
} else {
stp.Add(subtreepkg.Node{Hash: *hash, Fee: 1}, subtreepkg.TxInpoints{ParentTxHashes: []chainhash.Hash{*hash}})
}
Expand All @@ -608,7 +608,7 @@ func TestMoveForwardBlock(t *testing.T) {
// there should be 4 chained subtrees
assert.Equal(t, 4, len(stp.chainedSubtrees))
assert.Equal(t, 4, stp.chainedSubtrees[0].Size())
assert.Equal(t, 2, stp.currentSubtree.Length())
assert.Equal(t, 2, stp.currentSubtree.Load().Length())

assert.Equal(t, int(n-1), stp.currentTxMap.Length()) //nolint:gosec

Expand Down Expand Up @@ -637,7 +637,7 @@ func TestMoveForwardBlock(t *testing.T) {
// we added the coinbase placeholder
assert.Equal(t, 5, len(stp.chainedSubtrees))
assert.Equal(t, 2, stp.chainedSubtrees[0].Size())
assert.Equal(t, 1, stp.currentSubtree.Length())
assert.Equal(t, 1, stp.currentSubtree.Load().Length())

// check the currentTxMap, it will have 1 less than the tx count, which has the coinbase placeholder
assert.Equal(t, int(stp.TxCount()), stp.currentTxMap.Length()+1) // nolint:gosec
Expand Down Expand Up @@ -700,8 +700,8 @@ func TestMoveForwardBlock_LeftInQueue(t *testing.T) {
require.NoError(t, err)

assert.Len(t, subtreeProcessor.chainedSubtrees, 0)
assert.Len(t, subtreeProcessor.currentSubtree.Nodes, 1)
assert.Equal(t, *subtreepkg.CoinbasePlaceholderHash, subtreeProcessor.currentSubtree.Nodes[0].Hash)
assert.Len(t, subtreeProcessor.currentSubtree.Load().Nodes, 1)
assert.Equal(t, *subtreepkg.CoinbasePlaceholderHash, subtreeProcessor.currentSubtree.Load().Nodes[0].Hash)
}

func TestIncompleteSubtreeMoveForwardBlock(t *testing.T) {
Expand Down Expand Up @@ -759,7 +759,7 @@ func TestIncompleteSubtreeMoveForwardBlock(t *testing.T) {
require.NoError(t, err)

if i == 0 {
stp.currentSubtree.ReplaceRootNode(hash, 0, 0)
stp.currentSubtree.Load().ReplaceRootNode(hash, 0, 0)
} else {
stp.Add(subtreepkg.Node{Hash: *hash, Fee: 1}, subtreepkg.TxInpoints{ParentTxHashes: []chainhash.Hash{*hash}})
}
Expand All @@ -776,7 +776,7 @@ func TestIncompleteSubtreeMoveForwardBlock(t *testing.T) {
assert.Equal(t, 4, len(stp.chainedSubtrees))
assert.Equal(t, 4, stp.chainedSubtrees[0].Size())
// and 1 tx in the current subtree
assert.Equal(t, 1, stp.currentSubtree.Length())
assert.Equal(t, 1, stp.currentSubtree.Load().Length())

stp.currentItemsPerFile = 2
_ = stp.utxoStore.SetBlockHeight(1)
Expand All @@ -801,7 +801,7 @@ func TestIncompleteSubtreeMoveForwardBlock(t *testing.T) {

wg.Wait()
assert.Equal(t, 5, len(stp.chainedSubtrees))
assert.Equal(t, 0, stp.currentSubtree.Length())
assert.Equal(t, 0, stp.currentSubtree.Load().Length())
}

// current subtree should have 1 tx which due to the new added coinbase placeholder
Expand Down Expand Up @@ -862,7 +862,7 @@ func TestSubtreeMoveForwardBlockNewCurrent(t *testing.T) {
require.NoError(t, err)

if i == 0 {
stp.currentSubtree.ReplaceRootNode(hash, 0, 0)
stp.currentSubtree.Load().ReplaceRootNode(hash, 0, 0)
} else {
stp.Add(subtreepkg.Node{Hash: *hash, Fee: 1}, subtreepkg.TxInpoints{ParentTxHashes: []chainhash.Hash{*hash}})
}
Expand All @@ -877,7 +877,7 @@ func TestSubtreeMoveForwardBlockNewCurrent(t *testing.T) {
assert.Equal(t, 4, len(stp.chainedSubtrees))
assert.Equal(t, 4, stp.chainedSubtrees[0].Size())
// and 0 tx in the current subtree
assert.Equal(t, 0, stp.currentSubtree.Length())
assert.Equal(t, 0, stp.currentSubtree.Load().Length())

stp.currentItemsPerFile = 2
_ = stp.utxoStore.SetBlockHeight(1)
Expand All @@ -902,7 +902,7 @@ func TestSubtreeMoveForwardBlockNewCurrent(t *testing.T) {
wg.Wait()
require.NoError(t, err)
assert.Equal(t, 4, len(stp.chainedSubtrees))
assert.Equal(t, 1, stp.currentSubtree.Length())
assert.Equal(t, 1, stp.currentSubtree.Load().Length())
}

func TestCompareMerkleProofsToSubtrees(t *testing.T) {
Expand Down Expand Up @@ -948,7 +948,7 @@ func TestCompareMerkleProofsToSubtrees(t *testing.T) {

for i, hash := range hashes {
if i == 0 {
subtreeProcessor.currentSubtree.ReplaceRootNode(hash, 0, 0)
subtreeProcessor.currentSubtree.Load().ReplaceRootNode(hash, 0, 0)
} else {
subtreeProcessor.Add(subtreepkg.Node{Hash: *hash, Fee: 111}, subtreepkg.TxInpoints{ParentTxHashes: []chainhash.Hash{*hash}})
}
Expand Down Expand Up @@ -1081,10 +1081,11 @@ func TestSubtreeProcessor_getRemainderTxHashes(t *testing.T) {
chainedSubtrees = append(chainedSubtrees, lastSubtree)

// Setup fresh subtree processor state
subtreeProcessor.currentSubtree, err = subtreepkg.NewTree(4)
newSubtree, err := subtreepkg.NewTree(4)
require.NoError(t, err)
subtreeProcessor.currentSubtree.Store(newSubtree)
subtreeProcessor.chainedSubtrees = make([]*subtreepkg.Subtree, 0)
_ = subtreeProcessor.currentSubtree.AddCoinbaseNode()
_ = subtreeProcessor.currentSubtree.Load().AddCoinbaseNode()

// Setup maps
transactionMap := txmap.NewSplitSwissMap(4) // Transactions that are in the new block
Expand All @@ -1108,7 +1109,7 @@ func TestSubtreeProcessor_getRemainderTxHashes(t *testing.T) {
for _, subtree := range subtreeProcessor.chainedSubtrees {
remainder = append(remainder, subtree.Nodes...)
}
remainder = append(remainder, subtreeProcessor.currentSubtree.Nodes...)
remainder = append(remainder, subtreeProcessor.currentSubtree.Load().Nodes...)

// With duplicate detection, only the coinbase should remain (no duplicates added)
assert.Equal(t, 1, len(remainder))
Expand Down Expand Up @@ -1146,12 +1147,13 @@ func TestSubtreeProcessor_getRemainderTxHashes(t *testing.T) {
// "f923a14068167a9107a0b7cd6102bfa5c0a4c8a72726a82f12e91009fd7e33be",
}

subtreeProcessor.currentSubtree, err = subtreepkg.NewTree(4)
newSubtree2, err := subtreepkg.NewTree(4)
require.NoError(t, err)
subtreeProcessor.currentSubtree.Store(newSubtree2)

subtreeProcessor.chainedSubtrees = make([]*subtreepkg.Subtree, 0)

_ = subtreeProcessor.currentSubtree.AddCoinbaseNode()
_ = subtreeProcessor.currentSubtree.Load().AddCoinbaseNode()

// Test scenario 2: Some transactions are in the new block (transactionMap)
// Clear the subtreeProcessor state but keep currentTxMap populated
Expand Down Expand Up @@ -1180,7 +1182,7 @@ func TestSubtreeProcessor_getRemainderTxHashes(t *testing.T) {
remainder = append(remainder, subtree.Nodes...)
}

remainder = append(remainder, subtreeProcessor.currentSubtree.Nodes...)
remainder = append(remainder, subtreeProcessor.currentSubtree.Load().Nodes...)

// With duplicate detection, only the coinbase remains (no duplicates added)
assert.Equal(t, 1, len(remainder))
Expand Down Expand Up @@ -1346,7 +1348,7 @@ func TestSubtreeProcessor_moveBackBlock(t *testing.T) {
// there should be 4 chained subtrees
assert.Equal(t, 4, len(stp.chainedSubtrees))
assert.Equal(t, 4, stp.chainedSubtrees[0].Size())
assert.Equal(t, 3, stp.currentSubtree.Length())
assert.Equal(t, 3, stp.currentSubtree.Load().Length())

// create 2 subtrees from the previous block
subtree1 := createSubtree(t, 4, true)
Expand Down Expand Up @@ -1388,7 +1390,7 @@ func TestSubtreeProcessor_moveBackBlock(t *testing.T) {

assert.Equal(t, 6, len(stp.chainedSubtrees))
assert.Equal(t, 4, stp.chainedSubtrees[0].Size())
assert.Equal(t, 2, stp.currentSubtree.Length())
assert.Equal(t, 2, stp.currentSubtree.Load().Length())

// check that the nodes from subtree1 and subtree2 are the first nodes
for i := 0; i < 4; i++ {
Expand All @@ -1402,7 +1404,7 @@ func TestSubtreeProcessor_moveBackBlock(t *testing.T) {
shouldBeInNode := idx % 4

if shouldBeInSubtree > len(stp.chainedSubtrees)-1 {
assert.Equal(t, txHash, stp.currentSubtree.Nodes[shouldBeInNode].Hash)
assert.Equal(t, txHash, stp.currentSubtree.Load().Nodes[shouldBeInNode].Hash)
} else {
assert.Equal(t, txHash, stp.chainedSubtrees[shouldBeInSubtree].Nodes[shouldBeInNode].Hash)
}
Expand Down Expand Up @@ -1489,7 +1491,7 @@ func TestSubtreeProcessor_moveBackBlock(t *testing.T) {

// Verify state after processing empty block
assert.Equal(t, 0, len(stp.chainedSubtrees))
assert.Equal(t, 1, stp.currentSubtree.Length()) // Should only have coinbase placeholder
assert.Equal(t, 1, stp.currentSubtree.Load().Length()) // Should only have coinbase placeholder
})

// Test subtree store errors with state reset verification
Expand Down Expand Up @@ -1595,7 +1597,7 @@ func TestSubtreeProcessor_moveBackBlock(t *testing.T) {

// Verify the coinbase placeholder was handled correctly
assert.Equal(t, 0, len(stp.chainedSubtrees))
assert.Equal(t, 1, stp.currentSubtree.Length())
assert.Equal(t, 1, stp.currentSubtree.Load().Length())
})

// Test SetBlockProcessedAt error (non-critical path)
Expand Down Expand Up @@ -2879,7 +2881,7 @@ func TestRemoveTxsFromSubtreesBasic(t *testing.T) {
// Verify transaction was added
_, exists := stp.currentTxMap.Get(txHash)
require.True(t, exists, "Transaction should be in currentTxMap")
require.True(t, stp.currentSubtree.NodeIndex(txHash) >= 0, "Transaction should be in current subtree")
require.True(t, stp.currentSubtree.Load().NodeIndex(txHash) >= 0, "Transaction should be in current subtree")

initialTxCount := stp.TxCount()

Expand All @@ -2896,7 +2898,7 @@ func TestRemoveTxsFromSubtreesBasic(t *testing.T) {
t.Logf("Transaction still in currentTxMap: %v", stillExists)

// Check if transaction was removed from current subtree
indexAfter := stp.currentSubtree.NodeIndex(txHash)
indexAfter := stp.currentSubtree.Load().NodeIndex(txHash)
t.Logf("Transaction index in current subtree after removal: %d", indexAfter)
})

Expand Down Expand Up @@ -2930,7 +2932,7 @@ func TestRemoveTxsFromSubtreesBasic(t *testing.T) {

for _, hash := range txHashes {
_, stillExists := stp.currentTxMap.Get(hash)
indexAfter := stp.currentSubtree.NodeIndex(hash)
indexAfter := stp.currentSubtree.Load().NodeIndex(hash)
t.Logf("Hash %s - still in map: %v, index: %d", hash.String()[:8], stillExists, indexAfter)
}
})
Expand Down Expand Up @@ -2967,7 +2969,7 @@ func TestRemoveTxsFromSubtreesBasic(t *testing.T) {
t.Logf("Chained subtrees after: %d", len(stp.chainedSubtrees))

_, stillExists := stp.currentTxMap.Get(targetHash)
currentIndex := stp.currentSubtree.NodeIndex(targetHash)
currentIndex := stp.currentSubtree.Load().NodeIndex(targetHash)
t.Logf("Target hash still in map: %v, current subtree index: %d", stillExists, currentIndex)

// Check chained subtrees
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -224,8 +224,8 @@ func countTransactionInSubtreesForTest(stp *SubtreeProcessor, txHash chainhash.H
}

// Check current subtree
if stp.currentSubtree != nil {
for _, node := range stp.currentSubtree.Nodes {
if currentSubtree := stp.currentSubtree.Load(); currentSubtree != nil {
for _, node := range currentSubtree.Nodes {
if node.Hash.Equal(txHash) {
count++
}
Expand Down
Loading
Loading