diff --git a/scaletest/workspacetraffic/conn.go b/scaletest/workspacetraffic/conn.go index fd9bf93866..c6526dd172 100644 --- a/scaletest/workspacetraffic/conn.go +++ b/scaletest/workspacetraffic/conn.go @@ -21,8 +21,12 @@ import ( const ( // Set a timeout for graceful close of the connection. connCloseTimeout = 30 * time.Second + // Set a timeout for the read to unblock after a force close. Closing the + // connection unblocks a pending read, so this should never be hit unless + // the underlying connection misbehaves. + forceCloseReadTimeout = 5 * time.Second // Set a timeout for waiting for the connection to close. - waitCloseTimeout = connCloseTimeout + 5*time.Second + waitCloseTimeout = connCloseTimeout + forceCloseReadTimeout + 5*time.Second // In theory, we can send larger payloads to push bandwidth, but we need to // be careful not to send too much data at once or the server will close the @@ -53,10 +57,26 @@ func connectRPTY(ctx context.Context, client *codersdk.Client, agentID, reconnec return &crw, nil } +// errRPTYGracefulCloseTimeout indicates the server did not close the +// connection after Ctrl+C was sent and the connection was force closed +// instead. The connection is fully closed when this error is returned, so +// callers may treat it as a non-fatal warning. +var errRPTYGracefulCloseTimeout = xerrors.New("graceful close timed out, connection was force closed") + type rptyConn struct { conn io.ReadWriteCloser wenc *json.Encoder + // Both timeouts default to the package constants and are overridden + // only in tests. + // + // closeTimeout limits how long Close waits for the server to close the + // connection after Ctrl+C is sent. + closeTimeout time.Duration + // forceCloseReadTimeout limits how long Close waits for the read to + // unblock after the connection is force closed. + forceCloseReadTimeout time.Duration + readOnce sync.Once readErr chan error @@ -64,11 +84,16 @@ type rptyConn struct { closed bool } +// newPTYConn wraps conn for reconnecting PTY traffic. The caller must keep +// an active Read loop on the returned conn; Close waits for a read to +// observe the connection closing and will time out without one. func newPTYConn(conn io.ReadWriteCloser) *rptyConn { rc := &rptyConn{ - conn: conn, - wenc: json.NewEncoder(conn), - readErr: make(chan error, 1), + conn: conn, + wenc: json.NewEncoder(conn), + closeTimeout: connCloseTimeout, + forceCloseReadTimeout: forceCloseReadTimeout, + readErr: make(chan error, 1), } return rc } @@ -124,21 +149,52 @@ func (c *rptyConn) Close() (err error) { c.closed = true c.mu.Unlock() - defer c.conn.Close() - - // Send Ctrl+C to interrupt the command. - _, err = c.writeNoLock([]byte("\u0003")) - if err != nil { + // Send Ctrl+C to interrupt the command, giving the server a chance to + // flush remaining output and close the connection gracefully. + if _, err = c.writeNoLock([]byte("\u0003")); err != nil { + // We couldn't interrupt the command, force close the connection to + // unblock the read before returning. + if cerr := c.forceClose(); cerr != nil { + cerr = xerrors.Errorf("force close: %w", cerr) + return errors.Join(xerrors.Errorf("write ctrl+c: %w", err), cerr) + } return xerrors.Errorf("write ctrl+c: %w", err) } + + // Wait for the server to close the connection, which unblocks the read. If + // the server doesn't close in time, force close the connection ourselves. + t := time.NewTimer(c.closeTimeout) + defer t.Stop() select { - case <-time.After(connCloseTimeout): - return xerrors.Errorf("timeout waiting for read to finish") case err = <-c.readErr: + _ = c.conn.Close() if errors.Is(err, io.EOF) { return nil } return err + case <-t.C: + if err := c.forceClose(); err != nil { + return xerrors.Errorf("force close: %w", err) + } + return errRPTYGracefulCloseTimeout + } +} + +// forceClose closes the underlying connection and waits for the read to +// unblock. The read error is caused by the close, so it is expected and +// discarded. Returns an error if the read does not unblock within +// forceCloseReadTimeout, which also bounds a blocking close. +func (c *rptyConn) forceClose() error { + // Start the timer before closing so a blocking close cannot extend the + // total wait beyond forceCloseReadTimeout. + t := time.NewTimer(c.forceCloseReadTimeout) + defer t.Stop() + _ = c.conn.Close() + select { + case <-c.readErr: + return nil + case <-t.C: + return xerrors.New("timeout waiting for read to finish after close") } } diff --git a/scaletest/workspacetraffic/conn_internal_test.go b/scaletest/workspacetraffic/conn_internal_test.go new file mode 100644 index 0000000000..1903e51147 --- /dev/null +++ b/scaletest/workspacetraffic/conn_internal_test.go @@ -0,0 +1,172 @@ +package workspacetraffic + +import ( + "io" + "sync" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/testutil" +) + +// stubConn simulates the server side of a reconnecting PTY connection. +type stubConn struct { + // closeOnWrite closes the connection on the first write, simulating a + // server that closes gracefully in response to Ctrl+C. + closeOnWrite bool + // failWrites makes every write return an error, simulating a connection + // that can no longer send data. + failWrites bool + // readIgnoresClose prevents reads from unblocking when the connection is + // closed, simulating a misbehaving connection. + readIgnoresClose bool + + closeOnce sync.Once + closedCh chan struct{} + // releaseCh unblocks reads when readIgnoresClose is set, allowing the + // test to clean up the read goroutine. + releaseCh chan struct{} +} + +func newStubConn() *stubConn { + return &stubConn{ + closedCh: make(chan struct{}), + releaseCh: make(chan struct{}), + } +} + +func (s *stubConn) Read(_ []byte) (int, error) { + if s.readIgnoresClose { + <-s.releaseCh + return 0, io.EOF + } + <-s.closedCh + return 0, io.EOF +} + +func (s *stubConn) Write(p []byte) (int, error) { + if s.failWrites { + return 0, xerrors.New("write failed") + } + if s.closeOnWrite { + _ = s.Close() + } + return len(p), nil +} + +func (s *stubConn) Close() error { + s.closeOnce.Do(func() { + close(s.closedCh) + }) + return nil +} + +// startDrain reads from rc until it errors, mirroring the drain goroutine in +// Runner.Run. It returns a channel that is closed when the read finishes. +func startDrain(t *testing.T, rc *rptyConn) <-chan struct{} { + t.Helper() + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = io.Copy(io.Discard, rc) + }() + return done +} + +func waitDone(t *testing.T, done <-chan struct{}) { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + _ = testutil.TryReceive(ctx, t, done) +} + +func TestRPTYConn_Close(t *testing.T) { + t.Parallel() + + t.Run("Graceful", func(t *testing.T) { + t.Parallel() + + // The server closes the connection in response to Ctrl+C, the read + // unblocks with io.EOF and Close reports success. + stub := newStubConn() + stub.closeOnWrite = true + rc := newPTYConn(stub) + done := startDrain(t, rc) + + err := rc.Close() + require.NoError(t, err) + waitDone(t, done) + }) + + t.Run("ForceClose", func(t *testing.T) { + t.Parallel() + + // The server ignores Ctrl+C and never closes the connection. Close + // force closes the connection to unblock the read and reports a + // non-fatal graceful close timeout. + stub := newStubConn() + rc := newPTYConn(stub) + rc.closeTimeout = testutil.IntervalFast + done := startDrain(t, rc) + + err := rc.Close() + require.ErrorIs(t, err, errRPTYGracefulCloseTimeout) + waitDone(t, done) + }) + + t.Run("WriteFails", func(t *testing.T) { + t.Parallel() + + // The Ctrl+C write fails. Close force closes the connection to + // unblock the read and reports a hard error, not the non-fatal + // graceful close timeout. + stub := newStubConn() + stub.failWrites = true + rc := newPTYConn(stub) + done := startDrain(t, rc) + + err := rc.Close() + require.Error(t, err) + require.NotErrorIs(t, err, errRPTYGracefulCloseTimeout) + require.ErrorContains(t, err, "write ctrl+c") + waitDone(t, done) + }) + + t.Run("ReadStuckAfterClose", func(t *testing.T) { + t.Parallel() + + // The read doesn't unblock even after the connection is force + // closed. Close reports an error instead of blocking forever. + stub := newStubConn() + stub.readIgnoresClose = true + rc := newPTYConn(stub) + rc.closeTimeout = testutil.IntervalFast + rc.forceCloseReadTimeout = testutil.IntervalFast + done := startDrain(t, rc) + // Unblock the read goroutine at the end of the test. + t.Cleanup(func() { + close(stub.releaseCh) + waitDone(t, done) + }) + + err := rc.Close() + require.Error(t, err) + require.NotErrorIs(t, err, errRPTYGracefulCloseTimeout) + require.ErrorContains(t, err, "timeout waiting for read to finish after close") + }) + + t.Run("CloseTwice", func(t *testing.T) { + t.Parallel() + + // A second Close is a no-op and returns nil. + stub := newStubConn() + stub.closeOnWrite = true + rc := newPTYConn(stub) + done := startDrain(t, rc) + + require.NoError(t, rc.Close()) + require.NoError(t, rc.Close()) + waitDone(t, done) + }) +} diff --git a/scaletest/workspacetraffic/run.go b/scaletest/workspacetraffic/run.go index 80cb83fd43..0cb684e569 100644 --- a/scaletest/workspacetraffic/run.go +++ b/scaletest/workspacetraffic/run.go @@ -137,9 +137,15 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error) closeConn := func() error { closeOnce.Do(func() { closeErr = conn.Close() - if errors.Is(closeErr, io.EOF) { + switch { + case errors.Is(closeErr, io.EOF): closeErr = nil - } else if closeErr != nil { + case errors.Is(closeErr, errRPTYGracefulCloseTimeout): + // The connection was closed, just not gracefully. Surface it + // in the logs but don't fail the run. + logger.Warn(ctx, "close agent connection", slog.Error(closeErr)) + closeErr = nil + case closeErr != nil: logger.Error(ctx, "close agent connection", slog.Error(closeErr)) closeErr = xerrors.Errorf("close agent connection: %w", closeErr) }