fix: Unify context usage for agent cmd and logs (#5059)

agent/agent.go

@@ -174,7 +174,7 @@ func (a *agent) run(ctx context.Context) error {
 	if err != nil {
 		return xerrors.Errorf("fetch metadata: %w", err)
 	}
-	a.logger.Info(context.Background(), "fetched metadata")
+	a.logger.Info(ctx, "fetched metadata")
 	oldMetadata := a.metadata.Swap(metadata)
 
 	// The startup script should only execute on the first run!
@@ -208,7 +208,7 @@ func (a *agent) run(ctx context.Context) error {
 	a.closeMutex.Lock()
 	network := a.network
 	a.closeMutex.Unlock()
-	if a.network == nil {
+	if network == nil {
 		a.logger.Debug(ctx, "creating tailnet")
 		network, err = a.createTailnet(ctx, metadata.DERPMap)
 		if err != nil {
@@ -365,7 +365,7 @@ func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error
 		return err
 	}
 	defer coordinator.Close()
-	a.logger.Info(context.Background(), "connected to coordination server")
+	a.logger.Info(ctx, "connected to coordination server")
 	sendNodes, errChan := tailnet.ServeCoordinator(coordinator, network.UpdateNodes)
 	network.SetNodeCallback(sendNodes)
 	select {

cli/agent.go

@@ -36,6 +36,9 @@ func workspaceAgent() *cobra.Command {
 		// This command isn't useful to manually execute.
 		Hidden: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx, cancel := context.WithCancel(cmd.Context())
+			defer cancel()
+
 			rawURL, err := cmd.Flags().GetString(varAgentURL)
 			if err != nil {
 				return xerrors.Errorf("CODER_AGENT_URL must be set: %w", err)
@@ -57,22 +60,22 @@ func workspaceAgent() *cobra.Command {
 			// Spawn a reaper so that we don't accumulate a ton
 			// of zombie processes.
 			if reaper.IsInitProcess() && !noReap && isLinux {
-				logger.Info(cmd.Context(), "spawning reaper process")
+				logger.Info(ctx, "spawning reaper process")
 				// Do not start a reaper on the child process. It's important
 				// to do this else we fork bomb ourselves.
 				args := append(os.Args, "--no-reap")
 				err := reaper.ForkReap(reaper.WithExecArgs(args...))
 				if err != nil {
-					logger.Error(cmd.Context(), "failed to reap", slog.Error(err))
+					logger.Error(ctx, "failed to reap", slog.Error(err))
 					return xerrors.Errorf("fork reap: %w", err)
 				}
 
-				logger.Info(cmd.Context(), "reaper process exiting")
+				logger.Info(ctx, "reaper process exiting")
 				return nil
 			}
 
 			version := buildinfo.Version()
-			logger.Info(cmd.Context(), "starting agent",
+			logger.Info(ctx, "starting agent",
 				slog.F("url", coderURL),
 				slog.F("auth", auth),
 				slog.F("version", version),
@@ -84,7 +87,7 @@ func workspaceAgent() *cobra.Command {
 			// Enable pprof handler
 			// This prevents the pprof import from being accidentally deleted.
 			_ = pprof.Handler
-			pprofSrvClose := serveHandler(cmd.Context(), logger, nil, pprofAddress, "pprof")
+			pprofSrvClose := serveHandler(ctx, logger, nil, pprofAddress, "pprof")
 			defer pprofSrvClose()
 
 			// exchangeToken returns a session token.
@@ -102,7 +105,7 @@ func workspaceAgent() *cobra.Command {
 				// This is *only* done for testing to mock client authentication.
 				// This will never be set in a production scenario.
 				var gcpClient *metadata.Client
-				gcpClientRaw := cmd.Context().Value("gcp-client")
+				gcpClientRaw := ctx.Value("gcp-client")
 				if gcpClientRaw != nil {
 					gcpClient, _ = gcpClientRaw.(*metadata.Client)
 				}
@@ -113,7 +116,7 @@ func workspaceAgent() *cobra.Command {
 				// This is *only* done for testing to mock client authentication.
 				// This will never be set in a production scenario.
 				var awsClient *http.Client
-				awsClientRaw := cmd.Context().Value("aws-client")
+				awsClientRaw := ctx.Value("aws-client")
 				if awsClientRaw != nil {
 					awsClient, _ = awsClientRaw.(*http.Client)
 					if awsClient != nil {
@@ -127,7 +130,7 @@ func workspaceAgent() *cobra.Command {
 				// This is *only* done for testing to mock client authentication.
 				// This will never be set in a production scenario.
 				var azureClient *http.Client
-				azureClientRaw := cmd.Context().Value("azure-client")
+				azureClientRaw := ctx.Value("azure-client")
 				if azureClientRaw != nil {
 					azureClient, _ = azureClientRaw.(*http.Client)
 					if azureClient != nil {
@@ -166,7 +169,7 @@ func workspaceAgent() *cobra.Command {
 					"GIT_ASKPASS": executablePath,
 				},
 			})
-			<-cmd.Context().Done()
+			<-ctx.Done()
 			return closer.Close()
 		},
 	}