Pass Context from app.Run

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

cmd/mps-control-daemon/main.go

@@ -55,8 +55,8 @@ func main() {
 	c := cli.Command{}
 	c.Name = "NVIDIA MPS Control Daemon"
 	c.Version = info.GetVersionString()
-	c.Action = func(_ context.Context, cmd *cli.Command) error {
-		return start(cmd, config)
+	c.Action = func(ctx context.Context, cmd *cli.Command) error {
+		return start(ctx, cmd, config)
 	}
 	c.Commands = []*cli.Command{
 		mount.NewCommand(),
@@ -106,7 +106,7 @@ func (cfg *Config) loadConfig(c *cli.Command) (*spec.Config, error) {
 	return config, nil
 }
 
-func start(c *cli.Command, cfg *Config) error {
+func start(ctx context.Context, c *cli.Command, cfg *Config) error {
 	klog.Info("Starting OS watcher.")
 	sigs := watch.Signals(syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
 	var started bool
@@ -122,7 +122,7 @@ restart:
 	}
 
 	klog.Info("Starting Daemons.")
-	daemons, restartDaemons, err := startDaemons(c, cfg)
+	daemons, restartDaemons, err := startDaemons(ctx, c, cfg)
 	if err != nil {
 		return fmt.Errorf("error starting plugins: %v", err)
 	}
@@ -162,7 +162,7 @@ exit:
 	return nil
 }
 
-func startDaemons(c *cli.Command, cfg *Config) ([]*mps.Daemon, bool, error) {
+func startDaemons(ctx context.Context, c *cli.Command, cfg *Config) ([]*mps.Daemon, bool, error) {
 	// Load the configuration file
 	klog.Info("Loading configuration.")
 	config, err := cfg.loadConfig(c)
@@ -209,7 +209,7 @@ func startDaemons(c *cli.Command, cfg *Config) ([]*mps.Daemon, bool, error) {
 	// Loop through all MPS daemons and start them.
 	// If any daemon fails to start, all daemons are started again.
 	for _, mpsDaemon := range mpsDaemons {
-		if err := mpsDaemon.Start(); err != nil {
+		if err := mpsDaemon.Start(ctx); err != nil {
 			klog.Errorf("Failed to start MPS daemon: %v", err)
 			return mpsDaemons, true, nil
 		}

cmd/mps-control-daemon/mps/daemon.go

@@ -18,6 +18,7 @@ package mps
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -89,7 +90,7 @@ func (d *Daemon) EnvVars() envvars {
 }
 
 // Start starts the MPS deamon as a background process.
-func (d *Daemon) Start() error {
+func (d *Daemon) Start(ctx context.Context) error {
 	if err := d.setComputeMode(computeModeExclusiveProcess); err != nil {
 		return fmt.Errorf("error setting compute mode %v: %w", computeModeExclusiveProcess, err)
 	}
@@ -137,7 +138,7 @@ func (d *Daemon) Start() error {
 
 	d.logTailer = newTailer(filepath.Join(logDir, "control.log"))
 	klog.InfoS("Starting log tailer", "resource", d.rm.Resource())
-	if err := d.logTailer.Start(); err != nil {
+	if err := d.logTailer.Start(ctx); err != nil {
 		klog.ErrorS(err, "Could not start tail command on control.log; ignoring logs")
 	}
 

cmd/mps-control-daemon/mps/log-tailer.go

@@ -37,8 +37,8 @@ func newTailer(filename string) *tailer {
 }
 
 // Start starts tailing the specified filename.
-func (t *tailer) Start() error {
-	ctx, cancel := context.WithCancel(context.Background())
+func (t *tailer) Start(ctx context.Context) error {
+	ctx, cancel := context.WithCancel(ctx)
 	t.cancel = cancel
 
 	//nolint:gosec // G204: Subprocess launched with a potential tainted input or cmd arguments (gosec)