go/cmd: Audit and fix context.Background() usage (#15928)

Signed-off-by: Dirkjan Bussink <d.bussink@gmail.com>

go/cmd/mysqlctl/command/init.go

@@ -55,7 +55,7 @@ func commandInit(cmd *cobra.Command, args []string) error {
  }
  defer mysqld.Close()
 
- ctx, cancel := context.WithTimeout(context.Background(), initArgs.WaitTime)
+ ctx, cancel := context.WithTimeout(cmd.Context(), initArgs.WaitTime)
  defer cancel()
  if err := mysqld.Init(ctx, cnf, initArgs.InitDbSQLFile); err != nil {
  return fmt.Errorf("failed init mysql: %v", err)

go/cmd/mysqlctl/command/shutdown.go

@@ -50,7 +50,7 @@ func commandShutdown(cmd *cobra.Command, args []string) error {
  }
  defer mysqld.Close()
 
- ctx, cancel := context.WithTimeout(context.Background(), shutdownArgs.WaitTime+10*time.Second)
+ ctx, cancel := context.WithTimeout(cmd.Context(), shutdownArgs.WaitTime+10*time.Second)
  defer cancel()
  if err := mysqld.Shutdown(ctx, cnf, true, shutdownArgs.WaitTime); err != nil {
  return fmt.Errorf("failed shutdown mysql: %v", err)

go/cmd/mysqlctl/command/start.go

@@ -51,7 +51,7 @@ func commandStart(cmd *cobra.Command, args []string) error {
  }
  defer mysqld.Close()
 
- ctx, cancel := context.WithTimeout(context.Background(), startArgs.WaitTime)
+ ctx, cancel := context.WithTimeout(cmd.Context(), startArgs.WaitTime)
  defer cancel()
  if err := mysqld.Start(ctx, cnf, startArgs.MySQLdArgs...); err != nil {
  return fmt.Errorf("failed start mysql: %v", err)

go/cmd/mysqlctl/command/teardown.go

@@ -53,7 +53,7 @@ func commandTeardown(cmd *cobra.Command, args []string) error {
  }
  defer mysqld.Close()
 
- ctx, cancel := context.WithTimeout(context.Background(), teardownArgs.WaitTime+10*time.Second)
+ ctx, cancel := context.WithTimeout(cmd.Context(), teardownArgs.WaitTime+10*time.Second)
  defer cancel()
  if err := mysqld.Teardown(ctx, cnf, teardownArgs.Force, teardownArgs.WaitTime); err != nil {
  return fmt.Errorf("failed teardown mysql (forced? %v): %v", teardownArgs.Force, err)

go/cmd/mysqlctld/cli/mysqlctld.go

@@ -114,7 +114,7 @@ func run(cmd *cobra.Command, args []string) error {
  }
 
  // Start or Init mysqld as needed.
- ctx, cancel := context.WithTimeout(context.Background(), waitTime)
+ ctx, cancel := context.WithTimeout(cmd.Context(), waitTime)
  mycnfFile := mysqlctl.MycnfFile(tabletUID)
  if _, statErr := os.Stat(mycnfFile); os.IsNotExist(statErr) {
  // Generate my.cnf from scratch and use it to find mysqld.
@@ -167,7 +167,7 @@ func run(cmd *cobra.Command, args []string) error {
  // Take mysqld down with us on SIGTERM before entering lame duck.
  servenv.OnTermSync(func() {
  log.Infof("mysqlctl received SIGTERM, shutting down mysqld first")
- ctx, cancel := context.WithTimeout(context.Background(), shutdownWaitTime+10*time.Second)
+ ctx, cancel := context.WithTimeout(cmd.Context(), shutdownWaitTime+10*time.Second)
  defer cancel()
  if err := mysqld.Shutdown(ctx, cnf, true, shutdownWaitTime); err != nil {
  log.Errorf("failed to shutdown mysqld: %v", err)

go/cmd/topo2topo/cli/topo2topo.go

@@ -90,7 +90,7 @@ func run(cmd *cobra.Command, args []string) error {
  return fmt.Errorf("Cannot open 'to' topo %v: %w", toImplementation, err)
  }
 
- ctx := context.Background()
+ ctx := cmd.Context()
 
  if compare {
  return compareTopos(ctx, fromTS, toTS)

go/cmd/vtadmin/main.go

@@ -17,7 +17,6 @@ limitations under the License.
 package main
 
 import (
- "context"
  "flag"
  "io"
  "time"
@@ -97,7 +96,7 @@ func startTracing(cmd *cobra.Command) {
 }
 
 func run(cmd *cobra.Command, args []string) {
- bootSpan, ctx := trace.NewSpan(context.Background(), "vtadmin.boot")
+ bootSpan, ctx := trace.NewSpan(cmd.Context(), "vtadmin.boot")
  defer bootSpan.Finish()
 
  configs := clusterFileConfig.Combine(defaultClusterConfig, clusterConfigs)

go/cmd/vtbackup/cli/vtbackup.go

@@ -221,10 +221,10 @@ func init() {
  collationEnv = collations.NewEnvironment(servenv.MySQLServerVersion())
 }
 
-func run(_ *cobra.Command, args []string) error {
+func run(cc *cobra.Command, args []string) error {
  servenv.Init()
 
- ctx, cancel := context.WithCancel(context.Background())
+ ctx, cancel := context.WithCancel(cc.Context())
  servenv.OnClose(func() {
  cancel()
  })
@@ -282,7 +282,7 @@ func run(_ *cobra.Command, args []string) error {
  return fmt.Errorf("Can't take backup: %w", err)
  }
  if doBackup {
- if err := takeBackup(ctx, topoServer, backupStorage); err != nil {
+ if err := takeBackup(ctx, cc.Context(), topoServer, backupStorage); err != nil {
  return fmt.Errorf("Failed to take backup: %w", err)
  }
  }
@@ -304,7 +304,7 @@ func run(_ *cobra.Command, args []string) error {
  return nil
 }
 
-func takeBackup(ctx context.Context, topoServer *topo.Server, backupStorage backupstorage.BackupStorage) error {
+func takeBackup(ctx, backgroundCtx context.Context, topoServer *topo.Server, backupStorage backupstorage.BackupStorage) error {
  // This is an imaginary tablet alias. The value doesn't matter for anything,
  // except that we generate a random UID to ensure the target backup
  // directory is unique if multiple vtbackup instances are launched for the
@@ -344,9 +344,9 @@ func takeBackup(ctx context.Context, topoServer *topo.Server, backupStorage back
  deprecatedDurationByPhase.Set("InitMySQLd", int64(time.Since(initMysqldAt).Seconds()))
  // Shut down mysqld when we're done.
  defer func() {
- // Be careful not to use the original context, because we don't want to
- // skip shutdown just because we timed out waiting for other things.
- mysqlShutdownCtx, mysqlShutdownCancel := context.WithTimeout(context.Background(), mysqlShutdownTimeout+10*time.Second)
+ // Be careful use the background context, not the init one, because we don't want to
+ // skip shutdown just because we timed out waiting for init.
+ mysqlShutdownCtx, mysqlShutdownCancel := context.WithTimeout(backgroundCtx, mysqlShutdownTimeout+10*time.Second)
  defer mysqlShutdownCancel()
  if err := mysqld.Shutdown(mysqlShutdownCtx, mycnf, false, mysqlShutdownTimeout); err != nil {
  log.Errorf("failed to shutdown mysqld: %v", err)

go/cmd/vtbench/cli/vtbench.go

@@ -212,7 +212,7 @@ func run(cmd *cobra.Command, args []string) error {
 
  b := vtbench.NewBench(threads, count, connParams, sql)
 
- ctx, cancel := context.WithTimeout(context.Background(), deadline)
+ ctx, cancel := context.WithTimeout(cmd.Context(), deadline)
  defer cancel()
 
  fmt.Printf("Initializing test with %s protocol / %d threads / %d iterations\n",

go/cmd/vtclient/cli/vtclient.go

@@ -197,7 +197,7 @@ func _run(cmd *cobra.Command, args []string) (*results, error) {
 
  log.Infof("Sending the query...")
 
- ctx, cancel := context.WithTimeout(context.Background(), timeout)
+ ctx, cancel := context.WithTimeout(cmd.Context(), timeout)
  defer cancel()
  return execMulti(ctx, db, cmd.Flags().Arg(0))
 }

go/cmd/vtclient/cli/vtclient_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package cli
 
 import (
+ "context"
  "fmt"
  "os"
  "strings"
@@ -129,6 +130,7 @@ func TestVtclient(t *testing.T) {
  err := Main.ParseFlags(args)
  require.NoError(t, err)
 
+ Main.SetContext(context.Background())
  results, err := _run(Main, args)
  if q.errMsg != "" {
  if got, want := err.Error(), q.errMsg; !strings.Contains(got, want) {

go/cmd/vtcombo/cli/main.go

@@ -138,8 +138,8 @@ func init() {
  srvTopoCounts = stats.NewCountersWithSingleLabel("ResilientSrvTopoServer", "Resilient srvtopo server operations", "type")
 }
 
-func startMysqld(uid uint32) (mysqld *mysqlctl.Mysqld, cnf *mysqlctl.Mycnf, err error) {
- ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+func startMysqld(ctx context.Context, uid uint32) (mysqld *mysqlctl.Mysqld, cnf *mysqlctl.Mycnf, err error) {
+ ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
  defer cancel()
 
  mycnfFile := mysqlctl.MycnfFile(uid)
@@ -189,17 +189,20 @@ func run(cmd *cobra.Command, args []string) (err error) {
  cmd.Flags().Set("log_dir", "$VTDATAROOT/tmp")
  }
 
+ ctx, cancel := context.WithCancel(cmd.Context())
+ defer cancel()
  if externalTopoServer {
  // Open topo server based on the command line flags defined at topo/server.go
  // do not create cell info as it should be done by whoever sets up the external topo server
  ts = topo.Open()
  } else {
  // Create topo server. We use a 'memorytopo' implementation.
- ts = memorytopo.NewServer(context.Background(), tpb.Cells...)
+ ts = memorytopo.NewServer(ctx, tpb.Cells...)
  }
+ defer ts.Close()
 
  // attempt to load any routing rules specified by tpb
- if err := vtcombo.InitRoutingRules(context.Background(), ts, tpb.GetRoutingRules()); err != nil {
+ if err := vtcombo.InitRoutingRules(ctx, ts, tpb.GetRoutingRules()); err != nil {
  return fmt.Errorf("Failed to load routing rules: %w", err)
  }
 
@@ -212,17 +215,17 @@ func run(cmd *cobra.Command, args []string) (err error) {
  )
 
  if startMysql {
- mysqld.Mysqld, cnf, err = startMysqld(1)
+ mysqld.Mysqld, cnf, err = startMysqld(ctx, 1)
  if err != nil {
  return err
  }
  servenv.OnClose(func() {
- ctx, cancel := context.WithTimeout(cmd.Context(), mysqlctl.DefaultShutdownTimeout+10*time.Second)
- defer cancel()
- mysqld.Shutdown(ctx, cnf, true, mysqlctl.DefaultShutdownTimeout)
+ shutdownCtx, shutdownCancel := context.WithTimeout(cmd.Context(), mysqlctl.DefaultShutdownTimeout+10*time.Second)
+ defer shutdownCancel()
+ mysqld.Shutdown(shutdownCtx, cnf, true, mysqlctl.DefaultShutdownTimeout)
  })
  // We want to ensure we can write to this database
- mysqld.SetReadOnly(cmd.Context(), false)
+ mysqld.SetReadOnly(ctx, false)
 
  } else {
  dbconfigs.GlobalDBConfigs.InitWithSocket("", env.CollationEnv())
@@ -241,9 +244,9 @@ func run(cmd *cobra.Command, args []string) (err error) {
  if err != nil {
  // ensure we start mysql in the event we fail here
  if startMysql {
- ctx, cancel := context.WithTimeout(cmd.Context(), mysqlctl.DefaultShutdownTimeout+10*time.Second)
- defer cancel()
- mysqld.Shutdown(ctx, cnf, true, mysqlctl.DefaultShutdownTimeout)
+ startCtx, startCancel := context.WithTimeout(ctx, mysqlctl.DefaultShutdownTimeout+10*time.Second)
+ defer startCancel()
+ mysqld.Shutdown(startCtx, cnf, true, mysqlctl.DefaultShutdownTimeout)
  }
 
  return fmt.Errorf("initTabletMapProto failed: %w", err)
@@ -287,20 +290,21 @@ func run(cmd *cobra.Command, args []string) (err error) {
 
  // Now that we have fully initialized the tablets, rebuild the keyspace graph.
  for _, ks := range tpb.Keyspaces {
- err := topotools.RebuildKeyspace(context.Background(), logutil.NewConsoleLogger(), ts, ks.GetName(), tpb.Cells, false)
+ err := topotools.RebuildKeyspace(cmd.Context(), logutil.NewConsoleLogger(), ts, ks.GetName(), tpb.Cells, false)
  if err != nil {
  if startMysql {
- ctx, cancel := context.WithTimeout(context.Background(), mysqlctl.DefaultShutdownTimeout+10*time.Second)
- defer cancel()
- mysqld.Shutdown(ctx, cnf, true, mysqlctl.DefaultShutdownTimeout)
+ shutdownCtx, shutdownCancel := context.WithTimeout(cmd.Context(), mysqlctl.DefaultShutdownTimeout+10*time.Second)
+ defer shutdownCancel()
+ mysqld.Shutdown(shutdownCtx, cnf, true, mysqlctl.DefaultShutdownTimeout)
  }
 
  return fmt.Errorf("Couldn't build srv keyspace for (%v: %v). Got error: %w", ks, tpb.Cells, err)
  }
  }
 
  // vtgate configuration and init
- resilientServer = srvtopo.NewResilientServer(context.Background(), ts, srvTopoCounts)
+
+ resilientServer = srvtopo.NewResilientServer(ctx, ts, srvTopoCounts)
 
  tabletTypes := make([]topodatapb.TabletType, 0, 1)
  if len(tabletTypesToWait) != 0 {
@@ -324,7 +328,7 @@ func run(cmd *cobra.Command, args []string) (err error) {
  vtgate.QueryzHandler = "/debug/vtgate/queryz"
 
  // pass nil for healthcheck, it will get created
- vtg := vtgate.Init(context.Background(), env, nil, resilientServer, tpb.Cells[0], tabletTypes, plannerVersion)
+ vtg := vtgate.Init(ctx, env, nil, resilientServer, tpb.Cells[0], tabletTypes, plannerVersion)
 
  // vtctld configuration and init
  err = vtctld.InitVtctld(env, ts)
@@ -333,22 +337,13 @@ func run(cmd *cobra.Command, args []string) (err error) {
  }
 
  if vschemaPersistenceDir != "" && !externalTopoServer {
- startVschemaWatcher(vschemaPersistenceDir, tpb.Keyspaces, ts)
+ startVschemaWatcher(ctx, vschemaPersistenceDir, ts)
  }
 
  servenv.OnRun(func() {
  addStatusParts(vtg)
  })
 
- servenv.OnTerm(func() {
- log.Error("Terminating")
- // FIXME(alainjobart): stop vtgate
- })
- servenv.OnClose(func() {
- // We will still use the topo server during lameduck period
- // to update our state, so closing it in OnClose()
- ts.Close()
- })
  servenv.RunDefault()
 
  return nil

go/cmd/vtcombo/cli/vschema_watcher.go

@@ -27,28 +27,27 @@ import (
  "vitess.io/vitess/go/vt/vtgate/vindexes"
 
  vschemapb "vitess.io/vitess/go/vt/proto/vschema"
- vttestpb "vitess.io/vitess/go/vt/proto/vttest"
 )
 
-func startVschemaWatcher(vschemaPersistenceDir string, keyspaces []*vttestpb.Keyspace, ts *topo.Server) {
+func startVschemaWatcher(ctx context.Context, vschemaPersistenceDir string, ts *topo.Server) {
  // Create the directory if it doesn't exist.
  if err := createDirectoryIfNotExists(vschemaPersistenceDir); err != nil {
  log.Fatalf("Unable to create vschema persistence directory %v: %v", vschemaPersistenceDir, err)
  }
 
  // If there are keyspace files, load them.
- loadKeyspacesFromDir(vschemaPersistenceDir, keyspaces, ts)
+ loadKeyspacesFromDir(ctx, vschemaPersistenceDir, ts)
 
  // Rebuild the SrvVSchema object in case we loaded vschema from file
- if err := ts.RebuildSrvVSchema(context.Background(), tpb.Cells); err != nil {
+ if err := ts.RebuildSrvVSchema(ctx, tpb.Cells); err != nil {
  log.Fatalf("RebuildSrvVSchema failed: %v", err)
  }
 
  // Now watch for changes in the SrvVSchema object and persist them to disk.
- go watchSrvVSchema(context.Background(), ts, tpb.Cells[0])
+ go watchSrvVSchema(ctx, ts, tpb.Cells[0])
 }
 
-func loadKeyspacesFromDir(dir string, keyspaces []*vttestpb.Keyspace, ts *topo.Server) {
+func loadKeyspacesFromDir(ctx context.Context, dir string, ts *topo.Server) {
  for _, ks := range tpb.Keyspaces {
  ksFile := path.Join(dir, ks.Name+".json")
  if _, err := os.Stat(ksFile); err == nil {
@@ -67,14 +66,14 @@ func loadKeyspacesFromDir(dir string, keyspaces []*vttestpb.Keyspace, ts *topo.S
  if err != nil {
  log.Fatalf("Invalid keyspace definition: %v", err)
  }
- ts.SaveVSchema(context.Background(), ks.Name, keyspace)
+ ts.SaveVSchema(ctx, ks.Name, keyspace)
  log.Infof("Loaded keyspace %v from %v\n", ks.Name, ksFile)
  }
  }
 }
 
 func watchSrvVSchema(ctx context.Context, ts *topo.Server, cell string) {
- data, ch, err := ts.WatchSrvVSchema(context.Background(), tpb.Cells[0])
+ data, ch, err := ts.WatchSrvVSchema(ctx, tpb.Cells[0])
  if err != nil {
  log.Fatalf("WatchSrvVSchema failed: %v", err)
  }

go/cmd/vtctld/cli/cli.go

@@ -79,7 +79,7 @@ func run(cmd *cobra.Command, args []string) error {
  vtctld.RegisterDebugHealthHandler(ts)
 
  // Start schema manager service.
- initSchema()
+ initSchema(cmd.Context())
 
  // And run the server.
  servenv.RunDefault()

go/cmd/vtctld/cli/schema.go

@@ -47,7 +47,7 @@ func init() {
  Main.Flags().DurationVar(&schemaChangeReplicasTimeout, "schema_change_replicas_timeout", schemaChangeReplicasTimeout, "How long to wait for replicas to receive a schema change.")
 }
 
-func initSchema() {
+func initSchema(ctx context.Context) {
  // Start schema manager service if needed.
  if schemaChangeDir != "" {
  interval := schemaChangeCheckInterval
@@ -70,7 +70,6 @@ func initSchema() {
  log.Errorf("failed to get controller, error: %v", err)
  return
  }
- ctx := context.Background()
  wr := wrangler.New(env, logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient())
  _, err = schemamanager.Run(
  ctx,

go/cmd/vtctldclient/command/legacy_shim.go

@@ -43,7 +43,7 @@ var (
  Args: cobra.ArbitraryArgs,
  RunE: func(cmd *cobra.Command, args []string) error {
  cli.FinishedParsing(cmd)
- return runLegacyCommand(args)
+ return runLegacyCommand(cmd.Context(), args)
  },
  Long: strings.TrimSpace(`
 LegacyVtctlCommand uses the legacy vtctl grpc client to make an ExecuteVtctlCommand
@@ -76,11 +76,11 @@ LegacyVtctlCommand -- AddCellInfo --server_address "localhost:5678" --root "/vit
  }
 )
 
-func runLegacyCommand(args []string) error {
+func runLegacyCommand(ctx context.Context, args []string) error {
  // Duplicated (mostly) from go/cmd/vtctlclient/main.go.
  logger := logutil.NewConsoleLogger()
 
- ctx, cancel := context.WithTimeout(context.Background(), actionTimeout)
+ ctx, cancel := context.WithTimeout(ctx, actionTimeout)
  defer cancel()
 
  err := vtctlclient.RunCommandAndWait(ctx, server, args, func(e *logutilpb.Event) {

go/cmd/vtctldclient/command/root.go

@@ -111,7 +111,7 @@ connect directly to the topo server(s).`, useInternalVtctld),
  client, err = getClientForCommand(cmd)
  ctx := cmd.Context()
  if ctx == nil {
- ctx = context.Background()
+ ctx = cmd.Context()
  }
  commandCtx, commandCancel = context.WithTimeout(ctx, actionTimeout)
  if compactOutput {