chore(coderd): add logging to agent rpc yamux conn (#11965)

coadler · web-flow · commit 4ed1f5581a63 · 2024-01-31T23:17:20.000-06:00
diff --git a/coderd/files.go b/coderd/files.go
@@ -12,10 +12,10 @@ import (
 	"io"
 	"net/http"
 
-	"cdr.dev/slog"
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
 
+	"cdr.dev/slog"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/httpapi"
diff --git a/coderd/workspaceagentsrpc.go b/coderd/workspaceagentsrpc.go
@@ -37,6 +37,7 @@ import (
 // @x-apidocgen {"skip": true}
 func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
+	logger := api.Logger.Named("agentrpc")
 
 	version := r.URL.Query().Get("version")
 	if version == "" {
@@ -61,7 +62,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 	defer api.WebsocketWaitGroup.Done()
 	workspaceAgent := httpmw.WorkspaceAgent(r)
 
-	build, ok := ensureLatestBuild(ctx, api.Database, api.Logger, rw, workspaceAgent)
+	build, ok := ensureLatestBuild(ctx, api.Database, logger, rw, workspaceAgent)
 	if !ok {
 		return
 	}
@@ -84,6 +85,12 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	logger = logger.With(
+		slog.F("owner", owner.Username),
+		slog.F("workspace_name", workspace.Name),
+		slog.F("agent_name", workspaceAgent.Name),
+	)
+
 	conn, err := websocket.Accept(rw, r, nil)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
@@ -96,7 +103,11 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 	ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
 	defer wsNetConn.Close()
 
-	mux, err := yamux.Server(wsNetConn, nil)
+	ycfg := yamux.DefaultConfig()
+	ycfg.LogOutput = nil
+	ycfg.Logger = slog.Stdlib(ctx, logger.Named("yamux"), slog.LevelInfo)
+
+	mux, err := yamux.Server(wsNetConn, ycfg)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
 			Message: "Failed to start yamux over websocket.",
@@ -106,12 +117,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 	}
 	defer mux.Close()
 
-	api.Logger.Debug(ctx, "accepting agent RPC connection",
-		slog.F("owner", owner.Username),
-		slog.F("workspace", workspace.Name),
-		slog.F("name", workspaceAgent.Name),
-	)
-	api.Logger.Debug(ctx, "accepting agent details", slog.F("agent", workspaceAgent))
+	logger.Debug(ctx, "accepting agent RPC connection", slog.F("agent", workspaceAgent))
 
 	closeCtx, closeCtxCancel := context.WithCancel(ctx)
 	defer closeCtxCancel()
@@ -122,7 +128,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 		AgentID: workspaceAgent.ID,
 
 		Ctx:                               api.ctx,
-		Log:                               api.Logger,
+		Log:                               logger,
 		Database:                          api.Database,
 		Pubsub:                            api.Pubsub,
 		DerpMapFn:                         api.DERPMap,
@@ -155,7 +161,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 	ctx = agentapi.WithAPIVersion(ctx, version)
 	err = agentAPI.Serve(ctx, mux)
 	if err != nil {
-		api.Logger.Warn(ctx, "workspace agent RPC listen error", slog.Error(err))
+		logger.Warn(ctx, "workspace agent RPC listen error", slog.Error(err))
 		_ = conn.Close(websocket.StatusInternalError, err.Error())
 		return
 	}
