Cancel node context after StopAll (#8289)

* cancel node context after StopAll * build fix * cancel in validator Co-authored-by: Victor Farazdagi <simple.square@gmail.com> Co-authored-by: terence tsao <terence@prysmaticlabs.com>
prysmaticlabs · Jan 25, 2021 · fc8dc21 · fc8dc21
1 parent 7f5ffb7
commit fc8dc21
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/beacon-chain/node/node.go b/beacon-chain/node/node.go
@@ -60,6 +60,7 @@ const testSkipPowFlag = "test-skip-pow"
 type BeaconNode struct {
 	cliCtx          *cli.Context
 	ctx             context.Context
+	cancel          context.CancelFunc
 	services        *shared.ServiceRegistry
 	lock            sync.RWMutex
 	stop            chan struct{} // Channel to wait for termination notifications.
@@ -153,9 +154,11 @@ func NewBeaconNode(cliCtx *cli.Context) (*BeaconNode, error) {
 
 	registry := shared.NewServiceRegistry()
 
+	ctx, cancel := context.WithCancel(cliCtx.Context)
 	beacon := &BeaconNode{
 		cliCtx:          cliCtx,
-		ctx:             cliCtx.Context,
+		ctx:             ctx,
+		cancel:          cancel,
 		services:        registry,
 		stop:            make(chan struct{}),
 		stateFeed:       new(event.Feed),
@@ -278,6 +281,7 @@ func (b *BeaconNode) Close() {
 	if err := b.db.Close(); err != nil {
 		log.Errorf("Failed to close database: %v", err)
 	}
+	b.cancel()
 	close(b.stop)
 }
 

diff --git a/slasher/node/node.go b/slasher/node/node.go
@@ -150,11 +150,11 @@ func (n *SlasherNode) Close() {
 	defer n.lock.Unlock()
 
 	log.Info("Stopping hash slinging slasher")
-	n.cancel()
 	n.services.StopAll()
 	if err := n.db.Close(); err != nil {
 		log.Errorf("Failed to close database: %v", err)
 	}
+	n.cancel()
 	close(n.stop)
 }
 

diff --git a/validator/node/node.go b/validator/node/node.go
@@ -45,6 +45,8 @@ import (
 // the entire lifecycle of services attached to it participating in eth2.
 type ValidatorClient struct {
 	cliCtx            *cli.Context
+	ctx               context.Context
+	cancel            context.CancelFunc
 	db                *kv.Store
 	services          *shared.ServiceRegistry // Lifecycle and service store.
 	lock              sync.RWMutex
@@ -76,8 +78,11 @@ func NewValidatorClient(cliCtx *cli.Context) (*ValidatorClient, error) {
 	prereq.WarnIfNotSupported(cliCtx.Context)
 
 	registry := shared.NewServiceRegistry()
+	ctx, cancel := context.WithCancel(cliCtx.Context)
 	ValidatorClient := &ValidatorClient{
 		cliCtx:            cliCtx,
+		ctx:               ctx,
+		cancel:            cancel,
 		services:          registry,
 		walletInitialized: new(event.Feed),
 		stop:              make(chan struct{}),
@@ -154,6 +159,7 @@ func (c *ValidatorClient) Close() {
 
 	c.services.StopAll()
 	log.Info("Stopping Prysm validator")
+	c.cancel()
 	close(c.stop)
 }