refactor: reduce context.Background()

Signed-off-by: ashwat287 <[email protected]>

Author: ashwat287

cmd/limactl/clone.go

@@ -71,7 +71,7 @@ func cloneAction(cmd *cobra.Command, args []string) error {
 		if err != nil {
 			return err
 		}
-		yBytes, err := yqutil.EvaluateExpression(yq, yContent)
+		yBytes, err := yqutil.EvaluateExpression(cmd.Context(), yq, yContent)
 		if err != nil {
 			return err
 		}

cmd/limactl/edit.go

@@ -87,7 +87,7 @@ func editAction(cmd *cobra.Command, args []string) error {
 	var yBytes []byte
 	if len(yqExprs) > 0 {
 		yq := yqutil.Join(yqExprs)
-		yBytes, err = yqutil.EvaluateExpression(yq, yContent)
+		yBytes, err = yqutil.EvaluateExpression(cmd.Context(), yq, yContent)
 		if err != nil {
 			return err
 		}

cmd/limactl/network.go

@@ -4,6 +4,7 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -183,7 +184,7 @@ func networkCreateAction(cmd *cobra.Command, args []string) error {
 			return fmt.Errorf("network mode %q requires specifying interface", mode)
 		}
 		yq := fmt.Sprintf(`.networks.%q = {"mode":%q,"interface":%q}`, name, mode, intf)
-		return networkApplyYQ(yq)
+		return networkApplyYQ(cmd.Context(), yq)
 	default:
 		if gateway == "" {
 			return fmt.Errorf("network mode %q requires specifying gateway", mode)
@@ -208,11 +209,11 @@ func networkCreateAction(cmd *cobra.Command, args []string) error {
 		// TODO: check IP range collision
 
 		yq := fmt.Sprintf(`.networks.%q = {"mode":%q,"gateway":%q,"netmask":%q,"interface":%q}`, name, mode, gwIP.String(), gwMaskStr, intf)
-		return networkApplyYQ(yq)
+		return networkApplyYQ(cmd.Context(), yq)
 	}
 }
 
-func networkApplyYQ(yq string) error {
+func networkApplyYQ(ctx context.Context, yq string) error {
 	filePath, err := networks.ConfigFile()
 	if err != nil {
 		return err
@@ -221,7 +222,7 @@ func networkApplyYQ(yq string) error {
 	if err != nil {
 		return err
 	}
-	yBytes, err := yqutil.EvaluateExpression(yq, yContent)
+	yBytes, err := yqutil.EvaluateExpression(ctx, yq, yContent)
 	if err != nil {
 		return err
 	}
@@ -263,7 +264,7 @@ func networkDeleteAction(cmd *cobra.Command, args []string) error {
 			yq += " | "
 		}
 	}
-	return networkApplyYQ(yq)
+	return networkApplyYQ(cmd.Context(), yq)
 }
 
 func networkBashComplete(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {

cmd/limactl/start.go

@@ -234,7 +234,7 @@ func loadOrCreateInstance(cmd *cobra.Command, args []string, createOnly bool) (*
 		}
 	} else {
 		logrus.Info("Terminal is not available, proceeding without opening an editor")
-		if err := modifyInPlace(tmpl, yq); err != nil {
+		if err := modifyInPlace(cmd.Context(), tmpl, yq); err != nil {
 			return nil, err
 		}
 	}
@@ -252,7 +252,7 @@ func applyYQExpressionToExistingInstance(ctx context.Context, inst *store.Instan
 		return nil, err
 	}
 	logrus.Debugf("Applying yq expression %q to an existing instance %q", yq, inst.Name)
-	yBytes, err := yqutil.EvaluateExpression(yq, yContent)
+	yBytes, err := yqutil.EvaluateExpression(ctx, yq, yContent)
 	if err != nil {
 		return nil, err
 	}
@@ -275,8 +275,8 @@ func applyYQExpressionToExistingInstance(ctx context.Context, inst *store.Instan
 	return store.Inspect(ctx, inst.Name)
 }
 
-func modifyInPlace(st *limatmpl.Template, yq string) error {
-	out, err := yqutil.EvaluateExpression(yq, st.Bytes)
+func modifyInPlace(ctx context.Context, st *limatmpl.Template, yq string) error {
+	out, err := yqutil.EvaluateExpression(ctx, yq, st.Bytes)
 	if err != nil {
 		return err
 	}
@@ -301,7 +301,7 @@ func (exitSuccessError) ExitCode() int {
 
 func chooseNextCreatorState(ctx context.Context, tmpl *limatmpl.Template, yq string) (*limatmpl.Template, error) {
 	for {
-		if err := modifyInPlace(tmpl, yq); err != nil {
+		if err := modifyInPlace(ctx, tmpl, yq); err != nil {
 			logrus.WithError(err).Warn("Failed to evaluate yq expression")
 			return tmpl, err
 		}

cmd/limactl/template.go

@@ -119,7 +119,7 @@ func templateCopyAction(cmd *cobra.Command, args []string) error {
 	if embed && verbatim {
 		return errors.New("--verbatim cannot be used with any of --embed, --embed-all, or --fill")
 	}
-	tmpl, err := limatmpl.Read(cmd.Context(), "", source)
+	tmpl, err := limatmpl.Read(ctx, "", source)
 	if err != nil {
 		return err
 	}
@@ -129,11 +129,11 @@ func templateCopyAction(cmd *cobra.Command, args []string) error {
 	if !verbatim {
 		if embed {
 			// Embed default base.yaml only when fill is true.
-			if err := tmpl.Embed(cmd.Context(), embedAll, fill); err != nil {
+			if err := tmpl.Embed(ctx, embedAll, fill); err != nil {
 				return err
 			}
 		} else {
-			if err := tmpl.UseAbsLocators(); err != nil {
+			if err := tmpl.UseAbsLocators(ctx); err != nil {
 				return err
 			}
 		}

pkg/driver/driver.go

@@ -40,7 +40,7 @@ type Lifecycle interface {
 type GUI interface {
 	// RunGUI is for starting GUI synchronously by hostagent. This method should be wait and return only after vm terminates
 	// It returns error if there are any failures
-	RunGUI() error
+	RunGUI(ctx context.Context) error
 
 	ChangeDisplayPassword(ctx context.Context, password string) error
 	DisplayConnection(ctx context.Context) (string, error)
@@ -63,7 +63,7 @@ type Registration interface {
 // GuestAgent defines operations for the guest agent.
 type GuestAgent interface {
 	// ForwardGuestAgent returns if the guest agent sock needs forwarding by host agent.
-	ForwardGuestAgent() bool
+	ForwardGuestAgent(_ context.Context) bool
 
 	// GuestAgentConn returns the guest agent connection, or nil (if forwarded by ssh).
 	GuestAgentConn(_ context.Context) (net.Conn, string, error)
@@ -77,10 +77,10 @@ type Driver interface {
 	Registration
 	GuestAgent
 
-	Info() Info
+	Info(_ context.Context) Info
 
 	// SetConfig sets the configuration for the instance.
-	Configure(inst *store.Instance) *ConfiguredDriver
+	Configure(_ context.Context, inst *store.Instance) *ConfiguredDriver
 }
 
 type ConfiguredDriver struct {

pkg/driver/external/client/methods.go

@@ -102,10 +102,10 @@ func (d *DriverClient) Stop(ctx context.Context) error {
 	return nil
 }
 
-func (d *DriverClient) RunGUI() error {
+func (d *DriverClient) RunGUI(ctx context.Context) error {
 	d.logger.Debug("Running GUI for the driver instance")
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
 
 	_, err := d.DriverSvc.RunGUI(ctx, &emptypb.Empty{})
@@ -230,10 +230,10 @@ func (d *DriverClient) Unregister(ctx context.Context) error {
 	return nil
 }
 
-func (d *DriverClient) ForwardGuestAgent() bool {
+func (d *DriverClient) ForwardGuestAgent(ctx context.Context) bool {
 	d.logger.Debug("Checking if guest agent needs to be forwarded")
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
 
 	resp, err := d.DriverSvc.ForwardGuestAgent(ctx, &emptypb.Empty{})
@@ -256,10 +256,10 @@ func (d *DriverClient) GuestAgentConn(ctx context.Context) (net.Conn, string, er
 	return nil, "", nil
 }
 
-func (d *DriverClient) Info() driver.Info {
+func (d *DriverClient) Info(ctx context.Context) driver.Info {
 	d.logger.Debug("Getting driver info")
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
 
 	resp, err := d.DriverSvc.GetInfo(ctx, &emptypb.Empty{})
@@ -278,7 +278,7 @@ func (d *DriverClient) Info() driver.Info {
 	return info
 }
 
-func (d *DriverClient) Configure(inst *store.Instance) *driver.ConfiguredDriver {
+func (d *DriverClient) Configure(ctx context.Context, inst *store.Instance) *driver.ConfiguredDriver {
 	d.logger.Debugf("Setting config for instance %s with SSH local port %d", inst.Name, inst.SSHLocalPort)
 
 	instJSON, err := inst.MarshalJSON()
@@ -287,7 +287,7 @@ func (d *DriverClient) Configure(inst *store.Instance) *driver.ConfiguredDriver
 		return nil
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
 
 	_, err = d.DriverSvc.SetConfig(ctx, &pb.SetConfigRequest{

pkg/driver/external/server/methods.go

@@ -53,7 +53,7 @@ func (s *DriverServer) Start(_ *emptypb.Empty, stream pb.Driver_StartServer) err
 	}
 }
 
-func (s *DriverServer) SetConfig(_ context.Context, req *pb.SetConfigRequest) (*emptypb.Empty, error) {
+func (s *DriverServer) SetConfig(ctx context.Context, req *pb.SetConfigRequest) (*emptypb.Empty, error) {
 	s.logger.Debugf("Received SetConfig request")
 	var inst store.Instance
 
@@ -62,7 +62,7 @@ func (s *DriverServer) SetConfig(_ context.Context, req *pb.SetConfigRequest) (*
 		return &emptypb.Empty{}, err
 	}
 
-	_ = s.driver.Configure(&inst)
+	_ = s.driver.Configure(ctx, &inst)
 
 	return &emptypb.Empty{}, nil
 }
@@ -76,7 +76,7 @@ func (s *DriverServer) GuestAgentConn(ctx context.Context, _ *emptypb.Empty) (*e
 	}
 
 	if connType != "unix" {
-		proxySocketPath := filepath.Join(s.driver.Info().InstanceDir, filenames.GuestAgentSock)
+		proxySocketPath := filepath.Join(s.driver.Info(ctx).InstanceDir, filenames.GuestAgentSock)
 
 		var lc net.ListenConfig
 		listener, err := lc.Listen(ctx, "unix", proxySocketPath)
@@ -102,9 +102,9 @@ func (s *DriverServer) GuestAgentConn(ctx context.Context, _ *emptypb.Empty) (*e
 	return &emptypb.Empty{}, nil
 }
 
-func (s *DriverServer) GetInfo(_ context.Context, _ *emptypb.Empty) (*pb.InfoResponse, error) {
+func (s *DriverServer) GetInfo(ctx context.Context, _ *emptypb.Empty) (*pb.InfoResponse, error) {
 	s.logger.Debug("Received GetInfo request")
-	info := s.driver.Info()
+	info := s.driver.Info(ctx)
 
 	infoJSON, err := json.Marshal(info)
 	if err != nil {
@@ -161,9 +161,9 @@ func (s *DriverServer) Stop(ctx context.Context, empty *emptypb.Empty) (*emptypb
 	return empty, nil
 }
 
-func (s *DriverServer) RunGUI(_ context.Context, empty *emptypb.Empty) (*emptypb.Empty, error) {
+func (s *DriverServer) RunGUI(ctx context.Context, empty *emptypb.Empty) (*emptypb.Empty, error) {
 	s.logger.Debug("Received RunGUI request")
-	err := s.driver.RunGUI()
+	err := s.driver.RunGUI(ctx)
 	if err != nil {
 		s.logger.Errorf("RunGUI failed: %v", err)
 		return empty, err
@@ -260,7 +260,7 @@ func (s *DriverServer) Unregister(ctx context.Context, empty *emptypb.Empty) (*e
 	return empty, nil
 }
 
-func (s *DriverServer) ForwardGuestAgent(_ context.Context, _ *emptypb.Empty) (*pb.ForwardGuestAgentResponse, error) {
+func (s *DriverServer) ForwardGuestAgent(ctx context.Context, _ *emptypb.Empty) (*pb.ForwardGuestAgentResponse, error) {
 	s.logger.Debug("Received ForwardGuestAgent request")
-	return &pb.ForwardGuestAgentResponse{ShouldForward: s.driver.ForwardGuestAgent()}, nil
+	return &pb.ForwardGuestAgentResponse{ShouldForward: s.driver.ForwardGuestAgent(ctx)}, nil
 }

pkg/driver/external/server/server.go

@@ -54,7 +54,7 @@ func Serve(ctx context.Context, driver driver.Driver) {
 	logger := logrus.New()
 	logger.SetLevel(logrus.DebugLevel)
 
-	socketPath := filepath.Join(os.TempDir(), fmt.Sprintf("lima-driver-%s-%d.sock", driver.Info().DriverName, os.Getpid()))
+	socketPath := filepath.Join(os.TempDir(), fmt.Sprintf("lima-driver-%s-%d.sock", driver.Info(ctx).DriverName, os.Getpid()))
 
 	defer func() {
 		if err := os.Remove(socketPath); err != nil && !os.IsNotExist(err) {
@@ -132,8 +132,8 @@ func Serve(ctx context.Context, driver driver.Driver) {
 		}
 	}()
 
-	go func() {
-		logger.Infof("Starting external driver server for %s", driver.Info().DriverName)
+	go func(ctx context.Context) {
+		logger.Infof("Starting external driver server for %s", driver.Info(ctx).DriverName)
 		logger.Infof("Server starting on Unix socket: %s", socketPath)
 		if err := server.Serve(tListener); err != nil {
 			if errors.Is(err, grpc.ErrServerStopped) {
@@ -142,7 +142,7 @@ func Serve(ctx context.Context, driver driver.Driver) {
 				logger.Errorf("Failed to serve: %v", err)
 			}
 		}
-	}()
+	}(ctx)
 
 	<-shutdownCh
 	server.GracefulStop()
@@ -155,7 +155,7 @@ func Start(extDriver *registry.ExternalDriver, instName string) error {
 	}
 	extDriver.InstanceName = instName
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(extDriver.Ctx)
 	cmd := exec.CommandContext(ctx, extDriver.Path)
 
 	stdout, err := cmd.StdoutPipe()

pkg/driver/qemu/qemu_driver.go

@@ -64,7 +64,7 @@ func New() *LimaQemuDriver {
 	}
 }
 
-func (l *LimaQemuDriver) Configure(inst *store.Instance) *driver.ConfiguredDriver {
+func (l *LimaQemuDriver) Configure(_ context.Context, inst *store.Instance) *driver.ConfiguredDriver {
 	l.Instance = inst
 	l.SSHLocalPort = inst.SSHLocalPort
 
@@ -96,8 +96,8 @@ func (l *LimaQemuDriver) CreateDisk(ctx context.Context) error {
 	return EnsureDisk(ctx, qCfg)
 }
 
-func (l *LimaQemuDriver) Start(_ context.Context) (chan error, error) {
-	ctx, cancel := context.WithCancel(context.Background())
+func (l *LimaQemuDriver) Start(ctx context.Context) (chan error, error) {
+	ctx, cancel := context.WithCancel(ctx)
 	defer func() {
 		if l.qCmd == nil {
 			cancel()
@@ -381,7 +381,7 @@ func (l *LimaQemuDriver) shutdownQEMU(ctx context.Context, timeout time.Duration
 		logrus.WithError(err).Warnf("failed to send system_powerdown command via the QMP socket %q, forcibly killing QEMU", qmpSockPath)
 		return l.killQEMU(ctx, timeout, qCmd, qWaitCh)
 	}
-	timeoutCtx, timeoutCancel := context.WithTimeout(context.Background(), timeout)
+	timeoutCtx, timeoutCancel := context.WithTimeout(ctx, timeout)
 	defer timeoutCancel()
 
 	select {
@@ -528,7 +528,7 @@ func (a *qArgTemplateApplier) applyTemplate(qArg string) (string, error) {
 	return b.String(), nil
 }
 
-func (l *LimaQemuDriver) Info() driver.Info {
+func (l *LimaQemuDriver) Info(_ context.Context) driver.Info {
 	var info driver.Info
 	if l.Instance != nil && l.Instance.Dir != "" {
 		info.InstanceDir = l.Instance.Dir
@@ -544,7 +544,7 @@ func (l *LimaQemuDriver) Initialize(_ context.Context) error {
 	return nil
 }
 
-func (l *LimaQemuDriver) RunGUI() error {
+func (l *LimaQemuDriver) RunGUI(_ context.Context) error {
 	return nil
 }
 
@@ -556,7 +556,7 @@ func (l *LimaQemuDriver) Unregister(_ context.Context) error {
 	return nil
 }
 
-func (l *LimaQemuDriver) ForwardGuestAgent() bool {
+func (l *LimaQemuDriver) ForwardGuestAgent(_ context.Context) bool {
 	// if driver is not providing, use host agent
 	return l.vSockPort == 0 && l.virtioPort == ""
 }