fix(blocktree): stop and prune runtimes from memory

lib/blocktree/blocktree.go

@@ -260,6 +260,23 @@ func (bt *BlockTree) Prune(finalised Hash) (pruned []Hash) {
 		return pruned
 	}
 
+	// Cleanup in-memory runtimes from the canonical chain.
+	// The runtime used in the newly finalised block is kept
+	// instantiated in memory, and all other runtimes are
+	// stopped and removed from memory. Note these are still
+	// accessible through the storage as WASM blob.
+	previousFinalisedBlock := bt.root
+	newCanonicalChainBlocksCount := n.number - previousFinalisedBlock.number
+	if previousFinalisedBlock.number == 0 { // include the genesis block
+		newCanonicalChainBlocksCount++
+	}
+	canonicalChainBlock := n
+	newCanonicalChainBlockHashes := make([]Hash, newCanonicalChainBlocksCount)
+	for i := int(newCanonicalChainBlocksCount) - 1; i >= 0; i-- {
+		newCanonicalChainBlockHashes[i] = canonicalChainBlock.hash
+		canonicalChainBlock = canonicalChainBlock.parent
+	}
+
 	pruned = bt.root.prune(n, nil)
 	bt.root = n
 	bt.root.parent = nil
@@ -270,9 +287,7 @@ func (bt *BlockTree) Prune(finalised Hash) (pruned []Hash) {
 		bt.leaves.store(leaf.hash, leaf)
 	}
 
-	for _, hash := range pruned {
-		bt.runtimes.delete(hash)
-	}
+	bt.runtimes.onFinalisation(newCanonicalChainBlockHashes, pruned)
 
 	leavesGauge.Set(float64(len(bt.leaves.nodes())))
 	return pruned

lib/blocktree/blocktree_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -382,38 +383,281 @@ func Test_BlockTree_LowestCommonAncestor_SameChain(t *testing.T) {
 }
 
 func Test_BlockTree_Prune(t *testing.T) {
-	var bt *BlockTree
-	var branches []testBranch
+	t.Parallel()
 
-	for {
-		bt, branches = createTestBlockTree(t, testHeader, 5)
-		if len(branches) > 0 && len(bt.getNode(branches[0].hash).children) > 1 {
-			break
+	t.Run("finalised_hash_is_root_hash", func(t *testing.T) {
+		t.Parallel()
+
+		blockTree := &BlockTree{
+			root: &node{
+				hash: Hash{1},
+			},
 		}
-	}
 
-	blockTreeCopy := bt.DeepCopy()
+		pruned := blockTree.Prune(Hash{1})
 
-	// pick some block to finalise
-	finalised := bt.root.children[0].children[0].children[0]
-	pruned := bt.Prune(finalised.hash)
+		assert.Empty(t, pruned)
+	})
 
-	for _, prunedHash := range pruned {
-		prunedNode := blockTreeCopy.getNode(prunedHash)
-		if prunedNode.isDescendantOf(finalised) {
-			t.Fatal("pruned node that's descendant of finalised node!!")
+	t.Run("node_not_found", func(t *testing.T) {
+		t.Parallel()
+
+		blockTree := &BlockTree{
+			root:   &node{},
+			leaves: newEmptyLeafMap(),
 		}
 
-		if finalised.isDescendantOf(prunedNode) {
-			t.Fatal("pruned an ancestor of the finalised node!!")
+		pruned := blockTree.Prune(Hash{1})
+
+		assert.Empty(t, pruned)
+	})
+
+	t.Run("nothing_to_prune", func(t *testing.T) {
+		t.Parallel()
+
+		rootNode := &node{
+			hash:   common.Hash{1},
+			number: 0,
 		}
-	}
+		rootRuntime := NewMockRuntime(nil)
+		blockTree := &BlockTree{
+			root:     rootNode,
+			leaves:   newEmptyLeafMap(),
+			runtimes: newHashToRuntime(),
+		}
+		blockTree.runtimes.set(common.Hash{1}, rootRuntime)
+
+		// {1} -> {2}
+		parent := rootNode
+		newNode := &node{
+			parent: parent,
+			hash:   common.Hash{2},
+			number: 1,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		blockTree.runtimes.set(common.Hash{2}, rootRuntime)
 
-	require.NotEqual(t, 0, len(bt.leaves.nodes()))
-	for _, leaf := range bt.leaves.nodes() {
-		require.NotEqual(t, leaf.hash, finalised.hash)
-		require.True(t, leaf.isDescendantOf(finalised))
-	}
+		pruned := blockTree.Prune(common.Hash{2})
+		assert.Empty(t, pruned)
+
+		expectedHashToRuntime := &hashToRuntime{
+			mapping: map[common.Hash]Runtime{
+				{1}: rootRuntime,
+				{2}: rootRuntime,
+			},
+			finalisedRuntime:   rootRuntime,
+			currentBlockHashes: []Hash{{1}, {2}},
+		}
+		assert.Equal(t, expectedHashToRuntime, blockTree.runtimes)
+	})
+
+	t.Run("prune_canonical_runtimes", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+
+		rootNode := &node{
+			hash:   common.Hash{1},
+			number: 0,
+		}
+		rootRuntime := NewMockRuntime(ctrl)
+		blockTree := &BlockTree{
+			root:     rootNode,
+			leaves:   newEmptyLeafMap(),
+			runtimes: newHashToRuntime(),
+		}
+		blockTree.runtimes.set(common.Hash{1}, rootRuntime)
+
+		// {1} -> {2}
+		parent := rootNode
+		newNode := &node{
+			parent: parent,
+			hash:   common.Hash{2},
+			number: 1,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		leafRuntime := NewMockRuntime(ctrl)
+		blockTree.runtimes.set(common.Hash{2}, leafRuntime)
+
+		// Previous runtime is pruned
+		rootRuntime.EXPECT().Stop()
+
+		pruned := blockTree.Prune(common.Hash{2})
+		assert.Empty(t, pruned)
+
+		expectedHashToRuntime := &hashToRuntime{
+			mapping: map[common.Hash]Runtime{
+				{2}: leafRuntime,
+			},
+			finalisedRuntime:   leafRuntime,
+			currentBlockHashes: []Hash{{2}},
+		}
+		assert.Equal(t, expectedHashToRuntime, blockTree.runtimes)
+	})
+
+	t.Run("prune_fork", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+
+		rootNode := &node{
+			hash:   common.Hash{1},
+			number: 0,
+		}
+		rootRuntime := NewMockRuntime(ctrl)
+		blockTree := &BlockTree{
+			root:     rootNode,
+			leaves:   newEmptyLeafMap(),
+			runtimes: newHashToRuntime(),
+		}
+		blockTree.runtimes.set(common.Hash{1}, rootRuntime)
+
+		// {1} -> {2}
+		parent := rootNode
+		newNode := &node{
+			parent: parent,
+			hash:   common.Hash{2},
+			number: 1,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		blockTree.runtimes.set(common.Hash{2}, rootRuntime)
+
+		// {1} -> {2}
+		//     -> {3}
+		parent = rootNode
+		newNode = &node{
+			parent: parent,
+			hash:   common.Hash{3},
+			number: 1,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		blockTree.runtimes.set(common.Hash{3}, rootRuntime)
+
+		pruned := blockTree.Prune(common.Hash{2})
+		assert.Equal(t, []common.Hash{{3}}, pruned)
+
+		expectedHashToRuntime := &hashToRuntime{
+			mapping: map[common.Hash]Runtime{
+				{1}: rootRuntime,
+				{2}: rootRuntime,
+			},
+			finalisedRuntime:   rootRuntime,
+			currentBlockHashes: []Hash{{1}, {2}},
+		}
+		assert.Equal(t, expectedHashToRuntime, blockTree.runtimes)
+	})
+
+	t.Run("complex_example", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+
+		rootNode := &node{
+			hash:   common.Hash{1},
+			number: 100,
+		}
+		rootRuntime := NewMockRuntime(ctrl)
+		rootRuntime.EXPECT().Stop()
+
+		blockTree := &BlockTree{
+			root:   rootNode,
+			leaves: newEmptyLeafMap(),
+			runtimes: &hashToRuntime{
+				mapping: map[common.Hash]Runtime{
+					{1}: rootRuntime,
+				},
+				finalisedRuntime:   rootRuntime,
+				currentBlockHashes: []Hash{{1}},
+			},
+		}
+
+		// {1} -> {2}
+		parent := rootNode
+		newNode := &node{
+			parent: parent,
+			hash:   common.Hash{2},
+			number: 101,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		blockTree.runtimes.set(common.Hash{2}, rootRuntime)
+
+		// {1} -> {2} -> {3}
+		parent = newNode
+		newNode = &node{
+			parent: parent,
+			hash:   common.Hash{3},
+			number: 102,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		lastCanonicalRuntime := NewMockRuntime(ctrl)
+		blockTree.runtimes.set(common.Hash{3}, lastCanonicalRuntime)
+
+		// {1} -> {2} -> {3} -> {4}
+		parent = newNode
+		newNode = &node{
+			parent: parent,
+			hash:   common.Hash{4},
+			number: 103,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		blockTree.runtimes.set(common.Hash{4}, lastCanonicalRuntime)
+
+		// {1} -> {2} -> {3} -> {4}
+		//            -> {5}
+		parent = blockTree.getNode(common.Hash{2})
+		newNode = &node{
+			parent: parent,
+			hash:   common.Hash{5},
+			number: 102,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		blockTree.runtimes.set(common.Hash{5}, rootRuntime)
+
+		// {1} -> {2} -> {3} -> {4}
+		//            -> {5} -> {6}
+		parent = blockTree.getNode(common.Hash{5})
+		newNode = &node{
+			parent: parent,
+			hash:   common.Hash{6},
+			number: 103,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		blockTree.runtimes.set(common.Hash{6}, lastCanonicalRuntime)
+
+		// {1} -> {2} -> {3} -> {4}
+		//            -> {5} -> {6}
+		//						-> {7}
+		parent = blockTree.getNode(common.Hash{5})
+		newNode = &node{
+			parent: parent,
+			hash:   common.Hash{7},
+			number: 102,
+		}
+		parent.addChild(newNode)
+		blockTree.leaves.replace(parent, newNode)
+		forkedRuntime := NewMockRuntime(ctrl)
+		forkedRuntime.EXPECT().Stop()
+		blockTree.runtimes.set(common.Hash{7}, forkedRuntime)
+
+		pruned := blockTree.Prune(common.Hash{4})
+		assert.Equal(t, []common.Hash{{5}, {6}, {7}}, pruned)
+
+		expectedHashToRuntime := &hashToRuntime{
+			mapping: map[common.Hash]Runtime{
+				{3}: lastCanonicalRuntime,
+				{4}: lastCanonicalRuntime,
+			},
+			finalisedRuntime:   lastCanonicalRuntime,
+			currentBlockHashes: []Hash{{3}, {4}},
+		}
+		assert.Equal(t, expectedHashToRuntime, blockTree.runtimes)
+	})
 }
 
 func Test_BlockTree_GetHashByNumber(t *testing.T) {