diff --git a/agent/connection_manager.go b/agent/connection_manager.go index 8152817ee..3527c796c 100644 --- a/agent/connection_manager.go +++ b/agent/connection_manager.go @@ -8,6 +8,7 @@ import ( "os" "os/signal" "strings" + "sync" "syscall" "time" @@ -20,7 +21,10 @@ import ( // It handles both WebSocket and SSH connections, automatically switching between // them based on availability and managing reconnection attempts. type ConnectionManager struct { - agent *Agent // Reference to the parent agent + agent *Agent // Reference to the parent agent + // mu guards State and isConnecting, which are read and written from both + // the main event loop and the goroutine spawned by connect(). + mu sync.Mutex State ConnectionState // Current connection state eventChan chan ConnectionEvent // Channel for connection events wsClient *WebSocketClient // WebSocket client for hub communication @@ -78,6 +82,29 @@ func (c *ConnectionManager) stopWsTicker() { } } +// getState returns the current connection state. +func (c *ConnectionManager) getState() ConnectionState { + c.mu.Lock() + defer c.mu.Unlock() + return c.State +} + +// setConnecting sets the isConnecting flag and reports its previous value. +func (c *ConnectionManager) setConnecting(v bool) (previous bool) { + c.mu.Lock() + defer c.mu.Unlock() + previous = c.isConnecting + c.isConnecting = v + return previous +} + +// isConnectingNow reports whether a reconnection attempt is currently in flight. +func (c *ConnectionManager) isConnectingNow() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.isConnecting +} + // Start begins connection attempts and enters the main event loop. // It handles connection events, periodic health updates, and graceful shutdown. func (c *ConnectionManager) Start(serverOptions ServerOptions) error { @@ -122,7 +149,10 @@ func (c *ConnectionManager) Start(serverOptions ServerOptions) error { case connectionEvent := <-c.eventChan: c.handleEvent(connectionEvent) case <-c.wsTicker.C: - _ = c.startWebSocketConnection() + // skip if connect() is still running its own attempt + if !c.isConnectingNow() { + _ = c.startWebSocketConnection() + } case <-healthTicker: _ = health.Update() case <-sigCtx.Done(): @@ -165,15 +195,15 @@ func (c *ConnectionManager) handleEvent(event ConnectionEvent) { case WebSocketConnect: c.handleStateChange(WebSocketConnected) case SSHConnect: - if c.State == Disconnected { + if c.getState() == Disconnected { c.handleStateChange(SSHConnected) } case WebSocketDisconnect: - if c.State == WebSocketConnected { + if c.getState() == WebSocketConnected { c.handleStateChange(Disconnected) } case SSHDisconnect: - if c.State == SSHConnected { + if c.getState() == SSHConnected { c.handleStateChange(Disconnected) } } @@ -182,30 +212,40 @@ func (c *ConnectionManager) handleEvent(event ConnectionEvent) { // handleStateChange updates the connection state and performs necessary actions // based on the new state, including stopping services and initiating reconnections. func (c *ConnectionManager) handleStateChange(newState ConnectionState) { + c.mu.Lock() if c.State == newState { + c.mu.Unlock() return } c.State = newState + c.mu.Unlock() + switch newState { case WebSocketConnected: slog.Info("WebSocket connected", "host", c.wsClient.hubURL.Host) c.ConnectionType = system.ConnectionTypeWebSocket c.stopWsTicker() _ = c.agent.StopServer() - c.isConnecting = false + c.setConnecting(false) case SSHConnected: // stop new ws connection attempts slog.Info("SSH connection established") c.ConnectionType = system.ConnectionTypeSSH c.stopWsTicker() - c.isConnecting = false + c.setConnecting(false) case Disconnected: c.ConnectionType = system.ConnectionTypeNone - if c.isConnecting { + // Always keep the ticker running while disconnected. A pending WebSocket + // handshake started by connect() can fail asynchronously (e.g. the hub + // closes the socket, or the deadline set in OnOpen expires) after + // connect() has already returned with a nil error, in which case the + // ticker would otherwise never get re-armed and the agent would stop + // retrying entirely (#2326). + c.startWsTicker() + if c.setConnecting(true) { // Already handling reconnection, avoid duplicate attempts return } - c.isConnecting = true slog.Warn("Disconnected from hub") // make sure old ws connection is closed c.closeWebSocket() @@ -217,10 +257,8 @@ func (c *ConnectionManager) handleStateChange(newState ConnectionState) { // connect handles the connection logic with proper delays and priority. // It attempts WebSocket connection first, falling back to SSH server if needed. func (c *ConnectionManager) connect() { - c.isConnecting = true - defer func() { - c.isConnecting = false - }() + c.setConnecting(true) + defer c.setConnecting(false) if c.wsClient != nil && time.Since(c.wsClient.lastConnectAttempt) < 5*time.Second { time.Sleep(5 * time.Second) @@ -234,7 +272,7 @@ func (c *ConnectionManager) connect() { _ = c.stop() os.Exit(1) } - if c.State == Disconnected { + if c.getState() == Disconnected { c.startSSHServer() c.startWsTicker() } @@ -243,7 +281,7 @@ func (c *ConnectionManager) connect() { // startWebSocketConnection attempts to establish a WebSocket connection to the hub. func (c *ConnectionManager) startWebSocketConnection() error { - if c.State != Disconnected { + if c.getState() != Disconnected { return errors.New("already connected") } if c.wsClient == nil { @@ -263,7 +301,7 @@ func (c *ConnectionManager) startWebSocketConnection() error { // startSSHServer starts the SSH server if the agent is currently disconnected. func (c *ConnectionManager) startSSHServer() { - if c.State == Disconnected { + if c.getState() == Disconnected { go c.agent.StartServer(c.serverOptions) } } diff --git a/agent/connection_manager_test.go b/agent/connection_manager_test.go index ef9222bde..4a00c7d6e 100644 --- a/agent/connection_manager_test.go +++ b/agent/connection_manager_test.go @@ -9,6 +9,7 @@ import ( "net" "net/url" "testing" + "testing/synctest" "time" "github.com/stretchr/testify/assert" @@ -77,6 +78,10 @@ func TestConnectionManager_StateTransitions(t *testing.T) { cm.handleStateChange(SSHConnected) assert.Equal(t, SSHConnected, cm.State, "State should change to SSHConnected") + // Prevent handleStateChange from spawning its async reconnect goroutine: + // this test only checks the synchronous state machine, and the goroutine + // would otherwise race with the direct field writes below. + cm.setConnecting(true) cm.handleStateChange(Disconnected) assert.Equal(t, Disconnected, cm.State, "State should change to Disconnected") @@ -95,7 +100,6 @@ func TestConnectionManager_EventHandling(t *testing.T) { Host: "localhost:8080", }, } - testCases := []struct { name string initialState ConnectionState @@ -148,6 +152,11 @@ func TestConnectionManager_EventHandling(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { + // Prevent handleStateChange from spawning its async reconnect + // goroutine: this test only checks the synchronous state machine, + // and the goroutine would otherwise race with the direct field + // writes here and in later subtests. + cm.setConnecting(true) cm.State = tc.initialState cm.handleEvent(tc.event) assert.Equal(t, tc.expectedState, cm.State, "State should match expected after event") @@ -221,12 +230,56 @@ func TestConnectionManager_ReconnectionLogic(t *testing.T) { // Test that isConnecting flag prevents duplicate reconnection attempts // Start from connected state, then simulate disconnect cm.State = WebSocketConnected - cm.isConnecting = false + cm.setConnecting(false) // First disconnect should trigger reconnection logic cm.handleStateChange(Disconnected) assert.Equal(t, Disconnected, cm.State, "Should change to disconnected") - assert.True(t, cm.isConnecting, "Should set isConnecting flag") + assert.True(t, cm.isConnectingNow(), "Should set isConnecting flag") +} + +// TestConnectionManager_TickerSurvivesStaleDisconnect reproduces the freeze from +// https://github.com/henrygd/beszel/issues/2326: a reconnect attempt's handshake +// can fail asynchronously (after connect() already returned with a nil error) +// while the manager is still in the Disconnected state. Previously the ticker +// was only re-armed from connect()'s synchronous error branch, so once that +// window was missed, the agent stopped retrying forever. The ticker must keep +// running any time the manager transitions into Disconnected, regardless of +// what happens to the in-flight handshake afterwards. +func TestConnectionManager_TickerSurvivesStaleDisconnect(t *testing.T) { + agent := createTestAgent(t) + cm := agent.connectionManager + cm.eventChan = make(chan ConnectionEvent, 1) + + // Run on synctest's fake clock so the ticker fires without waiting a real + // wsTickerInterval. The ticker must be created inside the bubble. + synctest.Test(t, func(t *testing.T) { + // Simulate a healthy WebSocket connection, then a disconnect - mirroring + // handleStateChange's own Disconnected branch, but without launching the + // real async connect() goroutine so the ticker state can be asserted + // deterministically. + cm.State = WebSocketConnected + cm.stopWsTicker() + cm.setConnecting(true) + cm.handleStateChange(Disconnected) + require.NotNil(t, cm.wsTicker, "ticker must be armed as soon as the manager becomes Disconnected") + defer cm.stopWsTicker() + + // Now simulate connect()'s in-flight handshake dying asynchronously with the + // manager still Disconnected (e.g. a late OnClose on an unauthenticated + // connection). This event is dropped by handleEvent since State is not + // WebSocketConnected, but the ticker armed above must still be running so + // the manager keeps retrying. + cm.setConnecting(false) + cm.handleEvent(WebSocketDisconnect) + assert.Equal(t, Disconnected, cm.State) + + select { + case <-cm.wsTicker.C: + case <-time.After(wsTickerInterval + 2*time.Second): + t.Fatal("ticker did not fire after a stale disconnect event - agent would freeze forever") + } + }) } // TestConnectionManager_ConnectWithRateLimit tests connection rate limiting