@@ -11,8 +11,10 @@ import (
1111"net/http"
1212"net/netip"
1313"net/url"
14+ "runtime/pprof"
1415"strconv"
1516"strings"
17+ "sync"
1618"time"
1719
1820"github.com/google/uuid"
@@ -291,11 +293,12 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
291293})
292294return
293295}
294- go httpapi .Heartbeat (ctx ,conn )
295296
296- _ ,wsNetConn := websocketNetConn (ctx ,conn ,websocket .MessageBinary )
297+ ctx ,wsNetConn := websocketNetConn (ctx ,conn ,websocket .MessageBinary )
297298defer wsNetConn .Close ()// Also closes conn.
298299
300+ go httpapi .Heartbeat (ctx ,conn )
301+
299302agentConn ,release ,err := api .workspaceAgentCache .Acquire (r ,workspaceAgent .ID )
300303if err != nil {
301304_ = conn .Close (websocket .StatusInternalError ,httpapi .WebsocketCloseSprintf ("dial workspace agent: %s" ,err ))
@@ -606,11 +609,40 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
606609})
607610return
608611}
609- go httpapi .Heartbeat (ctx ,conn )
610612
611613ctx ,wsNetConn := websocketNetConn (ctx ,conn ,websocket .MessageBinary )
612614defer wsNetConn .Close ()
613615
616+ // We use a custom heartbeat routine here instead of `httpapi.Heartbeat`
617+ // because we want to log the agent's last ping time.
618+ var lastPing time.Time
619+ var pingMu sync.Mutex
620+ go pprof .Do (ctx ,pprof .Labels ("agent" ,workspaceAgent .ID .String ()),func (ctx context.Context ) {
621+ // TODO(mafredri): Is this too frequent? Use separate ping disconnect timeout?
622+ t := time .NewTicker (api .AgentConnectionUpdateFrequency )
623+ defer t .Stop ()
624+
625+ for {
626+ select {
627+ case <- t .C :
628+ case <- ctx .Done ():
629+ return
630+ }
631+
632+ // We don't need a context that times out here because the ping will
633+ // eventually go through. If the context times out, then other
634+ // websocket read operations will receive an error, obfuscating the
635+ // actual problem.
636+ err := conn .Ping (ctx )
637+ if err != nil {
638+ return
639+ }
640+ pingMu .Lock ()
641+ lastPing = time .Now ()
642+ pingMu .Unlock ()
643+ }
644+ })
645+
614646firstConnectedAt := workspaceAgent .FirstConnectedAt
615647if ! firstConnectedAt .Valid {
616648firstConnectedAt = sql.NullTime {
@@ -654,9 +686,12 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
654686ctx ,cancel := context .WithTimeout (dbauthz .AsSystemRestricted (api .ctx ),api .AgentInactiveDisconnectTimeout )
655687defer cancel ()
656688
657- disconnectedAt = sql.NullTime {
658- Time :database .Now (),
659- Valid :true ,
689+ // Only update timestamp if the disconnect is new.
690+ if ! disconnectedAt .Valid {
691+ disconnectedAt = sql.NullTime {
692+ Time :database .Now (),
693+ Valid :true ,
694+ }
660695}
661696err := updateConnectionTimes (ctx )
662697if err != nil {
@@ -711,15 +746,37 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
711746return
712747case <- ticker .C :
713748}
714- lastConnectedAt = sql.NullTime {
715- Time :database .Now (),
716- Valid :true ,
749+
750+ pingMu .Lock ()
751+ lastPing := lastPing
752+ pingMu .Unlock ()
753+
754+ var connectionStatusChanged bool
755+ if time .Since (lastPing )> api .AgentInactiveDisconnectTimeout {
756+ if ! disconnectedAt .Valid {
757+ connectionStatusChanged = true
758+ disconnectedAt = sql.NullTime {
759+ Time :database .Now (),
760+ Valid :true ,
761+ }
762+ }
763+ }else {
764+ connectionStatusChanged = disconnectedAt .Valid
765+ // TODO(mafredri): Should we update it here or allow lastConnectedAt to shadow it?
766+ disconnectedAt = sql.NullTime {}
767+ lastConnectedAt = sql.NullTime {
768+ Time :database .Now (),
769+ Valid :true ,
770+ }
717771}
718772err = updateConnectionTimes (ctx )
719773if err != nil {
720774_ = conn .Close (websocket .StatusGoingAway ,err .Error ())
721775return
722776}
777+ if connectionStatusChanged {
778+ api .publishWorkspaceUpdate (ctx ,build .WorkspaceID )
779+ }
723780err := ensureLatestBuild ()
724781if err != nil {
725782// Disconnect agents that are no longer valid.