@@ -56,7 +56,12 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) error {
5656bytesPerTick = r .cfg .BytesPerSecond / r .cfg .TicksPerSecond
5757)
5858
59+ // Set a deadline for stopping the text.
60+ start := time .Now ()
61+ deadlineCtx ,cancel := context .WithDeadline (ctx ,start .Add (r .cfg .Duration ))
62+ defer cancel ()
5963logger .Debug (ctx ,"connect to workspace agent" ,slog .F ("agent_id" ,agentID ))
64+
6065conn ,err := r .client .WorkspaceAgentReconnectingPTY (ctx , codersdk.WorkspaceAgentReconnectingPTYOpts {
6166AgentID :agentID ,
6267Reconnect :reconnect ,
@@ -74,11 +79,6 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) error {
7479_ = conn .Close ()
7580}()
7681
77- // Set a deadline for stopping the text.
78- start := time .Now ()
79- deadlineCtx ,cancel := context .WithDeadline (ctx ,start .Add (r .cfg .Duration ))
80- defer cancel ()
81-
8282// Wrap the conn in a countReadWriter so we can monitor bytes sent/rcvd.
8383crw := countReadWriter {ReadWriter :conn ,ctx :deadlineCtx }
8484
@@ -112,7 +112,7 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) error {
112112close (wch )
113113}()
114114
115- //Wait for both our reads and writes to be finished .
115+ //Write until the context is canceled .
116116if wErr := <- wch ;wErr != nil {
117117return xerrors .Errorf ("write to pty: %w" ,wErr )
118118}
@@ -138,7 +138,7 @@ func (*Runner) Cleanup(context.Context, string) error {
138138
139139// drainContext drains from src until it returns io.EOF or ctx times out.
140140func drainContext (ctx context.Context ,src io.Reader ,bufSize int64 )error {
141- errCh := make (chan error )
141+ errCh := make (chan error , 1 )
142142done := make (chan struct {})
143143go func () {
144144tmp := make ([]byte ,bufSize )
@@ -149,6 +149,9 @@ func drainContext(ctx context.Context, src io.Reader, bufSize int64) error {
149149return
150150default :
151151_ ,err := io .CopyN (buf ,src ,1 )
152+ if ctx .Err ()!= nil {
153+ return // context canceled while we were copying.
154+ }
152155if err != nil {
153156errCh <- err
154157close (errCh )