Jun 6, 2022 · Jun 2, 2022 · Jun 3, 2022 · Jun 3, 2022 · Jun 3, 2022 · Jun 3, 2022
diff --git a/cli/agent.go b/cli/agent.go
 import (
 "context"
 "net/http"
 _ "net/http/pprof" //nolint: gosec
 "net/url"
 "os"
 "path/filepath"

 func workspaceAgent() *cobra.Command {
 var (
 auth string
 auth         string
 pprofEnabled bool
 pprofAddress string
 )
 cmd := &cobra.Command{
 Use: "agent",
 logger := slog.Make(sloghuman.Sink(cmd.ErrOrStderr()), sloghuman.Sink(logWriter)).Leveled(slog.LevelDebug)
 client := codersdk.New(coderURL)

 if pprofEnabled {
 srvClose := serveHandler(cmd.Context(), logger, nil, pprofAddress, "pprof")
 defer srvClose()
 } else {
 // If pprof wasn't enabled at startup, allow a
 // `kill -USR1 $agent_pid` to start it (on Unix).
 srvClose := agentStartPPROFOnUSR1(cmd.Context(), logger, pprofAddress)
 defer srvClose()
 }

 // exchangeToken returns a session token.
 // This is abstracted to allow for the same looping condition
 // regardless of instance identity auth type.
 }

 cliflag.StringVarP(cmd.Flags(), &auth, "auth", "", "CODER_AGENT_AUTH", "token", "Specify the authentication type to use for the agent")
 cliflag.BoolVarP(cmd.Flags(), &pprofEnabled, "pprof-enable", "", "CODER_AGENT_PPROF_ENABLE", false, "Enable serving pprof metrics on the address defined by --pprof-address.")
 cliflag.StringVarP(cmd.Flags(), &pprofAddress, "pprof-address", "", "CODER_AGENT_PPROF_ADDRESS", "127.0.0.1:6060", "The address to serve pprof.")
 return cmd
 }
diff --git a/cli/agent_unix.go b/cli/agent_unix.go
 //go:build !windows

 package cli

 import (
 "context"
 "os"
 "os/signal"
 "syscall"

 "cdr.dev/slog"
 )

 func agentStartPPROFOnUSR1(ctx context.Context, logger slog.Logger, pprofAddress string) (srvClose func()) {
 ctx, cancel := context.WithCancel(ctx)

 usr1 := make(chan os.Signal, 1)
 signal.Notify(usr1, syscall.SIGUSR1)
 go func() {
 defer close(usr1)
 defer signal.Stop(usr1)

 select {
 case <-usr1:
 signal.Stop(usr1)
 srvClose := serveHandler(ctx, logger, nil, pprofAddress, "pprof")
 defer srvClose()
 case <-ctx.Done():
 return
 }
 <-ctx.Done() // Prevent defer close until done.
 }()

 return func() {
 cancel()
 <-usr1 // Wait until usr1 is closed, ensures srvClose was run.
 }
 }
diff --git a/cli/agent_windows.go b/cli/agent_windows.go
 package cli

 import (
 "context"

 "cdr.dev/slog"
 )

 // agentStartPPROFOnUSR1 is no-op on Windows (no SIGUSR1 signal).
 func agentStartPPROFOnUSR1(ctx context.Context, logger slog.Logger, pprofAddress string) (srvClose func()) {
 return func() {}
 }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -3,6 +3,7 @@ package cli
		import (
		"context"
		"net/http"
		_ "net/http/pprof" //nolint: gosec
		"net/url"
		"os"
		"path/filepath"
Expand All		@@ -25,7 +26,9 @@ import (

		func workspaceAgent() *cobra.Command {
		var (
		auth string
		auth string
		pprofEnabled bool
		pprofAddress string
		)
		cmd := &cobra.Command{
		Use: "agent",
Expand All		@@ -49,6 +52,16 @@ func workspaceAgent() *cobra.Command {
		logger := slog.Make(sloghuman.Sink(cmd.ErrOrStderr()), sloghuman.Sink(logWriter)).Leveled(slog.LevelDebug)
		client := codersdk.New(coderURL)

		if pprofEnabled {
		srvClose := serveHandler(cmd.Context(), logger, nil, pprofAddress, "pprof")
		defer srvClose()
		} else {
		// If pprof wasn't enabled at startup, allow a
		// `kill -USR1 $agent_pid` to start it (on Unix).
		srvClose := agentStartPPROFOnUSR1(cmd.Context(), logger, pprofAddress)
		defer srvClose()
		}

		// exchangeToken returns a session token.
		// This is abstracted to allow for the same looping condition
		// regardless of instance identity auth type.
Expand DownExpand Up		@@ -139,5 +152,7 @@ func workspaceAgent() *cobra.Command {
		}

		cliflag.StringVarP(cmd.Flags(), &auth, "auth", "", "CODER_AGENT_AUTH", "token", "Specify the authentication type to use for the agent")
		cliflag.BoolVarP(cmd.Flags(), &pprofEnabled, "pprof-enable", "", "CODER_AGENT_PPROF_ENABLE", false, "Enable serving pprof metrics on the address defined by --pprof-address.")
		cliflag.StringVarP(cmd.Flags(), &pprofAddress, "pprof-address", "", "CODER_AGENT_PPROF_ADDRESS", "127.0.0.1:6060", "The address to serve pprof.")
		return cmd
		}
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,38 @@
		//go:build !windows

		package cli

		import (
		"context"
		"os"
		"os/signal"
		"syscall"

		"cdr.dev/slog"
		)

		func agentStartPPROFOnUSR1(ctx context.Context, logger slog.Logger, pprofAddress string) (srvClose func()) {
		ctx, cancel := context.WithCancel(ctx)

		usr1 := make(chan os.Signal, 1)
		signal.Notify(usr1, syscall.SIGUSR1)
		go func() {
		defer close(usr1)
		defer signal.Stop(usr1)

		select {
		case <-usr1:
		signal.Stop(usr1)
		srvClose := serveHandler(ctx, logger, nil, pprofAddress, "pprof")
		defer srvClose()
		case <-ctx.Done():
		return
		}
		<-ctx.Done() // Prevent defer close until done.
		}()

		return func() {
		cancel()
		<-usr1 // Wait until usr1 is closed, ensures srvClose was run.
		}
		}
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,12 @@
		package cli

		import (
		"context"

		"cdr.dev/slog"
		)

		// agentStartPPROFOnUSR1 is no-op on Windows (no SIGUSR1 signal).
		func agentStartPPROFOnUSR1(ctx context.Context, logger slog.Logger, pprofAddress string) (srvClose func()) {
		return func() {}
		}