4390: tests for memory leaks (#251)

Author: liam

services/blockassembly/subtreeprocessor/SubtreeProcessor.go

@@ -835,6 +835,10 @@ func (stp *SubtreeProcessor) reset(blockHeader *model.BlockHeader, moveBackBlock
 	// clear current tx map
 	stp.currentTxMap.Clear()
 
+	// clear remove map to prevent memory leak - entries for transactions that were
+	// never dequeued would otherwise accumulate indefinitely across resets
+	stp.removeMap = txmap.NewSwissMap(0)
+
 	// reset tx count
 	stp.setTxCountFromSubtrees()
 

services/blockassembly/subtreeprocessor/SubtreeProcessor_test.go

@@ -2116,6 +2116,98 @@ func TestSubtreeProcessor_moveBackBlock(t *testing.T) {
 }
 
 func Test_removeMap(t *testing.T) {
+	t.Run("removeMap is cleared during reset - memory leak fix verification", func(t *testing.T) {
+		// This test verifies that removeMap entries ARE cleared during reset,
+		// preventing memory leaks from orphaned remove entries.
+
+		settings := test.CreateBaseTestSettings(t)
+		settings.BlockAssembly.InitialMerkleItemsPerSubtree = 128
+
+		newSubtreeChan := make(chan NewSubtreeRequest, 100)
+		done := make(chan struct{})
+		defer close(done)
+
+		go func() {
+			for {
+				select {
+				case newSubtreeRequest := <-newSubtreeChan:
+					if newSubtreeRequest.ErrChan != nil {
+						newSubtreeRequest.ErrChan <- nil
+					}
+				case <-done:
+					return
+				}
+			}
+		}()
+
+		// Create a mock blockchain client for reset operations
+		mockBlockchainClient := &blockchain.Mock{}
+
+		stp, err := NewSubtreeProcessor(t.Context(), ulogger.TestLogger{}, settings, nil, mockBlockchainClient, nil, newSubtreeChan)
+		require.NoError(t, err)
+		stp.Start(t.Context())
+
+		// Create some transaction hashes and add them ONLY to removeMap
+		// (not to the queue) - simulating a scenario where Remove() is called
+		// for transactions that were never queued
+		orphanedRemoveHashes := make([]chainhash.Hash, 50)
+		for i := 0; i < 50; i++ {
+			txHash := chainhash.HashH([]byte(fmt.Sprintf("orphaned-remove-tx-%d", i)))
+			orphanedRemoveHashes[i] = txHash
+			// Add directly to removeMap without ever queuing
+			err := stp.removeMap.Put(txHash)
+			require.NoError(t, err)
+		}
+
+		// Verify removeMap has entries before reset
+		removeMapLengthBeforeReset := stp.removeMap.Length()
+		assert.Equal(t, 50, removeMapLengthBeforeReset, "removeMap should have 50 entries before reset")
+
+		// Also add some transactions to currentTxMap to verify it IS cleared
+		for i := 0; i < 10; i++ {
+			txHash := chainhash.HashH([]byte(fmt.Sprintf("current-tx-%d", i)))
+			stp.currentTxMap.Set(txHash, subtreepkg.TxInpoints{})
+		}
+		currentTxMapLengthBeforeReset := stp.currentTxMap.Length()
+		assert.Equal(t, 10, currentTxMapLengthBeforeReset, "currentTxMap should have 10 entries before reset")
+
+		// Create a target header for reset
+		merkleRoot := chainhash.HashH([]byte("merkle"))
+		prevBlock := chainhash.HashH([]byte("prev"))
+		resetTargetHeader := &model.BlockHeader{
+			Version:        1,
+			HashMerkleRoot: &merkleRoot,
+			HashPrevBlock:  &prevBlock,
+			Timestamp:      1234567890,
+			Nonce:          0,
+		}
+
+		// Initialize the current block header (required for reset)
+		stp.InitCurrentBlockHeader(resetTargetHeader)
+
+		// Perform reset with no moveBack or moveForward blocks
+		response := stp.Reset(resetTargetHeader, nil, nil, false, nil)
+		require.NoError(t, response.Err, "Reset should succeed")
+
+		// Verify currentTxMap WAS cleared (expected behavior)
+		currentTxMapLengthAfterReset := stp.currentTxMap.Length()
+		assert.Equal(t, 0, currentTxMapLengthAfterReset, "currentTxMap should be cleared after reset")
+
+		// Verify removeMap WAS cleared (fix for memory leak)
+		removeMapLengthAfterReset := stp.removeMap.Length()
+		assert.Equal(t, 0, removeMapLengthAfterReset,
+			"removeMap should be cleared after reset to prevent memory leak")
+
+		// Verify the specific orphaned hashes are NO LONGER in removeMap
+		for i, hash := range orphanedRemoveHashes {
+			exists := stp.removeMap.Exists(hash)
+			assert.False(t, exists, "Orphaned hash %d should NOT exist in removeMap after reset", i)
+		}
+
+		t.Logf("Memory leak fix verified: removeMap cleared from %d to %d entries after reset",
+			removeMapLengthBeforeReset, removeMapLengthAfterReset)
+	})
+
 	t.Run("when adding from queue", func(t *testing.T) {
 		settings := test.CreateBaseTestSettings(t)
 		settings.BlockAssembly.InitialMerkleItemsPerSubtree = 128
@@ -2173,6 +2265,112 @@ func Test_removeMap(t *testing.T) {
 		assert.Equal(t, uint64(expectedNrTransactions-transactionsRemoved+1), stp.TxCount()) //nolint:gosec  // +1 for coinbase
 		assert.Equal(t, expectedNrTransactions-transactionsRemoved, stp.currentTxMap.Length())
 	})
+
+	t.Run("moveBackBlock adds to removeMap and reset clears it", func(t *testing.T) {
+		// This test verifies that when moveBackBlock processes coinbase child spends,
+		// entries are added to removeMap, and these entries are properly cleared during
+		// reset to prevent memory leaks.
+
+		ctx := context.Background()
+		utxoStoreURL, err := url.Parse("sqlitememory:///test")
+		require.NoError(t, err)
+
+		utxoStore, err := sql.New(ctx, ulogger.TestLogger{}, test.CreateBaseTestSettings(t), utxoStoreURL)
+		require.NoError(t, err)
+		require.NoError(t, utxoStore.SetBlockHeight(4))
+
+		blobStore := blob_memory.New()
+		settings := test.CreateBaseTestSettings(t)
+		settings.BlockAssembly.InitialMerkleItemsPerSubtree = 128
+
+		newSubtreeChan := make(chan NewSubtreeRequest, 10)
+		go func() {
+			for req := range newSubtreeChan {
+				if req.ErrChan != nil {
+					req.ErrChan <- nil
+				}
+			}
+		}()
+		defer close(newSubtreeChan)
+
+		stp, err := NewSubtreeProcessor(ctx, ulogger.TestLogger{}, settings, blobStore, nil, utxoStore, newSubtreeChan)
+		require.NoError(t, err)
+		stp.Start(ctx)
+
+		// Create coinbase transaction
+		coinbase := coinbaseTx
+		_, err = utxoStore.Create(ctx, coinbase, 1)
+		require.NoError(t, err)
+
+		// Create child transaction spending from coinbase
+		childTx := bt.NewTx()
+		err = childTx.From(coinbase.TxIDChainHash().String(), 0, coinbase.Outputs[0].LockingScript.String(), uint64(coinbase.Outputs[0].Satoshis))
+		require.NoError(t, err)
+		err = childTx.AddP2PKHOutputFromAddress("1A1zP1eP5QGefi2DMPTfTL5SLmv7DivfNa", 400000000)
+		require.NoError(t, err)
+		childTx.Inputs[0].UnlockingScript = bscript.NewFromBytes([]byte{})
+
+		// Create child in UTXO store and establish parent-child relationship
+		_, err = utxoStore.Create(ctx, childTx, 1)
+		require.NoError(t, err)
+		spends, err := utxoStore.Spend(ctx, childTx, 2, utxo.IgnoreFlags{})
+		require.NoError(t, err)
+		for _, spend := range spends {
+			require.NoError(t, spend.Err)
+		}
+
+		childHash := *childTx.TxIDChainHash()
+
+		// Verify removeMap is empty before
+		removeMapLengthBefore := stp.removeMap.Length()
+		assert.Equal(t, 0, removeMapLengthBefore, "removeMap should be empty initially")
+
+		// Create block for moveBack
+		block := &model.Block{
+			CoinbaseTx: coinbase,
+			Header: &model.BlockHeader{
+				Version:        1,
+				HashPrevBlock:  &chainhash.Hash{},
+				HashMerkleRoot: &chainhash.Hash{},
+				Timestamp:      1234567890,
+				Bits:           model.NBit{},
+				Nonce:          12345,
+			},
+			Subtrees: []*chainhash.Hash{},
+		}
+
+		// Call removeCoinbaseUtxos directly (this is what moveBackBlock calls)
+		err = stp.removeCoinbaseUtxos(ctx, block)
+		require.NoError(t, err)
+
+		// Verify child hash was added to removeMap
+		removeMapLengthAfter := stp.removeMap.Length()
+		assert.Equal(t, 1, removeMapLengthAfter, "removeMap should have 1 entry after removeCoinbaseUtxos")
+		assert.True(t, stp.removeMap.Exists(childHash), "childHash should be in removeMap")
+
+		// Now verify that reset clears the removeMap entry
+		merkleRoot := chainhash.HashH([]byte("merkle"))
+		prevBlock := chainhash.HashH([]byte("prev"))
+		resetTargetHeader := &model.BlockHeader{
+			Version:        1,
+			HashMerkleRoot: &merkleRoot,
+			HashPrevBlock:  &prevBlock,
+			Timestamp:      1234567890,
+			Nonce:          0,
+		}
+		stp.InitCurrentBlockHeader(resetTargetHeader)
+
+		// Reset with no moveBack/moveForward blocks
+		response := stp.Reset(resetTargetHeader, nil, nil, false, nil)
+		require.NoError(t, response.Err)
+
+		// Verify removeMap is cleared after reset (memory leak fix)
+		removeMapLengthAfterReset := stp.removeMap.Length()
+		assert.Equal(t, 0, removeMapLengthAfterReset,
+			"removeMap should be cleared after reset to prevent memory leak")
+		assert.False(t, stp.removeMap.Exists(childHash),
+			"childHash should be removed from removeMap after reset")
+	})
 }
 
 func waitForSubtreeProcessorQueueToEmpty(t *testing.T, stp *SubtreeProcessor) {