From 996ea99fb1f4aa6516bbf68086554a6817599797 Mon Sep 17 00:00:00 2001
From: Alex Willmer <alex@moreati.org.uk>
Date: Tue, 21 Apr 2020 01:34:22 +0100
Subject: [PATCH] vms: Ensure all spchainvm.VM instances in tests get shutdown

---
 vms/spchainvm/consensus_benchmark_test.go | 2 ++
 vms/spchainvm/vm_benchmark_test.go        | 7 +++++++
 vms/spchainvm/vm_test.go                  | 1 +
 3 files changed, 10 insertions(+)

diff --git a/vms/spchainvm/consensus_benchmark_test.go b/vms/spchainvm/consensus_benchmark_test.go
index aa80e6d..8b6352d 100644
--- a/vms/spchainvm/consensus_benchmark_test.go
+++ b/vms/spchainvm/consensus_benchmark_test.go
@@ -66,6 +66,7 @@ func ConsensusLeader(numBlocks, numTxsPerBlock int, b *testing.B) {
 
 		// Initialize the VM
 		vm := &VM{}
+		defer func() { ctx.Lock.Lock(); vm.Shutdown(); vm.ctx.Lock.Unlock() }()
 		ctx.Lock.Lock()
 		if err := vm.Initialize(ctx, vmDB, genesisData, msgChan, nil); err != nil {
 			b.Fatal(err)
@@ -198,6 +199,7 @@ func ConsensusFollower(numBlocks, numTxsPerBlock int, b *testing.B) {
 		vm := &VM{
 			onAccept: func(ids.ID) { wg.Done() },
 		}
+		defer func() { ctx.Lock.Lock(); vm.Shutdown(); vm.ctx.Lock.Unlock() }()
 		ctx.Lock.Lock()
 		if err := vm.Initialize(ctx, vmDB, genesisData, msgChan, nil); err != nil {
 			b.Fatal(err)
diff --git a/vms/spchainvm/vm_benchmark_test.go b/vms/spchainvm/vm_benchmark_test.go
index ac4ca13..ba450b1 100644
--- a/vms/spchainvm/vm_benchmark_test.go
+++ b/vms/spchainvm/vm_benchmark_test.go
@@ -73,6 +73,7 @@ func BenchmarkParseBlock(b *testing.B) {
 		/*testing=*/ b,
 	)
 	vm := &VM{}
+	defer func() { ctx.Lock.Lock(); vm.Shutdown(); vm.ctx.Lock.Unlock() }()
 	vm.Initialize(
 		/*ctx=*/ ctx,
 		/*db=*/ memdb.New(),
@@ -106,6 +107,7 @@ func BenchmarkParseAndVerify(b *testing.B) {
 
 	for n := 0; n < b.N; n++ {
 		vm := &VM{}
+		defer func() { ctx.Lock.Lock(); vm.Shutdown(); vm.ctx.Lock.Unlock() }()
 		vm.Initialize(
 			/*ctx=*/ snow.DefaultContextTest(),
 			/*db=*/ memdb.New(),
@@ -141,6 +143,8 @@ func BenchmarkAccept(b *testing.B) {
 
 	for n := 0; n < b.N; n++ {
 		vm := &VM{}
+		defer func() { ctx.Lock.Lock(); vm.Shutdown(); vm.ctx.Lock.Unlock() }()
+
 		vm.Initialize(
 			/*ctx=*/ snow.DefaultContextTest(),
 			/*db=*/ memdb.New(),
@@ -178,6 +182,7 @@ func ParseAndVerifyAndAccept(numBlocks, numTxsPerBlock int, b *testing.B) {
 
 	for n := 0; n < b.N; n++ {
 		vm := &VM{}
+		defer func() { ctx.Lock.Lock(); vm.Shutdown(); vm.ctx.Lock.Unlock() }()
 		vm.Initialize(
 			/*ctx=*/ snow.DefaultContextTest(),
 			/*db=*/ memdb.New(),
@@ -232,6 +237,7 @@ func ParseThenVerifyThenAccept(numBlocks, numTxsPerBlock int, b *testing.B) {
 
 	for n := 0; n < b.N; n++ {
 		vm := &VM{}
+		defer func() { ctx.Lock.Lock(); vm.Shutdown(); vm.ctx.Lock.Unlock() }()
 		vm.Initialize(
 			/*ctx=*/ snow.DefaultContextTest(),
 			/*db=*/ memdb.New(),
@@ -292,6 +298,7 @@ func IssueAndVerifyAndAccept(numBlocks, numTxsPerBlock int, b *testing.B) {
 
 	for n := 0; n < b.N; n++ {
 		vm := &VM{}
+		defer func() { ctx.Lock.Lock(); vm.Shutdown(); vm.ctx.Lock.Unlock() }()
 		vm.Initialize(
 			/*ctx=*/ snow.DefaultContextTest(),
 			/*db=*/ memdb.New(),
diff --git a/vms/spchainvm/vm_test.go b/vms/spchainvm/vm_test.go
index 4432d1e..fb54941 100644
--- a/vms/spchainvm/vm_test.go
+++ b/vms/spchainvm/vm_test.go
@@ -67,6 +67,7 @@ func TestPayments(t *testing.T) {
 	blocker, _ := queue.New(bootstrappingDB)
 
 	vm := &VM{}
+	defer func() { ctx.Lock.Lock(); vm.Shutdown(); vm.ctx.Lock.Unlock() }()
 	vm.Initialize(ctx, db, genesisData, msgChan, nil)
 
 	sender := &common.SenderTest{}