@@ -30,6 +30,7 @@ import (
3030 "github.com/spf13/afero"
3131 "go.uber.org/atomic"
3232 gossh "golang.org/x/crypto/ssh"
33+ "golang.org/x/exp/slices"
3334 "golang.org/x/xerrors"
3435 "tailscale.com/net/speedtest"
3536 "tailscale.com/tailcfg"
@@ -90,7 +91,7 @@ func New(options Options) io.Closer {
9091 }
9192 }
9293 ctx , cancelFunc := context .WithCancel (context .Background ())
93- server := & agent {
94+ a := & agent {
9495 reconnectingPTYTimeout : options .ReconnectingPTYTimeout ,
9596 logger : options .Logger ,
9697 closeCancel : cancelFunc ,
@@ -101,8 +102,8 @@ func New(options Options) io.Closer {
101102 filesystem : options .Filesystem ,
102103 tempDir : options .TempDir ,
103104 }
104- server .init (ctx )
105- return server
105+ a .init (ctx )
106+ return a
106107}
107108
108109type agent struct {
@@ -225,6 +226,25 @@ func (a *agent) run(ctx context.Context) error {
225226 _ = network .Close ()
226227 return xerrors .New ("agent is closed" )
227228 }
229+
230+ // Report statistics from the created network.
231+ cl , err := a .client .AgentReportStats (ctx , a .logger , func () * codersdk.AgentStats {
232+ stats := network .ExtractTrafficStats ()
233+ return convertAgentStats (stats )
234+ })
235+ if err != nil {
236+ a .logger .Error (ctx , "report stats" , slog .Error (err ))
237+ } else {
238+ if err = a .trackConnGoroutine (func () {
239+ // This is OK because the agent never re-creates the tailnet
240+ // and the only shutdown indicator is agent.Close().
241+ <- a .closed
242+ _ = cl .Close ()
243+ }); err != nil {
244+ a .logger .Debug (ctx , "report stats goroutine" , slog .Error (err ))
245+ _ = cl .Close ()
246+ }
247+ }
228248 } else {
229249 // Update the DERP map!
230250 network .SetDERPMap (metadata .DERPMap )
@@ -300,10 +320,12 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_
300320 }
301321 }()
302322 if err = a .trackConnGoroutine (func () {
323+ logger := a .logger .Named ("reconnecting-pty" )
324+
303325 for {
304326 conn , err := reconnectingPTYListener .Accept ()
305327 if err != nil {
306- a . logger .Debug (ctx , "accept pty failed" , slog .Error (err ))
328+ logger .Debug (ctx , "accept pty failed" , slog .Error (err ))
307329 return
308330 }
309331 // This cannot use a JSON decoder, since that can
@@ -324,7 +346,9 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_
324346 if err != nil {
325347 continue
326348 }
327- go a .handleReconnectingPTY (ctx , msg , conn )
349+ go func () {
350+ _ = a .handleReconnectingPTY (ctx , logger , msg , conn )
351+ }()
328352 }
329353 }); err != nil {
330354 return nil , err
@@ -556,28 +580,6 @@ func (a *agent) init(ctx context.Context) {
556580 }
557581
558582 go a .runLoop (ctx )
559- cl , err := a .client .AgentReportStats (ctx , a .logger , func () * codersdk.AgentStats {
560- stats := map [netlogtype.Connection ]netlogtype.Counts {}
561- a .closeMutex .Lock ()
562- if a .network != nil {
563- stats = a .network .ExtractTrafficStats ()
564- }
565- a .closeMutex .Unlock ()
566- return convertAgentStats (stats )
567- })
568- if err != nil {
569- a .logger .Error (ctx , "report stats" , slog .Error (err ))
570- return
571- }
572-
573- if err = a .trackConnGoroutine (func () {
574- <- a .closed
575- _ = cl .Close ()
576- }); err != nil {
577- a .logger .Error (ctx , "report stats goroutine" , slog .Error (err ))
578- _ = cl .Close ()
579- return
580- }
581583}
582584
583585func convertAgentStats (counts map [netlogtype.Connection ]netlogtype.Counts ) * codersdk.AgentStats {
@@ -798,38 +800,56 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
798800 return cmd .Wait ()
799801}
800802
801- func (a * agent ) handleReconnectingPTY (ctx context.Context , msg codersdk.ReconnectingPTYInit , conn net.Conn ) {
803+ func (a * agent ) handleReconnectingPTY (ctx context.Context , logger slog. Logger , msg codersdk.ReconnectingPTYInit , conn net.Conn ) ( retErr error ) {
802804 defer conn .Close ()
803805
804806 connectionID := uuid .NewString ()
807+ logger = logger .With (slog .F ("id" , msg .ID ), slog .F ("connection_id" , connectionID ))
808+
809+ defer func () {
810+ if err := retErr ; err != nil {
811+ a .closeMutex .Lock ()
812+ closed := a .isClosed ()
813+ a .closeMutex .Unlock ()
814+
815+ // If the agent is closed, we don't want to
816+ // log this as an error since it's expected.
817+ if closed {
818+ logger .Debug (ctx , "session error after agent close" , slog .Error (err ))
819+ } else {
820+ logger .Error (ctx , "session error" , slog .Error (err ))
821+ }
822+ }
823+ logger .Debug (ctx , "session closed" )
824+ }()
825+
805826 var rpty * reconnectingPTY
806827 rawRPTY , ok := a .reconnectingPTYs .Load (msg .ID )
807828 if ok {
829+ logger .Debug (ctx , "connecting to existing session" )
808830 rpty , ok = rawRPTY .(* reconnectingPTY )
809831 if ! ok {
810- a .logger .Error (ctx , "found invalid type in reconnecting pty map" , slog .F ("id" , msg .ID ))
811- return
832+ return xerrors .Errorf ("found invalid type in reconnecting pty map: %T" , rawRPTY )
812833 }
813834 } else {
835+ logger .Debug (ctx , "creating new session" )
836+
814837 // Empty command will default to the users shell!
815838 cmd , err := a .createCommand (ctx , msg .Command , nil )
816839 if err != nil {
817- a .logger .Error (ctx , "create reconnecting pty command" , slog .Error (err ))
818- return
840+ return xerrors .Errorf ("create command: %w" , err )
819841 }
820842 cmd .Env = append (cmd .Env , "TERM=xterm-256color" )
821843
822844 // Default to buffer 64KiB.
823845 circularBuffer , err := circbuf .NewBuffer (64 << 10 )
824846 if err != nil {
825- a .logger .Error (ctx , "create circular buffer" , slog .Error (err ))
826- return
847+ return xerrors .Errorf ("create circular buffer: %w" , err )
827848 }
828849
829850 ptty , process , err := pty .Start (cmd )
830851 if err != nil {
831- a .logger .Error (ctx , "start reconnecting pty command" , slog .F ("id" , msg .ID ), slog .Error (err ))
832- return
852+ return xerrors .Errorf ("start command: %w" , err )
833853 }
834854
835855 ctx , cancelFunc := context .WithCancel (ctx )
@@ -873,7 +893,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
873893 _ , err = rpty .circularBuffer .Write (part )
874894 rpty .circularBufferMutex .Unlock ()
875895 if err != nil {
876- a . logger .Error (ctx , "reconnecting pty write buffer" , slog .Error (err ), slog . F ( "id" , msg . ID ))
896+ logger .Error (ctx , "write to circular buffer" , slog .Error (err ))
877897 break
878898 }
879899 rpty .activeConnsMutex .Lock ()
@@ -889,23 +909,27 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
889909 rpty .Close ()
890910 a .reconnectingPTYs .Delete (msg .ID )
891911 }); err != nil {
892- a .logger .Error (ctx , "start reconnecting pty routine" , slog .F ("id" , msg .ID ), slog .Error (err ))
893- return
912+ return xerrors .Errorf ("start routine: %w" , err )
894913 }
895914 }
896915 // Resize the PTY to initial height + width.
897916 err := rpty .ptty .Resize (msg .Height , msg .Width )
898917 if err != nil {
899918 // We can continue after this, it's not fatal!
900- a . logger .Error (ctx , "resize reconnecting pty" , slog . F ( "id" , msg . ID ) , slog .Error (err ))
919+ logger .Error (ctx , "resize" , slog .Error (err ))
901920 }
902921 // Write any previously stored data for the TTY.
903922 rpty .circularBufferMutex .RLock ()
904- _ , err = conn . Write (rpty .circularBuffer .Bytes ())
923+ prevBuf := slices . Clone (rpty .circularBuffer .Bytes ())
905924 rpty .circularBufferMutex .RUnlock ()
925+ // Note that there is a small race here between writing buffered
926+ // data and storing conn in activeConns. This is likely a very minor
927+ // edge case, but we should look into ways to avoid it. Holding
928+ // activeConnsMutex would be one option, but holding this mutex
929+ // while also holding circularBufferMutex seems dangerous.
930+ _ , err = conn .Write (prevBuf )
906931 if err != nil {
907- a .logger .Warn (ctx , "write reconnecting pty buffer" , slog .F ("id" , msg .ID ), slog .Error (err ))
908- return
932+ return xerrors .Errorf ("write buffer to conn: %w" , err )
909933 }
910934 // Multiple connections to the same TTY are permitted.
911935 // This could easily be used for terminal sharing, but
@@ -946,16 +970,16 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
946970 for {
947971 err = decoder .Decode (& req )
948972 if xerrors .Is (err , io .EOF ) {
949- return
973+ return nil
950974 }
951975 if err != nil {
952- a . logger .Warn (ctx , "reconnecting pty buffer read error" , slog . F ( "id" , msg . ID ) , slog .Error (err ))
953- return
976+ logger .Warn (ctx , "read conn" , slog .Error (err ))
977+ return nil
954978 }
955979 _ , err = rpty .ptty .Input ().Write ([]byte (req .Data ))
956980 if err != nil {
957- a . logger .Warn (ctx , "write to reconnecting pty" , slog . F ( "id" , msg . ID ) , slog .Error (err ))
958- return
981+ logger .Warn (ctx , "write to pty" , slog .Error (err ))
982+ return nil
959983 }
960984 // Check if a resize needs to happen!
961985 if req .Height == 0 || req .Width == 0 {
@@ -964,7 +988,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
964988 err = rpty .ptty .Resize (req .Height , req .Width )
965989 if err != nil {
966990 // We can continue after this, it's not fatal!
967- a . logger .Error (ctx , "resize reconnecting pty" , slog . F ( "id" , msg . ID ) , slog .Error (err ))
991+ logger .Error (ctx , "resize" , slog .Error (err ))
968992 }
969993 }
970994}
0 commit comments