diff --git a/config.go b/config.go index af91112..4f72ff9 100644 --- a/config.go +++ b/config.go @@ -15,6 +15,15 @@ type Config struct { Location LocationConfig `yaml:"location"` DatabasePath string `yaml:"databasePath"` ReconnectInterval time.Duration `yaml:"reconnectInterval"` + + // PingInterval is how often to send a websocket ping to Home Assistant to + // verify the connection is still alive. Defaults to 30s. + PingInterval time.Duration `yaml:"pingInterval"` + + // ReadTimeout is how long to wait for any data (including a pong response + // to our ping) before considering the connection dead and forcing a + // reconnect. Must be greater than PingInterval. Defaults to 60s. + ReadTimeout time.Duration `yaml:"readTimeout"` } type HomeAssistantConfig struct { diff --git a/connection.go b/connection.go index 63c3a13..251c0da 100644 --- a/connection.go +++ b/connection.go @@ -58,8 +58,10 @@ func NewConnection(cfg Config) *Connection { } api := hassws.NewClient(hassws.ClientConfig{ - Host: cfg.HomeAssistant.Host, - Token: cfg.HomeAssistant.Token, + Host: cfg.HomeAssistant.Host, + Token: cfg.HomeAssistant.Token, + PingInterval: cfg.PingInterval, + PongTimeout: cfg.ReadTimeout, }) // Set the database on the global logger diff --git a/connection_reconnection_test.go b/connection_reconnection_test.go index c6f9ba9..f2ed44a 100644 --- a/connection_reconnection_test.go +++ b/connection_reconnection_test.go @@ -361,3 +361,60 @@ func TestReconnectionAttemptCounter(t *testing.T) { waitForEventSubscription(t, server, 1) assert.Equal(t, 1, conn.GetReconnectAttempts()) } + +// TestStuckConnectionTriggersReconnect verifies that when the websocket is +// still open but no longer responsive (no data and no pong responses to our +// pings), the client detects this via the read timeout and reconnects. +func TestStuckConnectionTriggersReconnect(t *testing.T) { + conn, server, cleanup := newClientServerWithConfig(t, Config{ + DatabasePath: ":memory:", + ReconnectInterval: 100 * time.Millisecond, + PingInterval: 100 * time.Millisecond, + ReadTimeout: 300 * time.Millisecond, + }) + defer cleanup() + + testEntity := NewEntity("test.entity") + conn.RegisterEntities(testEntity) + + // Verify the connection works initially. + server.SendEvent(homeassistant.Event{ + EventData: homeassistant.EventData{ + EntityID: "test.entity", + NewState: &homeassistant.State{ + EntityID: "test.entity", + State: "on", + }, + }, + }) + + waitFor(t, "initial state update", func() bool { + return testEntity.GetState().State == "on" + }, func() { + t.Logf("Entity state: %v", testEntity.GetState()) + }) + + // Simulate a stuck connection: still open, but no data and no pong replies. + server.SimulateStuckConnection() + + // The read timeout should fire and trigger a reconnect. + waitForReconnection(t, conn, 1) + waitForEventSubscription(t, server, 1) + + // Verify the connection is healthy again. + server.SendEvent(homeassistant.Event{ + EventData: homeassistant.EventData{ + EntityID: "test.entity", + NewState: &homeassistant.State{ + EntityID: "test.entity", + State: "off", + }, + }, + }) + + waitFor(t, "state update after reconnection", func() bool { + return testEntity.GetState().State == "off" + }, func() { + t.Logf("Entity state: %v", testEntity.GetState()) + }) +} diff --git a/hassws/client.go b/hassws/client.go index 8ac2a48..63923ce 100644 --- a/hassws/client.go +++ b/hassws/client.go @@ -15,6 +15,17 @@ import ( const readTimeoutSeconds = 3 +const ( + // defaultPingInterval is how often a websocket ping is sent to Home + // Assistant to verify the connection is still alive. + defaultPingInterval = 30 * time.Second + + // defaultPongTimeout is how long we wait to receive any data (including a + // pong response to our ping) before considering the connection dead and + // forcing a reconnect. It must be greater than defaultPingInterval. + defaultPongTimeout = 60 * time.Second +) + // Connection states type connectionState int @@ -35,6 +46,11 @@ type Client struct { cfg ClientConfig conn *websocket.Conn + // writeMutex serializes writes to the websocket. gorilla/websocket does not + // support concurrent writers, and the keepAlive goroutine writes ping + // frames concurrently with other senders. + writeMutex sync.Mutex + msgID atomic.Int64 // Each request has a unique ID and any response will have the same ID. To @@ -51,9 +67,41 @@ type Client struct { type ClientConfig struct { Host string Token string + + // PingInterval is how often to send a websocket ping to verify the + // connection is alive. Defaults to defaultPingInterval. + PingInterval time.Duration + + // PongTimeout is how long to wait for any data (including a pong response) + // before considering the connection dead and forcing a reconnect. Defaults + // to defaultPongTimeout. Must be greater than PingInterval. + PongTimeout time.Duration } func NewClient(config ClientConfig) *Client { + if config.PingInterval <= 0 { + config.PingInterval = defaultPingInterval + } + + if config.PongTimeout <= 0 { + config.PongTimeout = defaultPongTimeout + } + + // The pong timeout must be greater than the ping interval, otherwise the + // read deadline can expire on a quiet-but-healthy connection before a ping + // has a chance to elicit a pong and refresh it, causing a reconnect loop. + if config.PongTimeout <= config.PingInterval { + logger.Warn( + "PongTimeout must be greater than PingInterval; adjusting", + "", + "pingInterval", config.PingInterval, + "pongTimeout", config.PongTimeout, + "adjustedPongTimeout", 2*config.PingInterval, + ) + + config.PongTimeout = 2 * config.PingInterval + } + c := &Client{ cfg: config, responses: make(map[int]chan []byte), @@ -121,6 +169,10 @@ func (c *Client) authenticate() error { func (c *Client) Close() error { c.setState(stateDisconnected) + + c.writeMutex.Lock() + defer c.writeMutex.Unlock() + return c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "bye")) } @@ -151,17 +203,78 @@ func (c *Client) Connect() error { c.setState(stateConnected) + // Set up the initial read deadline and a pong handler that extends it + // whenever a pong is received. Any data received in listen() also extends + // the deadline. If neither arrives within PongTimeout the read will error + // out, triggering a reconnect. + _ = conn.SetReadDeadline(time.Now().Add(c.cfg.PongTimeout)) + conn.SetPongHandler(func(string) error { + return conn.SetReadDeadline(time.Now().Add(c.cfg.PongTimeout)) + }) + + // done is closed when listen() exits so the keepAlive goroutine for this + // connection stops. + done := make(chan struct{}) + // Start listening for messages - go c.listen() + go c.listen(conn, done) + + // Start sending periodic pings to verify the connection is alive + go c.keepAlive(conn, done) return nil } +// keepAlive periodically sends websocket ping frames to Home Assistant. If a +// ping cannot be written the connection is considered dead and the goroutine +// exits; the corresponding read error in listen() will trigger a reconnect. +func (c *Client) keepAlive(conn *websocket.Conn, done chan struct{}) { + ticker := time.NewTicker(c.cfg.PingInterval) + defer ticker.Stop() + + for { + select { + case <-done: + return + case <-ticker.C: + c.writeMutex.Lock() + err := conn.WriteControl( + websocket.PingMessage, + []byte{}, + time.Now().Add(readTimeoutSeconds*time.Second), + ) + c.writeMutex.Unlock() + + if err != nil { + logger.Error("Failed to send ping, closing connection", "", "error", err) + + // Force the read loop to error out so the disconnection is + // signalled and a reconnect happens. Without this, an + // asymmetric failure (writes failing but reads still + // arriving) would keep the client in stateConnected while + // service calls silently fail. + _ = conn.Close() + + return + } + + logger.Debug("Sent ping", "") + } + } +} + // Listen for messages from the websocket and dispatch to listener channels. -func (c *Client) listen() { +func (c *Client) listen(conn *websocket.Conn, done chan struct{}) { logger.Info("Connection established", "") defer func() { + // Signal the keepAlive goroutine for this connection to stop + close(done) + + // Close the underlying connection so we don't leak the socket (and so + // the remote end sees the connection go away) before reconnecting. + _ = conn.Close() + // Close all pending response channels to unblock waiting callers c.mutex.Lock() for msgID, ch := range c.responses { @@ -186,7 +299,7 @@ func (c *Client) listen() { }() for { - _, msgBytes, err := c.conn.ReadMessage() + _, msgBytes, err := conn.ReadMessage() if err != nil { if websocket.IsCloseError(err, websocket.CloseNormalClosure) { logger.Info("Received close message", "") @@ -198,11 +311,16 @@ func (c *Client) listen() { return } - // Log error and return gracefully instead of panicking + // This includes read deadline (idle) timeouts: if no data or pong + // has been received within PongTimeout, ReadMessage returns an + // error and we return here to trigger a reconnect. logger.Error("Error reading from websocket", "", "error", err) return } + // Extend the read deadline: we received data, so the connection is alive. + _ = conn.SetReadDeadline(time.Now().Add(c.cfg.PongTimeout)) + logger.DebugJSON("Received message", "", string(msgBytes)) // Get message ID @@ -262,6 +380,9 @@ func (c *Client) send(msg any) error { logger.DebugJSON("Writing message", "", string(msgBytes)) + c.writeMutex.Lock() + defer c.writeMutex.Unlock() + return c.conn.WriteMessage(websocket.TextMessage, msgBytes) } diff --git a/hassws/client_test.go b/hassws/client_test.go new file mode 100644 index 0000000..4898310 --- /dev/null +++ b/hassws/client_test.go @@ -0,0 +1,38 @@ +package hassws + +import ( + "testing" + "time" + + "gotest.tools/v3/assert" +) + +func TestNewClientDefaults(t *testing.T) { + c := NewClient(ClientConfig{}) + + assert.Equal(t, defaultPingInterval, c.cfg.PingInterval) + assert.Equal(t, defaultPongTimeout, c.cfg.PongTimeout) +} + +func TestNewClientAdjustsPongTimeoutBelowPingInterval(t *testing.T) { + // PongTimeout <= PingInterval would cause a reconnect loop on a + // quiet-but-healthy connection, so it should be bumped above the ping + // interval. + c := NewClient(ClientConfig{ + PingInterval: 30 * time.Second, + PongTimeout: 10 * time.Second, + }) + + assert.Assert(t, c.cfg.PongTimeout > c.cfg.PingInterval) + assert.Equal(t, 60*time.Second, c.cfg.PongTimeout) +} + +func TestNewClientKeepsValidPongTimeout(t *testing.T) { + c := NewClient(ClientConfig{ + PingInterval: 5 * time.Second, + PongTimeout: 20 * time.Second, + }) + + assert.Equal(t, 5*time.Second, c.cfg.PingInterval) + assert.Equal(t, 20*time.Second, c.cfg.PongTimeout) +} diff --git a/hassws/server.go b/hassws/server.go index b263682..c8ac935 100644 --- a/hassws/server.go +++ b/hassws/server.go @@ -7,6 +7,7 @@ import ( "net" "net/http" "sync" + "sync/atomic" "time" "github.com/dansimau/hal/homeassistant" @@ -33,6 +34,11 @@ type Server struct { // authenticatedUserID stores the user ID of the authenticated client authenticatedUserID string + // stuck, when true, makes the current connection stop responding to pings, + // simulating a connection that is open but unresponsive. It is reset for + // each new connection. + stuck atomic.Bool + lock sync.RWMutex } @@ -70,7 +76,34 @@ func (s *Server) handler(w http.ResponseWriter, r *http.Request) { return } + s.lock.Lock() s.websocket = conn + // A new connection replaces any previous one, so reset subscribers. + s.subscribers = nil + s.lock.Unlock() + + // A fresh connection is responsive by default. Install a ping handler that + // mimics gorilla's default behaviour (reply with a pong) unless the + // connection has been marked stuck via SimulateStuckConnection. + s.stuck.Store(false) + conn.SetPingHandler(func(appData string) error { + if s.stuck.Load() { + return nil + } + + err := conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second)) + if errors.Is(err, websocket.ErrCloseSent) { + return nil + } + + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return nil + } + + return err + }) + defer conn.Close() if err := s.handleAuthentication(conn); err != nil { @@ -79,19 +112,23 @@ func (s *Server) handler(w http.ResponseWriter, r *http.Request) { return } - s.listen() + s.listen(conn) } -func (s *Server) listen() { - // Clear subscribers when connection closes +func (s *Server) listen(conn *websocket.Conn) { + // Clear subscribers when connection closes, but only if this is still the + // active connection. A stale/stuck connection that closes after a client + // has already reconnected must not clobber the new connection's state. defer func() { s.lock.Lock() - s.subscribers = nil + if s.websocket == conn { + s.subscribers = nil + } s.lock.Unlock() }() for { - _, messageBytes, err := s.websocket.ReadMessage() + _, messageBytes, err := conn.ReadMessage() if err != nil { if websocket.IsCloseError(err, websocket.CloseNormalClosure) { log.Println("[Server] Received close message, bye") @@ -246,6 +283,14 @@ func (s *Server) Close() error { ) } +// SimulateStuckConnection simulates a connection that is still open but no +// longer responsive: it stops replying to ping frames (and, being a test +// server, sends no further data). This is used to verify the client's +// heartbeat / read-timeout based reconnection. +func (s *Server) SimulateStuckConnection() { + s.stuck.Store(true) +} + // DisconnectClient forcibly closes the WebSocket connection to simulate network failure. // Unlike Close(), this does not send a graceful close message. // Returns nil if connection is already closed.