Simplify shutdown

2 files changed, 27 insertions, 16 deletions

diff --git a/plugin/evm/vm.go b/plugin/evm/vm.go
index 5d379b7..c9bdad9 100644
--- a/plugin/evm/vm.go
+++ b/plugin/evm/vm.go
@@ -156,10 +156,9 @@ type VM struct {
 
 	acceptedDB database.Database
 
-	txPoolStabilizedLock         sync.Mutex
-	txPoolStabilizedHead         common.Hash
-	txPoolStabilizedOk           chan struct{}
-	txPoolStabilizedShutdownChan chan struct{}
+	txPoolStabilizedLock sync.Mutex
+	txPoolStabilizedHead common.Hash
+	txPoolStabilizedOk   chan struct{}
 
 	metalock                     sync.Mutex
 	blockCache, blockStatusCache cache.LRU
@@ -175,7 +174,6 @@ type VM struct {
 	genlock               sync.Mutex
 	txSubmitChan          <-chan struct{}
 	atomicTxSubmitChan    chan struct{}
-	shutdownSubmitChan    chan struct{}
 	baseCodec             codec.Codec
 	codec                 codec.Manager
 	clock                 timer.Clock
@@ -183,7 +181,8 @@ type VM struct {
 	pendingAtomicTxs      chan *Tx
 	blockAtomicInputCache cache.LRU
 
-	shutdownWg sync.WaitGroup
+	shutdownChan chan struct{}
+	shutdownWg   sync.WaitGroup
 
 	fx secp256k1fx.Fx
 }
@@ -232,6 +231,7 @@ func (vm *VM) Initialize(
 		return errUnsupportedFXs
 	}
 
+	vm.shutdownChan = make(chan struct{}, 1)
 	vm.ctx = ctx
 	vm.chaindb = Database{db}
 	g := new(core.Genesis)
@@ -350,11 +350,9 @@ func (vm *VM) Initialize(
 	vm.bdTimerState = bdTimerStateLong
 	vm.bdGenWaitFlag = true
 	vm.txPoolStabilizedOk = make(chan struct{}, 1)
-	vm.txPoolStabilizedShutdownChan = make(chan struct{}, 1) // Signal goroutine to shutdown
 	// TODO: read size from options
 	vm.pendingAtomicTxs = make(chan *Tx, 1024)
 	vm.atomicTxSubmitChan = make(chan struct{}, 1)
-	vm.shutdownSubmitChan = make(chan struct{}, 1)
 	vm.newMinedBlockSub = vm.chain.SubscribeNewMinedBlockEvent()
 	vm.shutdownWg.Add(1)
 	go ctx.Log.RecoverAndPanic(vm.awaitTxPoolStabilized)
@@ -413,8 +411,8 @@ func (vm *VM) Shutdown() error {
 	}
 
 	vm.writeBackMetadata()
-	close(vm.txPoolStabilizedShutdownChan)
-	close(vm.shutdownSubmitChan)
+	vm.blockDelayTimer.Stop()
+	close(vm.shutdownChan)
 	vm.chain.Stop()
 	vm.shutdownWg.Wait()
 	return nil
@@ -763,7 +761,7 @@ func (vm *VM) awaitTxPoolStabilized() {
 				vm.txPoolStabilizedLock.Unlock()
 			default:
 			}
-		case <-vm.txPoolStabilizedShutdownChan:
+		case <-vm.shutdownChan:
 			return
 		}
 	}
@@ -782,7 +780,7 @@ func (vm *VM) awaitSubmittedTxs() {
 			vm.tryBlockGen()
 		case <-time.After(5 * time.Second):
 			vm.tryBlockGen()
-		case <-vm.shutdownSubmitChan:
+		case <-vm.shutdownChan:
 			return
 		}
 	}
diff --git a/plugin/evm/vm_test.go b/plugin/evm/vm_test.go
index 4593404..dbceb1b 100644
--- a/plugin/evm/vm_test.go
+++ b/plugin/evm/vm_test.go
@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"math/big"
 	"testing"
+	"time"
 
 	"github.com/ava-labs/avalanchego/api/keystore"
 	"github.com/ava-labs/avalanchego/cache"
@@ -155,11 +156,23 @@ func GenesisVM(t *testing.T, finishBootstrapping bool) (chan engCommon.Message,
 func TestVMGenesis(t *testing.T) {
 	_, vm, _, _ := GenesisVM(t, true)
 
-	defer func() {
-		if err := vm.Shutdown(); err != nil {
-			t.Fatal(err)
+	shutdownChan := make(chan error, 1)
+	shutdownFunc := func() {
+		err := vm.Shutdown()
+		shutdownChan <- err
+	}
+
+	go shutdownFunc()
+	shutdownTimeout := 10 * time.Millisecond
+	ticker := time.NewTicker(shutdownTimeout)
+	select {
+	case <-ticker.C:
+		t.Fatalf("VM shutdown took longer than timeout: %v", shutdownTimeout)
+	case err := <-shutdownChan:
+		if err != nil {
+			t.Fatalf("Shutdown errored: %s", err)
 		}
-	}()
+	}
 }
 
 func TestIssueAtomicTxs(t *testing.T) {