add context to node run

Author: janezpodhostnik

cmd/access/main.go

@@ -1,6 +1,8 @@
 package main
 
 import (
+	"context"
+
 	"github.com/onflow/flow-go/cmd"
 	nodebuilder "github.com/onflow/flow-go/cmd/access/node_builder"
 	"github.com/onflow/flow-go/model/flow"
@@ -24,5 +26,5 @@ func main() {
 	if err != nil {
 		builder.Logger.Fatal().Err(err).Send()
 	}
-	node.Run()
+	node.Run(context.Background())
 }

cmd/collection/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -646,7 +647,7 @@ func main() {
 	if err != nil {
 		nodeBuilder.Logger.Fatal().Err(err).Send()
 	}
-	node.Run()
+	node.Run(context.Background())
 }
 
 // createQCContractClient creates QC contract client

cmd/consensus/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -933,7 +934,7 @@ func main() {
 	if err != nil {
 		nodeBuilder.Logger.Fatal().Err(err).Send()
 	}
-	node.Run()
+	node.Run(context.Background())
 }
 
 func loadBeaconPrivateKey(dir string, myID flow.Identifier) (*encodable.RandomBeaconPrivKey, error) {

cmd/execution/main.go

@@ -1,6 +1,8 @@
 package main
 
 import (
+	"context"
+
 	"github.com/onflow/flow-go/cmd"
 	"github.com/onflow/flow-go/model/flow"
 )
@@ -19,5 +21,5 @@ func main() {
 	if err != nil {
 		exeBuilder.FlowNodeBuilder.Logger.Fatal().Err(err).Send()
 	}
-	node.Run()
+	node.Run(context.Background())
 }

cmd/ghost/main.go

@@ -1,6 +1,8 @@
 package main
 
 import (
+	"context"
+
 	"github.com/spf13/pflag"
 
 	"github.com/onflow/flow-go/cmd"
@@ -45,5 +47,5 @@ func main() {
 	if err != nil {
 		nodeBuilder.Logger.Fatal().Err(err).Send()
 	}
-	node.Run()
+	node.Run(context.Background())
 }

cmd/node.go

@@ -22,7 +22,8 @@ type Node interface {
 	// Run initiates all common components (logger, database, protocol state etc.)
 	// then starts each component. It also sets up a channel to gracefully shut
 	// down each component if a SIGINT is received.
-	Run()
+	// The context can also be used to signal the node to shutdown.
+	Run(ctx context.Context)
 }
 
 // FlowNodeImp is created by the FlowNodeBuilder with all components ready to be started.
@@ -86,8 +87,7 @@ func NewBaseNode(
 // which point it gracefully shuts down.
 // Any unhandled irrecoverable errors thrown in child components will propagate up to here and
 // result in a fatal error.
-func (node *NodeImp) Run() {
-	ctx := context.Background()
+func (node *NodeImp) Run(ctx context.Context) {
 
 	// Block until node is shutting down
 	err := node.run(ctx)

cmd/node_test.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"os"
 	"syscall"
@@ -42,7 +43,7 @@ func TestRunShutsDownCleanly(t *testing.T) {
 
 		finished := make(chan struct{})
 		go func() {
-			node.Run()
+			node.Run(context.Background())
 			close(finished)
 		}()
 
@@ -62,6 +63,44 @@ func TestRunShutsDownCleanly(t *testing.T) {
 		}, testLogger.logs)
 	})
 
+	t.Run("Run shuts down gracefully on context cancel", func(t *testing.T) {
+		testLogger.Reset()
+		manager := component.NewComponentManagerBuilder().
+			AddWorker(func(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
+				testLogger.Log("worker starting up")
+				ready()
+				testLogger.Log("worker startup complete")
+
+				<-ctx.Done()
+				testLogger.Log("worker shutting down")
+				testLogger.Log("worker shutdown complete")
+			}).
+			Build()
+		node := NewNode(manager, nodeConfig, logger, postShutdown, fatalHandler)
+
+		ctx, cancel := context.WithCancel(context.Background())
+
+		finished := make(chan struct{})
+		go func() {
+			node.Run(ctx)
+			close(finished)
+		}()
+
+		<-node.Ready()
+
+		cancel()
+
+		<-finished
+
+		assert.Equal(t, []string{
+			"worker starting up",
+			"worker startup complete",
+			"worker shutting down",
+			"worker shutdown complete",
+			"running cleanup",
+		}, testLogger.logs)
+	})
+
 	t.Run("Run encounters error during postShutdown", func(t *testing.T) {
 		testLogger.Reset()
 		manager := component.NewComponentManagerBuilder().
@@ -82,7 +121,7 @@ func TestRunShutsDownCleanly(t *testing.T) {
 
 		finished := make(chan struct{})
 		go func() {
-			node.Run()
+			node.Run(context.Background())
 			close(finished)
 		}()
 
@@ -123,7 +162,7 @@ func TestRunShutsDownCleanly(t *testing.T) {
 
 		finished := make(chan struct{})
 		go func() {
-			node.Run()
+			node.Run(context.Background())
 			close(finished)
 		}()
 
@@ -157,7 +196,7 @@ func TestRunShutsDownCleanly(t *testing.T) {
 
 		finished := make(chan struct{})
 		go func() {
-			node.Run()
+			node.Run(context.Background())
 			close(finished)
 		}()
 
@@ -191,7 +230,7 @@ func TestRunShutsDownCleanly(t *testing.T) {
 
 		finished := make(chan struct{})
 		go func() {
-			node.Run()
+			node.Run(context.Background())
 			close(finished)
 		}()
 

cmd/observer/main.go

@@ -1,6 +1,8 @@
 package main
 
 import (
+	"context"
+
 	nodebuilder "github.com/onflow/flow-go/cmd/observer/node_builder"
 )
 
@@ -22,5 +24,5 @@ func main() {
 	if err != nil {
 		anb.Logger.Fatal().Err(err).Send()
 	}
-	node.Run()
+	node.Run(context.Background())
 }

cmd/verification/main.go

@@ -1,6 +1,8 @@
 package main
 
 import (
+	"context"
+
 	"github.com/onflow/flow-go/cmd"
 	"github.com/onflow/flow-go/model/flow"
 )
@@ -20,5 +22,5 @@ func main() {
 	if err != nil {
 		verificationBuilder.FlowNodeBuilder.Logger.Fatal().Err(err).Send()
 	}
-	node.Run()
+	node.Run(context.Background())
 }

insecure/cmd/access/main.go

@@ -1,6 +1,8 @@
 package main
 
 import (
+	"context"
+
 	nodebuilder "github.com/onflow/flow-go/cmd/access/node_builder"
 	insecmd "github.com/onflow/flow-go/insecure/cmd"
 	"github.com/onflow/flow-go/model/flow"
@@ -31,5 +33,5 @@ func main() {
 		builder.Logger.Fatal().Err(err).Send()
 	}
 
-	node.Run()
+	node.Run(context.Background())
 }