fix(codersdk): abort in-progress writes/reads when closing websocket

ammario · ammario · commit d195ae679cf1 · 2024-03-18T11:52:19.000-05:00
Fixes #9203 Also, adds some basic tracing infrastructure that we can build upon for more improvements.
diff --git a/cli/root.go b/cli/root.go
@@ -18,6 +18,7 @@ import (
 	"os/signal"
 	"path/filepath"
 	"runtime"
+	"runtime/trace"
 	"strings"
 	"syscall"
 	"text/tabwriter"
@@ -139,6 +140,22 @@ func (r *RootCmd) AGPL() []*serpent.Command {
 
 // Main is the entrypoint for the Coder CLI.
 func (r *RootCmd) RunMain(subcommands []*serpent.Command) {
+	// This configuration is not available as a standard option because we
+	// want to trace the entire program, including Options parsing.
+	goTraceFilePath, ok := os.LookupEnv("CODER_GO_TRACE")
+	if ok {
+		traceFile, err := os.OpenFile(goTraceFilePath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
+		if err != nil {
+			panic(fmt.Sprintf("failed to open trace file: %v", err))
+		}
+		defer traceFile.Close()
+
+		if err := trace.Start(traceFile); err != nil {
+			panic(fmt.Sprintf("failed to start trace: %v", err))
+		}
+		defer trace.Stop()
+	}
+
 	rand.Seed(time.Now().UnixMicro())
 
 	cmd, err := r.Command(subcommands)
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -11,6 +11,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"runtime/trace"
 	"strings"
 	"sync"
 	"time"
@@ -67,6 +68,8 @@ func (r *RootCmd) ssh() *serpent.Command {
 			r.InitClient(client),
 		),
 		Handler: func(inv *serpent.Invocation) (retErr error) {
+			defer trace.StartRegion(inv.Context(), "ssh").End()
+
 			// Before dialing the SSH server over TCP, capture Interrupt signals
 			// so that if we are interrupted, we have a chance to tear down the
 			// TCP session cleanly before exiting.  If we don't, then the TCP
diff --git a/codersdk/websocket.go b/codersdk/websocket.go
@@ -32,7 +32,7 @@ func (c *wsNetConn) Write(b []byte) (n int, err error) {
 }
 
 func (c *wsNetConn) Close() error {
-	defer c.cancel()
+	c.cancel()
 	return c.Conn.Close()
 }
 

Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ func (c *wsNetConn) Write(b []byte) (n int, err error) {`
`32`	`32`	`}`
`33`	`33`
`34`	`34`	`func (c *wsNetConn) Close() error {`
`35`		`- defer c.cancel()`
	`35`	`+ c.cancel()`
`36`	`36`	`return c.Conn.Close()`
`37`	`37`	`}`
`38`	`38`