feat: add logging to client tailnet yamux (#11908)

spikecurtis · web-flow · commit d3983e4dba0b · 2024-01-30T09:58:59.000+04:00
Adds logging to yamux when used for tailnet client connections, e.g. CLI and wsproxy.  This could be useful for debugging connection issues with tailnet v2 API.
diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -436,7 +436,10 @@ func (tac *tailnetAPIConnector) dial() (proto.DRPCTailnetClient, error) {
 		}
 		return nil, err
 	}
-	client, err := tailnet.NewDRPCClient(websocket.NetConn(tac.ctx, ws, websocket.MessageBinary))
+	client, err := tailnet.NewDRPCClient(
+		websocket.NetConn(tac.ctx, ws, websocket.MessageBinary),
+		tac.logger,
+	)
 	if err != nil {
 		tac.logger.Debug(tac.ctx, "failed to create DRPCClient", slog.Error(err))
 		_ = ws.Close(websocket.StatusInternalError, "")
diff --git a/enterprise/wsproxy/wsproxysdk/wsproxysdk.go b/enterprise/wsproxy/wsproxysdk/wsproxysdk.go
@@ -458,7 +458,7 @@ func (c *Client) DialCoordinator(ctx context.Context) (agpl.MultiAgentConn, erro
 	go httpapi.HeartbeatClose(ctx, logger, cancel, conn)
 
 	nc := websocket.NetConn(ctx, conn, websocket.MessageBinary)
-	client, err := agpl.NewDRPCClient(nc)
+	client, err := agpl.NewDRPCClient(nc, logger)
 	if err != nil {
 		logger.Debug(ctx, "failed to create DRPCClient", slog.Error(err))
 		_ = conn.Close(websocket.StatusInternalError, "")
@@ -488,7 +488,9 @@ func (c *Client) DialCoordinator(ctx context.Context) (agpl.MultiAgentConn, erro
 
 	go func() {
 		<-ctx.Done()
-		ma.Close()
+		_ = ma.Close()
+		_ = client.DRPCConn().Close()
+		<-client.DRPCConn().Closed()
 		_ = conn.Close(websocket.StatusGoingAway, "closed")
 	}()
 
diff --git a/tailnet/client.go b/tailnet/client.go
@@ -1,19 +1,21 @@
 package tailnet
 
 import (
-	"io"
+	"context"
 	"net"
 
 	"github.com/hashicorp/yamux"
 	"golang.org/x/xerrors"
 
+	"cdr.dev/slog"
 	"github.com/coder/coder/v2/codersdk/drpc"
 	"github.com/coder/coder/v2/tailnet/proto"
 )
 
-func NewDRPCClient(conn net.Conn) (proto.DRPCTailnetClient, error) {
+func NewDRPCClient(conn net.Conn, logger slog.Logger) (proto.DRPCTailnetClient, error) {
 	config := yamux.DefaultConfig()
-	config.LogOutput = io.Discard
+	config.LogOutput = nil
+	config.Logger = slog.Stdlib(context.Background(), logger, slog.LevelInfo)
 	session, err := yamux.Client(conn, config)
 	if err != nil {
 		return nil, xerrors.Errorf("multiplex client: %w", err)
diff --git a/tailnet/coordinator_test.go b/tailnet/coordinator_test.go
@@ -464,7 +464,7 @@ func TestRemoteCoordination(t *testing.T) {
 		serveErr <- err
 	}()
 
-	client, err := tailnet.NewDRPCClient(cC)
+	client, err := tailnet.NewDRPCClient(cC, logger)
 	require.NoError(t, err)
 	protocol, err := client.Coordinate(ctx)
 	require.NoError(t, err)
diff --git a/tailnet/service_test.go b/tailnet/service_test.go
@@ -51,7 +51,7 @@ func TestClientService_ServeClient_V2(t *testing.T) {
 		errCh <- err
 	}()
 
-	client, err := tailnet.NewDRPCClient(c)
+	client, err := tailnet.NewDRPCClient(c, logger)
 	require.NoError(t, err)
 
 	// Coordinate

Original file line number	Diff line number	Diff line change
`@@ -436,7 +436,10 @@ func (tac *tailnetAPIConnector) dial() (proto.DRPCTailnetClient, error) {`
`436`	`436`	`}`
`437`	`437`	`return nil, err`
`438`	`438`	`}`
`439`		`- client, err := tailnet.NewDRPCClient(websocket.NetConn(tac.ctx, ws, websocket.MessageBinary))`
	`439`	`+ client, err := tailnet.NewDRPCClient(`
	`440`	`+ websocket.NetConn(tac.ctx, ws, websocket.MessageBinary),`
	`441`	`+ tac.logger,`
	`442`	`+ )`
`440`	`443`	`if err != nil {`
`441`	`444`	`tac.logger.Debug(tac.ctx, "failed to create DRPCClient", slog.Error(err))`
`442`	`445`	`_ = ws.Close(websocket.StatusInternalError, "")`