Add maxprocs to the agent

kylecarbs · kylecarbs · commit 3aed5c5bb898 · 2023-11-08T16:47:50.000Z
diff --git a/cli/agent.go b/cli/agent.go
@@ -17,6 +17,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/compute/metadata"
+	"go.uber.org/automaxprocs/maxprocs"
 	"golang.org/x/xerrors"
 	"gopkg.in/natefinch/lumberjack.v2"
 	"tailscale.com/util/clientmetric"
@@ -162,6 +163,16 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 			sinks = append(sinks, sloghuman.Sink(logWriter))
 			logger := slog.Make(sinks...).Leveled(slog.LevelDebug)
 
+			// This can improve the performance of Coder inside of a container.
+			// See: https://github.com/uber-go/automaxprocs
+			undoMacProcs, err := maxprocs.Set(maxprocs.Logger(func(format string, args ...interface{}) {
+				logger.Debug(ctx, fmt.Sprintf(format, args...))
+			}))
+			if err != nil {
+				return xerrors.Errorf("set maxprocs: %w", err)
+			}
+			defer undoMacProcs()
+
 			version := buildinfo.Version()
 			logger.Info(ctx, "agent is starting now",
 				slog.F("url", r.agentURL),
diff --git a/cli/server.go b/cli/server.go
@@ -326,9 +326,13 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 
 			// This can improve the performance of Coder inside of a container.
 			// See: https://github.com/uber-go/automaxprocs
-			maxprocs.Set(maxprocs.Logger(func(format string, args ...interface{}) {
+			undoMacProcs, err := maxprocs.Set(maxprocs.Logger(func(format string, args ...interface{}) {
 				logger.Debug(ctx, fmt.Sprintf(format, args...))
 			}))
+			if err != nil {
+				return xerrors.Errorf("set maxprocs: %w", err)
+			}
+			defer undoMacProcs()
 
 			// Register signals early on so that graceful shutdown can't
 			// be interrupted by additional signals. Note that we avoid
