Initialize cancellable root context in main.go (#13252)

Co-authored-by: Kasey Kirkham <[email protected]>
Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/node/node.go

@@ -118,7 +118,7 @@ type BeaconNode struct {
 
 // New creates a new node instance, sets up configuration options, and registers
 // every required service to the node.
-func New(cliCtx *cli.Context, opts ...Option) (*BeaconNode, error) {
+func New(cliCtx *cli.Context, cancel context.CancelFunc, opts ...Option) (*BeaconNode, error) {
 	if err := configureTracing(cliCtx); err != nil {
 		return nil, err
 	}
@@ -166,7 +166,7 @@ func New(cliCtx *cli.Context, opts ...Option) (*BeaconNode, error) {
 
 	registry := runtime.NewServiceRegistry()
 
-	ctx, cancel := context.WithCancel(cliCtx.Context)
+	ctx := cliCtx.Context
 	beacon := &BeaconNode{
 		cliCtx:                  cliCtx,
 		ctx:                     ctx,

beacon-chain/node/node_test.go

@@ -33,6 +33,12 @@ import (
 // Ensure BeaconNode implements interfaces.
 var _ statefeed.Notifier = (*BeaconNode)(nil)
 
+func newCliContextWithCancel(app *cli.App, set *flag.FlagSet) (*cli.Context, context.CancelFunc) {
+	context, cancel := context.WithCancel(context.Background())
+	parent := &cli.Context{Context: context}
+	return cli.NewContext(app, set, parent), cancel
+}
+
 // Test that beacon chain node can close.
 func TestNodeClose_OK(t *testing.T) {
 	hook := logTest.NewGlobal()
@@ -49,9 +55,9 @@ func TestNodeClose_OK(t *testing.T) {
 	require.NoError(t, set.Set("suggested-fee-recipient", "0x6e35733c5af9B61374A128e6F85f553aF09ff89A"))
 	cmd.ValidatorMonitorIndicesFlag.Value = &cli.IntSlice{}
 	cmd.ValidatorMonitorIndicesFlag.Value.SetInt(1)
-	ctx := cli.NewContext(&app, set, nil)
+	ctx, cancel := newCliContextWithCancel(&app, set)
 
-	node, err := New(ctx)
+	node, err := New(ctx, cancel)
 	require.NoError(t, err)
 
 	node.Close()
@@ -68,8 +74,8 @@ func TestNodeStart_Ok(t *testing.T) {
 	set.String("suggested-fee-recipient", "0x6e35733c5af9B61374A128e6F85f553aF09ff89A", "fee recipient")
 	require.NoError(t, set.Set("suggested-fee-recipient", "0x6e35733c5af9B61374A128e6F85f553aF09ff89A"))
 
-	ctx := cli.NewContext(&app, set, nil)
-	node, err := New(ctx, WithBlockchainFlagOptions([]blockchain.Option{}),
+	ctx, cancel := newCliContextWithCancel(&app, set)
+	node, err := New(ctx, cancel, WithBlockchainFlagOptions([]blockchain.Option{}),
 		WithBuilderFlagOptions([]builder.Option{}),
 		WithExecutionChainOptions([]execution.Option{}),
 		WithBlobStorage(filesystem.NewEphemeralBlobStorage(t)))
@@ -81,7 +87,6 @@ func TestNodeStart_Ok(t *testing.T) {
 	time.Sleep(3 * time.Second)
 	node.Close()
 	require.LogsContain(t, hook, "Starting beacon node")
-
 }
 
 func TestNodeStart_Ok_registerDeterministicGenesisService(t *testing.T) {
@@ -117,8 +122,8 @@ func TestNodeStart_Ok_registerDeterministicGenesisService(t *testing.T) {
 	require.NoError(t, err)
 	require.NoError(t, os.WriteFile("genesis_ssz.json", genesisBytes, 0666))
 	set.String("genesis-state", "genesis_ssz.json", "")
-	ctx := cli.NewContext(&app, set, nil)
-	node, err := New(ctx, WithBlockchainFlagOptions([]blockchain.Option{}),
+	ctx, cancel := newCliContextWithCancel(&app, set)
+	node, err := New(ctx, cancel, WithBlockchainFlagOptions([]blockchain.Option{}),
 		WithBuilderFlagOptions([]builder.Option{}),
 		WithExecutionChainOptions([]execution.Option{}))
 	require.NoError(t, err)
@@ -149,12 +154,12 @@ func TestClearDB(t *testing.T) {
 	set.Bool(cmd.ForceClearDB.Name, true, "force clear db")
 	set.String("suggested-fee-recipient", "0x6e35733c5af9B61374A128e6F85f553aF09ff89A", "fee recipient")
 	require.NoError(t, set.Set("suggested-fee-recipient", "0x6e35733c5af9B61374A128e6F85f553aF09ff89A"))
-	context := cli.NewContext(&app, set, nil)
+	context, cancel := newCliContextWithCancel(&app, set)
 	options := []Option{
 		WithExecutionChainOptions([]execution.Option{execution.WithHttpEndpoint(endpoint)}),
 		WithBlobStorage(filesystem.NewEphemeralBlobStorage(t)),
 	}
-	_, err = New(context, options...)
+	_, err = New(context, cancel, options...)
 	require.NoError(t, err)
 	require.LogsContain(t, hook, "Removing database")
 }

cmd/beacon-chain/main.go

@@ -2,6 +2,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -142,11 +143,14 @@ func init() {
 }
 
 func main() {
+	// rctx = root context with cancellation.
+	// note other instances of ctx in this func are *cli.Context.
+	rctx, cancel := context.WithCancel(context.Background())
 	app := cli.App{}
 	app.Name = "beacon-chain"
 	app.Usage = "this is a beacon chain implementation for Ethereum"
 	app.Action = func(ctx *cli.Context) error {
-		if err := startNode(ctx); err != nil {
+		if err := startNode(ctx, cancel); err != nil {
 			return cli.Exit(err.Error(), 1)
 		}
 		return nil
@@ -219,12 +223,12 @@ func main() {
 		}
 	}()
 
-	if err := app.Run(os.Args); err != nil {
+	if err := app.RunContext(rctx, os.Args); err != nil {
 		log.Error(err.Error())
 	}
 }
 
-func startNode(ctx *cli.Context) error {
+func startNode(ctx *cli.Context, cancel context.CancelFunc) error {
 	// Fix data dir for Windows users.
 	outdatedDataDir := filepath.Join(file.HomeDir(), "AppData", "Roaming", "Eth2")
 	currentDataDir := ctx.String(cmd.DataDirFlag.Name)
@@ -292,7 +296,7 @@ func startNode(ctx *cli.Context) error {
 		}
 	}
 
-	beacon, err := node.New(ctx, opts...)
+	beacon, err := node.New(ctx, cancel, opts...)
 	if err != nil {
 		return fmt.Errorf("unable to start beacon node: %w", err)
 	}