coder · code-asher · Aug 14, 2023 · Jul 18, 2023 · Aug 8, 2023 · Aug 8, 2023
diff --git a/agent/agent.go b/agent/agent.go
@@ -21,7 +21,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/armon/circbuf"
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
@@ -36,12 +35,12 @@ import (
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/agent/agentssh"
+	"github.com/coder/coder/agent/reconnectingpty"
 	"github.com/coder/coder/buildinfo"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/gitauth"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/codersdk/agentsdk"
-	"github.com/coder/coder/pty"
 	"github.com/coder/coder/tailnet"
 	"github.com/coder/retry"
 )
@@ -92,9 +91,6 @@ type Agent interface {
 }
 
 func New(options Options) Agent {
-	if options.ReconnectingPTYTimeout == 0 {
-		options.ReconnectingPTYTimeout = 5 * time.Minute
-	}
 	if options.Filesystem == nil {
 		options.Filesystem = afero.NewOsFs()
 	}
@@ -1075,8 +1071,8 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 	defer a.connCountReconnectingPTY.Add(-1)
 
 	connectionID := uuid.NewString()
-	logger = logger.With(slog.F("message_id", msg.ID), slog.F("connection_id", connectionID))
-	logger.Debug(ctx, "starting handler")
+	connLogger := logger.With(slog.F("message_id", msg.ID), slog.F("connection_id", connectionID))
+	connLogger.Debug(ctx, "starting handler")
 
 	defer func() {
 		if err := retErr; err != nil {
@@ -1087,22 +1083,22 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 			// If the agent is closed, we don't want to
 			// log this as an error since it's expected.
 			if closed {
-				logger.Debug(ctx, "reconnecting PTY failed with session error (agent closed)", slog.Error(err))
+				connLogger.Debug(ctx, "reconnecting pty failed with attach error (agent closed)", slog.Error(err))
 			} else {
-				logger.Error(ctx, "reconnecting PTY failed with session error", slog.Error(err))
+				connLogger.Error(ctx, "reconnecting pty failed with attach error", slog.Error(err))
 			}
 		}
-		logger.Debug(ctx, "session closed")
+		connLogger.Debug(ctx, "reconnecting pty connection closed")
 	}()
 
-	var rpty *reconnectingPTY
-	sendConnected := make(chan *reconnectingPTY, 1)
+	var rpty reconnectingpty.ReconnectingPTY
+	sendConnected := make(chan reconnectingpty.ReconnectingPTY, 1)
 	// On store, reserve this ID to prevent multiple concurrent new connections.
 	waitReady, ok := a.reconnectingPTYs.LoadOrStore(msg.ID, sendConnected)
 	if ok {
 		close(sendConnected) // Unused.
-		logger.Debug(ctx, "connecting to existing session")
-		c, ok := waitReady.(chan *reconnectingPTY)
+		connLogger.Debug(ctx, "connecting to existing reconnecting pty")
+		c, ok := waitReady.(chan reconnectingpty.ReconnectingPTY)
 		if !ok {
 			return xerrors.Errorf("found invalid type in reconnecting pty map: %T", waitReady)
 		}
@@ -1112,7 +1108,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 		}
 		c <- rpty // Put it back for the next reconnect.
 	} else {
-		logger.Debug(ctx, "creating new session")
+		connLogger.Debug(ctx, "creating new reconnecting pty")
 
 		connected := false
 		defer func() {
@@ -1128,169 +1124,24 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 			a.metrics.reconnectingPTYErrors.WithLabelValues("create_command").Add(1)
 			return xerrors.Errorf("create command: %w", err)
 		}
-		cmd.Env = append(cmd.Env, "TERM=xterm-256color")
-
-		// Default to buffer 64KiB.
-		circularBuffer, err := circbuf.NewBuffer(64 << 10)
-		if err != nil {
-			return xerrors.Errorf("create circular buffer: %w", err)
-		}
 
-		ptty, process, err := pty.Start(cmd)
-		if err != nil {
-			a.metrics.reconnectingPTYErrors.WithLabelValues("start_command").Add(1)
-			return xerrors.Errorf("start command: %w", err)
-		}
+		rpty = reconnectingpty.New(ctx, cmd, &reconnectingpty.Options{
+			Timeout: a.reconnectingPTYTimeout,
+			Metrics: a.metrics.reconnectingPTYErrors,
+		}, logger.With(slog.F("message_id", msg.ID)))
 
-		ctx, cancel := context.WithCancel(ctx)
-		rpty = &reconnectingPTY{
-			activeConns: map[string]net.Conn{
-				// We have to put the connection in the map instantly otherwise
-				// the connection won't be closed if the process instantly dies.
-				connectionID: conn,
-			},
-			ptty: ptty,
-			// Timeouts created with an after func can be reset!
-			timeout:        time.AfterFunc(a.reconnectingPTYTimeout, cancel),
-			circularBuffer: circularBuffer,
-		}
-		// We don't need to separately monitor for the process exiting.
-		// When it exits, our ptty.OutputReader() will return EOF after
-		// reading all process output.
 		if err = a.trackConnGoroutine(func() {
-			buffer := make([]byte, 1024)
-			for {
-				read, err := rpty.ptty.OutputReader().Read(buffer)
-				if err != nil {
-					// When the PTY is closed, this is triggered.
-					// Error is typically a benign EOF, so only log for debugging.
-					if errors.Is(err, io.EOF) {
-						logger.Debug(ctx, "unable to read pty output, command might have exited", slog.Error(err))
-					} else {
-						logger.Warn(ctx, "unable to read pty output, command might have exited", slog.Error(err))
-						a.metrics.reconnectingPTYErrors.WithLabelValues("output_reader").Add(1)
-					}
-					break
-				}
-				part := buffer[:read]
-				rpty.circularBufferMutex.Lock()
-				_, err = rpty.circularBuffer.Write(part)
-				rpty.circularBufferMutex.Unlock()
-				if err != nil {
-					logger.Error(ctx, "write to circular buffer", slog.Error(err))
-					break
-				}
-				rpty.activeConnsMutex.Lock()
-				for cid, conn := range rpty.activeConns {
-					_, err = conn.Write(part)
-					if err != nil {
-						logger.Warn(ctx,
-							"error writing to active conn",
-							slog.F("other_conn_id", cid),
-							slog.Error(err),
-						)
-						a.metrics.reconnectingPTYErrors.WithLabelValues("write").Add(1)
-					}
-				}
-				rpty.activeConnsMutex.Unlock()
-			}
-
-			// Cleanup the process, PTY, and delete it's
-			// ID from memory.
-			_ = process.Kill()
-			rpty.Close()
+			rpty.Wait()
 			a.reconnectingPTYs.Delete(msg.ID)
 		}); err != nil {
-			_ = process.Kill()
-			_ = ptty.Close()
+			rpty.Close(err.Error())
 			return xerrors.Errorf("start routine: %w", err)
 		}
+
 		connected = true
 		sendConnected <- rpty
 	}
-	// Resize the PTY to initial height + width.
-	err := rpty.ptty.Resize(msg.Height, msg.Width)
-	if err != nil {
-		// We can continue after this, it's not fatal!
-		logger.Error(ctx, "reconnecting PTY initial resize failed, but will continue", slog.Error(err))
-		a.metrics.reconnectingPTYErrors.WithLabelValues("resize").Add(1)
-	}
-	// Write any previously stored data for the TTY.
-	rpty.circularBufferMutex.RLock()
-	prevBuf := slices.Clone(rpty.circularBuffer.Bytes())
-	rpty.circularBufferMutex.RUnlock()
-	// Note that there is a small race here between writing buffered
-	// data and storing conn in activeConns. This is likely a very minor
-	// edge case, but we should look into ways to avoid it. Holding
-	// activeConnsMutex would be one option, but holding this mutex
-	// while also holding circularBufferMutex seems dangerous.
-	_, err = conn.Write(prevBuf)
-	if err != nil {
-		a.metrics.reconnectingPTYErrors.WithLabelValues("write").Add(1)
-		return xerrors.Errorf("write buffer to conn: %w", err)
-	}
-	// Multiple connections to the same TTY are permitted.
-	// This could easily be used for terminal sharing, but
-	// we do it because it's a nice user experience to
-	// copy/paste a terminal URL and have it _just work_.
-	rpty.activeConnsMutex.Lock()
-	rpty.activeConns[connectionID] = conn
-	rpty.activeConnsMutex.Unlock()
-	// Resetting this timeout prevents the PTY from exiting.
-	rpty.timeout.Reset(a.reconnectingPTYTimeout)
-
-	ctx, cancelFunc := context.WithCancel(ctx)
-	defer cancelFunc()
-	heartbeat := time.NewTicker(a.reconnectingPTYTimeout / 2)
-	defer heartbeat.Stop()
-	go func() {
-		// Keep updating the activity while this
-		// connection is alive!
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case <-heartbeat.C:
-			}
-			rpty.timeout.Reset(a.reconnectingPTYTimeout)
-		}
-	}()
-	defer func() {
-		// After this connection ends, remove it from
-		// the PTYs active connections. If it isn't
-		// removed, all PTY data will be sent to it.
-		rpty.activeConnsMutex.Lock()
-		delete(rpty.activeConns, connectionID)
-		rpty.activeConnsMutex.Unlock()
-	}()
-	decoder := json.NewDecoder(conn)
-	var req codersdk.ReconnectingPTYRequest
-	for {
-		err = decoder.Decode(&req)
-		if xerrors.Is(err, io.EOF) {
-			return nil
-		}
-		if err != nil {
-			logger.Warn(ctx, "reconnecting PTY failed with read error", slog.Error(err))
-			return nil
-		}
-		_, err = rpty.ptty.InputWriter().Write([]byte(req.Data))
-		if err != nil {
-			logger.Warn(ctx, "reconnecting PTY failed with write error", slog.Error(err))
-			a.metrics.reconnectingPTYErrors.WithLabelValues("input_writer").Add(1)
-			return nil
-		}
-		// Check if a resize needs to happen!
-		if req.Height == 0 || req.Width == 0 {
-			continue
-		}
-		err = rpty.ptty.Resize(req.Height, req.Width)
-		if err != nil {
-			// We can continue after this, it's not fatal!
-			logger.Error(ctx, "reconnecting PTY resize failed, but will continue", slog.Error(err))
-			a.metrics.reconnectingPTYErrors.WithLabelValues("resize").Add(1)
-		}
-	}
+	return rpty.Attach(ctx, connectionID, conn, msg.Height, msg.Width, connLogger)
 }
 
 // startReportingConnectionStats runs the connection stats reporting goroutine.
@@ -1541,31 +1392,6 @@ lifecycleWaitLoop:
 	return nil
 }
 
-type reconnectingPTY struct {
-	activeConnsMutex sync.Mutex
-	activeConns      map[string]net.Conn
-
-	circularBuffer      *circbuf.Buffer
-	circularBufferMutex sync.RWMutex
-	timeout             *time.Timer
-	ptty                pty.PTYCmd
-}
-
-// Close ends all connections to the reconnecting
-// PTY and clear the circular buffer.
-func (r *reconnectingPTY) Close() {
-	r.activeConnsMutex.Lock()
-	defer r.activeConnsMutex.Unlock()
-	for _, conn := range r.activeConns {
-		_ = conn.Close()
-	}
-	_ = r.ptty.Close()
-	r.circularBufferMutex.Lock()
-	r.circularBuffer.Reset()
-	r.circularBufferMutex.Unlock()
-	r.timeout.Stop()
-}
-
 // userHomeDir returns the home directory of the current user, giving
 // priority to the $HOME environment variable.
 func userHomeDir() (string, error) {