fix issue with writing getting caught up under very rare conditions

Overlord-Client/cmd/agent/handlers/ping.go

@@ -9,7 +9,6 @@ import (
 )
 
 func HandlePing(ctx context.Context, env *runtime.Env, envelope map[string]interface{}) error {
-	env.SetLastPong(time.Now().UnixMilli())
 
 	ts := extractTimestamp(envelope["ts"])
 	if ts == 0 {

Overlord-Client/cmd/agent/session.go

@@ -589,6 +589,7 @@ func runSession(ctx context.Context, cancel context.CancelFunc, conn *websocket.
 					last := env.LastPong()
 					if !last.IsZero() && time.Since(last) > grace {
 						log.Printf("ping: no pong for %s, forcing reconnect", time.Since(last))
+						conn.Close(websocket.StatusGoingAway, "pong timeout")
 						cancel()
 						return
 					}

Overlord-Client/cmd/agent/wire/safe_writer.go

@@ -2,22 +2,33 @@ package wire
 
 import (
 	"context"
-	"sync"
+	"time"
 
 	"nhooyr.io/websocket"
 )
 
+const writeTimeout = 30 * time.Second
+
 type SafeWriter struct {
-	mu sync.Mutex
+	sem chan struct{}
-	w  Writer
+	w   Writer
 }
 
 func NewSafeWriter(w Writer) *SafeWriter {
-	return &SafeWriter{w: w}
+	sem := make(chan struct{}, 1)
+	sem <- struct{}{}
+	return &SafeWriter{sem: sem, w: w}
 }
 
 func (s *SafeWriter) Write(ctx context.Context, messageType websocket.MessageType, p []byte) error {
-	s.mu.Lock()
+	select {
-	defer s.mu.Unlock()
+	case <-s.sem:
-	return s.w.Write(ctx, messageType, p)
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+	defer func() { s.sem <- struct{}{} }()
+
+	writeCtx, cancel := context.WithTimeout(ctx, writeTimeout)
+	defer cancel()
+	return s.w.Write(writeCtx, messageType, p)
 }