diff --git a/pkg/i18n/en_base_config_descriptions.go b/pkg/i18n/en_base_config_descriptions.go index 3f6f8df..6b1f7c6 100644 --- a/pkg/i18n/en_base_config_descriptions.go +++ b/pkg/i18n/en_base_config_descriptions.go @@ -59,14 +59,16 @@ var ( ConfigGlobalSize = ffc("config.global.cache.size", "The size of the cache", ByteSizeType) ConfigGlobalTTL = ffc("config.global.cache.ttl", "The time to live (TTL) for the cache", TimeDurationType) - ConfigGlobalWsConnectionTimeout = ffc("config.global.ws.connectionTimeout", "The amount of time to wait while establishing a connection (or auto-reconnection)", TimeDurationType) - ConfigGlobalWsHeartbeatInterval = ffc("config.global.ws.heartbeatInterval", "The amount of time to wait between heartbeat signals on the WebSocket connection", TimeDurationType) - ConfigGlobalWsBackgroundConnect = ffc("config.global.ws.backgroundConnect", "When true the connection is established in the background with infinite reconnect (makes initialConnectAttempts redundant when set)", BooleanType) - ConfigGlobalWsInitialConnectAttempts = ffc("config.global.ws.initialConnectAttempts", "The number of attempts FireFly will make to connect to the WebSocket when starting up, before failing", IntType) - ConfigGlobalWsPath = ffc("config.global.ws.path", "The WebSocket sever URL to which FireFly should connect", "WebSocket URL "+StringType) - ConfigGlobalWsReadBufferSize = ffc("config.global.ws.readBufferSize", "The size in bytes of the read buffer for the WebSocket connection", ByteSizeType) - ConfigGlobalWsWriteBufferSize = ffc("config.global.ws.writeBufferSize", "The size in bytes of the write buffer for the WebSocket connection", ByteSizeType) - ConfigGlobalWsURL = ffc("config.global.ws.url", "URL to use for WebSocket - overrides url one level up (in the HTTP config)", StringType) + ConfigGlobalWsConnectionTimeout = ffc("config.global.ws.connectionTimeout", "The amount of time to wait while establishing a connection (or auto-reconnection)", TimeDurationType) + ConfigGlobalWsHeartbeatInterval = ffc("config.global.ws.heartbeatInterval", "The amount of time to wait between heartbeat signals on the WebSocket connection", TimeDurationType) + ConfigGlobalWsBackgroundConnect = ffc("config.global.ws.backgroundConnect", "When true the connection is established in the background with infinite reconnect (makes initialConnectAttempts redundant when set)", BooleanType) + ConfigGlobalWsInitialConnectAttempts = ffc("config.global.ws.initialConnectAttempts", "The number of attempts FireFly will make to connect to the WebSocket when starting up, before failing", IntType) + ConfigGlobalWsPath = ffc("config.global.ws.path", "The WebSocket sever URL to which FireFly should connect", "WebSocket URL "+StringType) + ConfigGlobalWsReadBufferSize = ffc("config.global.ws.readBufferSize", "The size in bytes of the read buffer for the WebSocket connection", ByteSizeType) + ConfigGlobalWsWriteBufferSize = ffc("config.global.ws.writeBufferSize", "The size in bytes of the write buffer for the WebSocket connection", ByteSizeType) + ConfigGlobalWsURL = ffc("config.global.ws.url", "URL to use for WebSocket - overrides url one level up (in the HTTP config)", StringType) + ConfigGlobalWsConnectionCycleInterval = ffc("config.global.ws.connectionCycleInterval", "When non-zero, the client proactively replaces its WebSocket connection on this interval - the new connection is fully established (including post-connect setup) before the old one is quiesced and closed", TimeDurationType) + ConfigGlobalWsConnectionCycleQuiesceTime = ffc("config.global.ws.connectionCycleQuiesceTime", "The amount of time the old WebSocket connection continues to deliver inbound messages after a connection cycle switches sending to the new connection, before it is closed", TimeDurationType) ConfigGlobalTLSCA = ffc("config.global.tls.ca", "The TLS certificate authority in PEM format (this option is ignored if caFile is also set)", StringType) ConfigGlobalTLSCaFile = ffc("config.global.tls.caFile", "The path to the CA file for TLS on this API", StringType) diff --git a/pkg/wsclient/wsclient.go b/pkg/wsclient/wsclient.go index be2c931..d0f3789 100644 --- a/pkg/wsclient/wsclient.go +++ b/pkg/wsclient/wsclient.go @@ -62,6 +62,19 @@ type WSConfig struct { // underlying TCP connection. Built by GenerateConfig from the net config; cannot be set in // JSON. Left nil for hand-built configs, in which case the default net dialer is used. NetDialer *net.Dialer `json:"-"` + // ConnectionCycleInterval when non-zero enables proactive replacement of the connection + // on this interval - the new connection is fully established (including afterConnect) + // before sends switch over and the old connection is quiesced and closed. + // The interval restarts from the end of each quiesce/close, and from any reconnect + // due to a connection error - so at most two connections ever exist concurrently. + ConnectionCycleInterval time.Duration `json:"connectionCycleInterval,omitempty"` + // ConnectionCycleQuiesceTime is how long the old connection continues to deliver inbound + // messages after a connection cycle switches sends to the new connection, before it is closed + ConnectionCycleQuiesceTime time.Duration `json:"connectionCycleQuiesceTime,omitempty"` + // The lifecycle handlers cannot be set in JSON - they must be configured on the code interface + PreConnectHandler WSPreConnectHandler `json:"-"` + PostConnectHandler WSPostConnectHandler `json:"-"` + PreDisconnectHandler WSPreDisconnectHandler `json:"-"` // This one cannot be set in JSON - must be configured on the code interface ReceiveExt bool } @@ -113,25 +126,24 @@ type wsClient struct { backgroundConnect bool initialRetryAttempts int wsdialer *websocket.Dialer - wsconn *websocket.Conn + current *wsConnection connRetry retry.Retry closed bool useReceiveExt bool receive chan []byte receiveExt chan *WSPayload send chan []byte - sendDone chan []byte bgConnCancelCtx context.CancelFunc bgConnDone chan struct{} closing chan struct{} beforeConnect WSPreConnectHandler afterConnect WSPostConnectHandler + beforeDisconnect WSPreDisconnectHandler disableReconnect bool heartbeatInterval time.Duration - stateMux sync.Mutex // guards closed, wsconn, bgConnDone and bgConnCancelCtx - heartbeatMux sync.Mutex - activePingSent *time.Time - lastPingCompleted time.Time + connCycleInterval time.Duration + connCycleQuiesce time.Duration + stateMux sync.Mutex // guards closed, current, bgConnDone and bgConnCancelCtx rateLimiter *rate.Limiter } @@ -139,10 +151,26 @@ type wsClient struct { type WSPreConnectHandler func(ctx context.Context, w WSClient) error // WSPostConnectHandler will be called after every connect/reconnect. Can send data over ws, but must not block listening for data on the ws. +// Note: During auto-cycle this is called on the new connection, after the WSPreDisconnectHandler is called on the old one, but before the old connection is closed. type WSPostConnectHandler func(ctx context.Context, w WSClient) error -// Creates a new outbound client that can be connected to a remote server +// WSPreDisconnectHandler is called before a graceful close, to allow cleanup (such as unsubscribe): +// - When closed explicitly +// - When cycling the connection (after the new connection is established, before post-connect is called) +type WSPreDisconnectHandler func(ctx context.Context, w WSClient) error + +// New creates a new outbound client that can be connected to a remote server. +// ** Recommend using NewWithConfig directly ** func New(ctx context.Context, config *WSConfig, beforeConnect WSPreConnectHandler, afterConnect WSPostConnectHandler) (WSClient, error) { + conf := *config // copy, so we don't modify the supplied config with the handler overrides + conf.PreConnectHandler = beforeConnect + conf.PostConnectHandler = afterConnect + return NewWithConfig(ctx, &conf) +} + +// NewWithConfig creates a new outbound WebSocket client with configuration, +// including lifecycle hooks +func NewWithConfig(ctx context.Context, config *WSConfig) (WSClient, error) { l := log.L(ctx) wsURL, err := buildWSUrl(ctx, config) if err != nil { @@ -175,13 +203,19 @@ func New(ctx context.Context, config *WSConfig, beforeConnect WSPreConnectHandle headers: make(http.Header), send: make(chan []byte), closing: make(chan struct{}), - beforeConnect: beforeConnect, - afterConnect: afterConnect, + beforeConnect: config.PreConnectHandler, + afterConnect: config.PostConnectHandler, + beforeDisconnect: config.PreDisconnectHandler, heartbeatInterval: config.HeartbeatInterval, + connCycleInterval: config.ConnectionCycleInterval, + connCycleQuiesce: config.ConnectionCycleQuiesceTime, useReceiveExt: config.ReceiveExt, disableReconnect: config.DisableReconnect, rateLimiter: ffresty.GetRateLimiter(config.ThrottleRequestsPerSecond, config.ThrottleBurst), } + if w.connCycleInterval > 0 && w.disableReconnect { + l.Warnf("WS %s connection cycling configured, but inactive as reconnect is disabled", w.url) + } w.setupReceiveChannel() for k, v := range config.HTTPHeaders { if vs, ok := v.(string); ok { @@ -214,7 +248,6 @@ func Wrap(ctx context.Context, config WSWrapConfig, wsconn *websocket.Conn, onCl w := &wsClient{ ctx: ctx, url: wsconn.LocalAddr().String(), - wsconn: wsconn, disableReconnect: true, heartbeatInterval: config.HeartbeatInterval, rateLimiter: ffresty.GetRateLimiter(config.ThrottleRequestsPerSecond, config.ThrottleBurst), @@ -223,8 +256,9 @@ func Wrap(ctx context.Context, config WSWrapConfig, wsconn *websocket.Conn, onCl closing: make(chan struct{}), } w.setupReceiveChannel() - w.pongReceivedOrReset(false) - w.wsconn.SetPongHandler(w.pongHandler) + c := w.newConnection(wsconn) + close(c.promoted) // sole connection - immediately the consumer of the shared send channel + w.current = c log.L(ctx).Infof("WS %s wrapped", w.url) go func() { w.receiveReconnectLoop() @@ -283,12 +317,25 @@ func (w *wsClient) initialConnect() error { } func (w *wsClient) Close() { + // Run the pre-disconnect handler for an orderly close of the current connection, before + // we start closing - the same handler that runs on a connection cycle, so pre-close + // processing lives in one place. The claim ensures at most one invocation per connection + // (a re-entrant or concurrent Close proceeds straight to closing), and any error just + // means we fall back to relying on server-side cleanup. + if w.beforeDisconnect != nil { + if c := w.currentConn(); c != nil && !w.isClosed() && c.claimPreDisconnect() { + if err := w.beforeDisconnect(w.ctx, w.boundTo(c)); err != nil { + log.L(w.ctx).Warnf("WS %s pre-disconnect handler failed in close: %s", w.url, err) + } + } + } + c, bgConnDone, bgConnCancelCtx, alreadyClosed := w.markClosed() if alreadyClosed { return } if c != nil { - _ = c.Close() + c.closeConn() } if bgConnDone != nil { // Cancel the background connect routine and wait for it to exit. Note we must not @@ -300,7 +347,8 @@ func (w *wsClient) Close() { // markClosed transitions the client to closed exactly once, returning the resources the // caller must then clean up outside of the lock. -func (w *wsClient) markClosed() (c *websocket.Conn, bgConnDone chan struct{}, bgConnCancelCtx context.CancelFunc, alreadyClosed bool) { +// Note an old connection mid-quiesce during a cycle is not returned here (handled by cycleConnection) +func (w *wsClient) markClosed() (c *wsConnection, bgConnDone chan struct{}, bgConnCancelCtx context.CancelFunc, alreadyClosed bool) { w.stateMux.Lock() defer w.stateMux.Unlock() @@ -311,7 +359,7 @@ func (w *wsClient) markClosed() (c *websocket.Conn, bgConnDone chan struct{}, bg close(w.closing) // bgConnCancelCtx+bgConnDone are both set as a pair in the stateMux, so if one is non-nil they both are - c, bgConnDone, bgConnCancelCtx = w.wsconn, w.bgConnDone, w.bgConnCancelCtx + c, bgConnDone, bgConnCancelCtx = w.current, w.bgConnDone, w.bgConnCancelCtx w.bgConnDone = nil return c, bgConnDone, bgConnCancelCtx, false } @@ -323,24 +371,46 @@ func (w *wsClient) isClosed() bool { return w.closed } -// called when we've just connected a new underlying websocket.Conn, returning false +// currentConn gets the current connection under the stateMux +func (w *wsClient) currentConn() *wsConnection { + w.stateMux.Lock() + defer w.stateMux.Unlock() + return w.current +} + +// called when we've just connected a new underlying connection, returning false // if the wsClient was cleaned up in the meantime - meaning the caller has an orphaned // connection they need to close. -func (w *wsClient) setWSConnIfNotClosed(conn *websocket.Conn) bool { +func (w *wsClient) setCurrentIfNotClosed(c *wsConnection) bool { w.stateMux.Lock() defer w.stateMux.Unlock() if w.closed { return false } - w.wsconn = conn + w.current = c return true } -func (w *wsClient) clearWSConn() { +func (w *wsClient) clearCurrentIf(c *wsConnection) { w.stateMux.Lock() defer w.stateMux.Unlock() - w.wsconn = nil + if w.current == c { + w.current = nil + } +} + +// promoteConnection atomically switches the connection from old to newC during a cycle +func (w *wsClient) promoteConnection(old, newC *wsConnection) bool { + w.stateMux.Lock() + defer w.stateMux.Unlock() + if w.closed { + return false + } + w.current = newC + close(old.demoted) + close(newC.promoted) + return true } func (w *wsClient) Receive() <-chan []byte { @@ -364,42 +434,81 @@ func (w *wsClient) SetHeader(header, value string) { w.headers.Set(header, value) } -func (w *wsClient) Send(ctx context.Context, message []byte) error { +func (w *wsClient) waitRateLimiter(ctx context.Context) error { if w.rateLimiter != nil { // Wait for permission to proceed with the request - err := w.rateLimiter.Wait(ctx) - if err != nil { - return err - } + return w.rateLimiter.Wait(ctx) + } + return nil +} + +func (w *wsClient) Send(ctx context.Context, message []byte) error { + if err := w.waitRateLimiter(ctx); err != nil { + return err } // Send + for { + // The sendDone of the current connection guards against blocking forever when that + // connection's sender loop has exited - needed because the receiver can actually + // call the sender indirectly on reconnect, so if the sender loop fails the + // receiver can get blocked + var sendDone chan []byte + c := w.currentConn() + if c != nil { + sendDone = c.sendDone + } + select { + case w.send <- message: + return nil + case <-ctx.Done(): + return i18n.NewError(ctx, i18n.MsgWSSendTimedOut) + case <-sendDone: + if w.currentConn() != c { + continue // the connection was replaced under us (reconnect/cycle) - retry against the new one + } + return i18n.NewError(ctx, i18n.MsgWSClosing) + case <-w.closing: + return i18n.NewError(ctx, i18n.MsgWSClosing) + } + } +} + +// connBoundClient is the WSClient facade supplied to the connect/disconnect handlers - +// its Send() is bound to one specific connection (as is promised during a cycle). +type connBoundClient struct { + *wsClient + c *wsConnection +} + +func (bc *connBoundClient) Send(ctx context.Context, message []byte) error { + if err := bc.waitRateLimiter(ctx); err != nil { + return err + } + bs := &trackedSend{message: message, sent: make(chan bool, 1)} select { - case w.send <- message: - return nil + case bc.c.send <- bs: + // Handed off - now wait for the write to complete, so a pre-disconnect handler's + // messages are on the wire before the connection is closed behind it + select { + case ok := <-bs.sent: + if !ok { + return i18n.NewError(ctx, i18n.MsgWSClosing) + } + return nil + case <-ctx.Done(): + return i18n.NewError(ctx, i18n.MsgWSSendTimedOut) + } case <-ctx.Done(): return i18n.NewError(ctx, i18n.MsgWSSendTimedOut) - case <-w.sendDone: - // Need this case because the receiver can actually call the sender indirectly on reconnect, - // so if the sender loop fails the receiver can get blocked + case <-bc.c.sendDone: return i18n.NewError(ctx, i18n.MsgWSClosing) - case <-w.closing: + case <-bc.closing: return i18n.NewError(ctx, i18n.MsgWSClosing) } } -func (w *wsClient) heartbeatTimeout(ctx context.Context) (context.Context, context.CancelFunc) { - if w.heartbeatInterval > 0 { - w.heartbeatMux.Lock() - baseTime := w.lastPingCompleted - if w.activePingSent != nil { - // We're waiting for a pong - baseTime = *w.activePingSent - } - waitTime := w.heartbeatInterval - time.Since(baseTime) // if negative, will pop immediately - w.heartbeatMux.Unlock() - return context.WithTimeout(ctx, waitTime) - } - return context.WithCancel(ctx) +func (w *wsClient) boundTo(c *wsConnection) WSClient { + return &connBoundClient{wsClient: w, c: c} } func buildWSUrl(ctx context.Context, config *WSConfig) (string, error) { @@ -432,6 +541,36 @@ func buildWSUrl(ctx context.Context, config *WSConfig) (string, error) { return u.String(), nil } +// dialConnectionAttempt makes a single connect attempt (including the beforeConnect handler). +// Does not start send/receive loops, or set it as the active connection. +func (w *wsClient) dialConnectionAttempt(attempt int) (*wsConnection, error) { + l := log.L(w.ctx) + if w.beforeConnect != nil { + if err := w.beforeConnect(w.ctx, w); err != nil { + l.Warnf("WS %s connect attempt %d failed in beforeConnect", w.url, attempt) + return nil, err + } + } + + conn, res, err := w.wsdialer.DialContext(w.ctx, w.url, w.headers) + if err != nil { + errMsg := err.Error() + var status = -1 + if res != nil { + b, readErr := io.ReadAll(res.Body) + res.Body.Close() + if readErr == nil && len(b) > 0 { + // The info we need is what the server returned and the status + errMsg = string(b) + } + status = res.StatusCode + } + l.Warnf("WS %s connect attempt %d failed [%d]: %s", w.url, attempt, status, errMsg) + return nil, i18n.WrapError(w.ctx, err, i18n.MsgWSConnectFailed) + } + return w.newConnection(conn), nil +} + func (w *wsClient) connect(initial bool) error { l := log.L(w.ctx) l.Debugf("WS %s connecting, isInitial: %t", w.url, initial) @@ -442,169 +581,21 @@ func (w *wsClient) connect(initial bool) error { } l.Debugf("WS %s connect attempt %d", w.url, attempt) retry = w.backgroundConnect || !initial || attempt < w.initialRetryAttempts - if w.beforeConnect != nil { - if err = w.beforeConnect(w.ctx, w); err != nil { - l.Warnf("WS %s connect attempt %d failed in beforeConnect", w.url, attempt) - return retry, err - } - } - - var res *http.Response - var conn *websocket.Conn - conn, res, err = w.wsdialer.DialContext(w.ctx, w.url, w.headers) + c, err := w.dialConnectionAttempt(attempt) if err != nil { - errMsg := err.Error() - var status = -1 - if res != nil { - b, readErr := io.ReadAll(res.Body) - res.Body.Close() - if readErr == nil && len(b) > 0 { - // The info we need is what the server returned and the status - errMsg = string(b) - } - status = res.StatusCode - } - l.Warnf("WS %s connect attempt %d failed [%d]: %s", w.url, attempt, status, errMsg) - return retry, i18n.WrapError(w.ctx, err, i18n.MsgWSConnectFailed) + return retry, err } - if !w.setWSConnIfNotClosed(conn) { - _ = conn.Close() // we have to clean up the orphan we just created + if !w.setCurrentIfNotClosed(c) { + c.closeConn() // we have to clean up the orphan we just created return false, i18n.NewError(w.ctx, i18n.MsgWSClosing) } l.Debugf("WS %s connect attempt %d succeeded", w.url, attempt) - w.pongReceivedOrReset(false) - conn.SetPongHandler(w.pongHandler) + close(c.promoted) // sole connection - immediately the consumer of the shared send channel l.Infof("WS %s connected", w.url) return false, nil }) } -func (w *wsClient) readLoop() { - l := log.L(w.ctx) - for { - mt, message, err := w.wsconn.ReadMessage() - if err != nil { - // We treat this as informational, as it's normal for the client to disconnect here - l.Infof("WS %s closed: %s", w.url, err) - return - } - - // Pass the message to the consumer - l.Tracef("WS %s read (mt=%d): %s", w.url, mt, message) - select { - case <-w.sendDone: - l.Debugf("WS %s closing reader after send error", w.url) - return - case w.receive <- message: - } - } -} - -func (w *wsClient) readLoopExt() { - l := log.L(w.ctx) - for { - // We set a deadline for twice the heartbeat interval - note we bump this on pong - if w.heartbeatInterval > 0 { - _ = w.wsconn.SetReadDeadline(time.Now().Add(2 * w.heartbeatInterval)) - } - - mt, r, err := w.wsconn.NextReader() - if err != nil { - // We treat this as informational, as it's normal for the client to disconnect here - l.Infof("WS %s closed: %s", w.url, err) - return - } - - // Pass the message to the consumer - l.Tracef("WS %s read (mt=%d)", w.url, mt) - payload := NewWSPayload(mt, r) - select { - case <-w.sendDone: - l.Debugf("WS %s closing reader after send error (waiting for data)", w.url) - return - case w.receiveExt <- payload: - } - select { - case <-payload.processed: - // It's the callers responsibility to ensure they call done on this before we can get the next payload - case <-w.sendDone: - l.Debugf("WS %s closing reader after send error (waiting for processing of data by client)", w.url) - return - } - } -} - -func (w *wsClient) pongHandler(_ string) error { - w.pongReceivedOrReset(true) - return nil -} - -func (w *wsClient) pongReceivedOrReset(isPong bool) { - w.heartbeatMux.Lock() - defer w.heartbeatMux.Unlock() - - if isPong && w.activePingSent != nil { - log.L(w.ctx).Debugf("WS %s heartbeat completed (pong) after %.2fms", w.url, float64(time.Since(*w.activePingSent))/float64(time.Millisecond)) - } - w.lastPingCompleted = time.Now() // in new connection case we still want to consider now the time we completed the ping - w.activePingSent = nil - - // We set a deadline for twice the heartbeat interval - if w.heartbeatInterval > 0 { - _ = w.wsconn.SetReadDeadline(time.Now().Add(2 * w.heartbeatInterval)) - } - -} - -func (w *wsClient) heartbeatCheck() error { - w.heartbeatMux.Lock() - defer w.heartbeatMux.Unlock() - - if w.activePingSent != nil { - return i18n.NewError(w.ctx, i18n.MsgWSHeartbeatTimeout, float64(time.Since(*w.activePingSent))/float64(time.Millisecond)) - } - log.L(w.ctx).Debugf("WS %s heartbeat timer popped (ping) after %.2fms", w.url, float64(time.Since(w.lastPingCompleted))/float64(time.Millisecond)) - now := time.Now() - w.activePingSent = &now - return nil -} - -func (w *wsClient) sendLoop(receiverDone chan struct{}) { - l := log.L(w.ctx) - defer close(w.sendDone) - - disconnecting := false - for !disconnecting { - timeoutContext, timeoutCancel := w.heartbeatTimeout(w.ctx) - - select { - case message := <-w.send: - l.Tracef("WS sending: %s", message) - if err := w.wsconn.WriteMessage(websocket.TextMessage, message); err != nil { - l.Errorf("WS %s send failed: %s", w.url, err) - disconnecting = true - } - case <-timeoutContext.Done(): - wsconn := w.wsconn - if err := w.heartbeatCheck(); err != nil { - l.Errorf("WS %s closing: %s", w.url, err) - disconnecting = true - } else if wsconn != nil { - l.Debugf("WS %s send heartbeat ping", w.url) - if err := wsconn.WriteMessage(websocket.PingMessage, []byte{}); err != nil { - l.Errorf("WS %s heartbeat send failed: %s", w.url, err) - disconnecting = true - } - } - case <-receiverDone: - l.Debugf("WS %s send loop exiting", w.url) - disconnecting = true - } - - timeoutCancel() - } -} - func (w *wsClient) receiveReconnectLoop() { l := log.L(w.ctx) if w.useReceiveExt { @@ -612,37 +603,59 @@ func (w *wsClient) receiveReconnectLoop() { } else { defer close(w.receive) } + + // Connection cycling proactively replaces the connection on a regular interval, when + // enabled. The timer restarts after the completion of each cycle (the end of the + // quiesce period), and after any reconnect due to a connection error. + var cycleC <-chan time.Time + var cycleTimer *time.Timer + cyclingEnabled := w.connCycleInterval > 0 && !w.disableReconnect + if cyclingEnabled { + cycleTimer = time.NewTimer(w.connCycleInterval) + defer cycleTimer.Stop() + cycleC = cycleTimer.C + } + resetCycleTimer := func() { + if cyclingEnabled { + cycleTimer.Reset(w.connCycleInterval) + } + } + for !w.isClosed() { // Start the sender, letting it close without blocking sending a notification on the sendDone - w.sendDone = make(chan []byte, 1) - receiverDone := make(chan struct{}) - go w.sendLoop(receiverDone) + c := w.currentConn() + c.startSender() - // Call the reconnect processor + // Call the reconnect processor, bound to this connection var err error if w.afterConnect != nil { - err = w.afterConnect(w.ctx, w) + err = w.afterConnect(w.ctx, w.boundTo(c)) l.Debugf("WS %s afterConnect (error: %v)", w.url, err) } if err == nil { - // Synchronously invoke the reader, as it's important we react immediately to any error there. - if w.useReceiveExt { - w.readLoopExt() - } else { - w.readLoop() - } - close(receiverDone) - <-w.sendDone - - // Ensure the connection is closed after the sender and receivers exit - err = w.wsconn.Close() - if err != nil { - l.Warnf("WS %s ignoring websocket connection close error: %v", w.url, err) + c.startReader() + connected: + for { + select { + case <-c.readDone: + // The reader exited - a connection error, server close, or Close() + w.teardownConnection(c) + w.clearCurrentIf(c) + l.Debugf("WS %s reset the connection", w.url) + break connected + case <-cycleC: + if newC, ok := w.cycleConnection(c); ok { + c = newC // adopt the new connection - its send/read loops are already running + } + // if !ok the client is closing - the readDone/isClosed checks handle the exit + resetCycleTimer() + } } - l.Debugf("WS %s reset the connection", w.url) - w.sendDone = nil - w.clearWSConn() + } else { + // Ensure the connection and its sender are fully cleaned up before we reconnect + w.teardownConnection(c) + w.clearCurrentIf(c) } if w.disableReconnect { @@ -657,6 +670,75 @@ func (w *wsClient) receiveReconnectLoop() { l.Errorf("WS %s exiting due to connect error: %v", w.url, err) return } + resetCycleTimer() + } + } +} + +// cycleConnection runs on the receiveReconnectLoop goroutine when a cycle is due. +// It does not return until either the old connection is fully torn down, or the client is +// closing. This ensures we never have more than two connections. +func (w *wsClient) cycleConnection(old *wsConnection) (*wsConnection, bool) { + l := log.L(w.ctx) + l.Debugf("WS %s connection cycle starting", w.url) + + // Establish the new connection - with infinite retry (like reconnect), keeping the + // old connection fully active (sending, receiving and heartbeating) throughout. + var newC *wsConnection + err := w.connRetry.DoCustomLog(w.ctx, func(attempt int) (retry bool, err error) { + if w.isClosed() { + return false, i18n.NewError(w.ctx, i18n.MsgWSClosing) + } + c, err := w.dialConnectionAttempt(attempt) + if err != nil { + return true, err // the old connection remains active while we retry + } + c.startSender() // services connection-bound sends from the hooks (not yet promoted) + + // The pre-disconnect handler gets to run on cycle, as well as close. + // This happens before the quiesce cycle, so things like subscriptions are only active + // on a single connection. In-flight request/reply exchanges can continue on the old connection. + if w.beforeDisconnect != nil && old.claimPreDisconnect() { + if pdErr := w.beforeDisconnect(w.ctx, w.boundTo(old)); pdErr != nil { + l.Warnf("WS %s pre-disconnect handler failed (continuing connection cycle): %s", w.url, pdErr) + } } + + // Call the connect processor against the new connection + if w.afterConnect != nil { + if acErr := w.afterConnect(w.ctx, w.boundTo(c)); acErr != nil { + l.Warnf("WS %s connection cycle attempt %d failed in afterConnect: %s", w.url, attempt, acErr) + w.teardownConnection(c) + return true, acErr + } + } + newC = c + return false, nil + }) + if err != nil { + // Only reachable when the client is closing / the context is cancelled + return nil, false + } + + // Start reading, and atomically switch all new sends over to the new connection + newC.startReader() + if !w.promoteConnection(old, newC) { + // The client was closed while we were cycling - clean up the orphaned new connection + w.teardownConnection(newC) + return nil, false + } + + // Quiesce period - the old connection continues to deliver any in-flight inbound + // messages to the receive channel, before we close it. + old.setQuiesceDeadline(w.connCycleQuiesce) + l.Infof("WS %s connection cycled, quiescing old connection for %s", w.url, w.connCycleQuiesce) + quiesce := time.NewTimer(w.connCycleQuiesce) + defer quiesce.Stop() + select { + case <-quiesce.C: + case <-old.readDone: // the old connection failed during quiesce - just close it early + case <-w.closing: // Close() handles the current (new) connection - we clean up the old one below } + w.teardownConnection(old) + return newC, true } diff --git a/pkg/wsclient/wsclient_cycle_test.go b/pkg/wsclient/wsclient_cycle_test.go new file mode 100644 index 0000000..9ef0fc1 --- /dev/null +++ b/pkg/wsclient/wsclient_cycle_test.go @@ -0,0 +1,736 @@ +// Copyright © 2026 Kaleido, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package wsclient + +import ( + "context" + "fmt" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "golang.org/x/time/rate" +) + +func cycleTestConfig(url string) *WSConfig { + return &WSConfig{ + HTTPURL: url, + InitialDelay: 5 * time.Millisecond, + MaximumDelay: 20 * time.Millisecond, + ConnectionCycleInterval: 75 * time.Millisecond, + ConnectionCycleQuiesceTime: 150 * time.Millisecond, + } +} + +func nextConn(t *testing.T, connections chan *TestWSConnection) *TestWSConnection { + t.Helper() + select { + case c := <-connections: + return c + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for a connection to the test server") + return nil + } +} + +func expectMsg(t *testing.T, c *TestWSConnection, expected string) { + t.Helper() + select { + case msg := <-c.ToServer: + assert.Equal(t, expected, msg) + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for message %q", expected) + } +} + +func waitDone(t *testing.T, c *TestWSConnection) { + t.Helper() + select { + case <-c.Done: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for connection to close") + } +} + +// waitForSwitch polls (white box) until the current connection is no longer the supplied +// one - so the test knows promotion has completed and the old connection is quiescing +func waitForSwitch(t *testing.T, wsc WSClient, from *wsConnection) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + c := wsc.(*wsClient).currentConn() + if c != nil && c != from { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatal("timed out waiting for connection switch") +} + +func TestWSConnectionCycleE2E(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := cycleTestConfig(url) + wsConfig.HeartbeatInterval = 25 * time.Millisecond + wsConfig.PreDisconnectHandler = func(ctx context.Context, w WSClient) error { + return w.Send(ctx, []byte(`unsubscribe`)) + } + wsConfig.PostConnectHandler = func(ctx context.Context, w WSClient) error { + return w.Send(ctx, []byte(`subscribe`)) + } + + wsc, err := NewWithConfig(context.Background(), wsConfig) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + expectMsg(t, conn1, `subscribe`) + + // The first cycle establishes a second connection - the pre-disconnect handler sends + // on the OLD connection, and the post-connect handler sends on the NEW connection + conn2 := nextConn(t, connections) + expectMsg(t, conn1, `unsubscribe`) + expectMsg(t, conn2, `subscribe`) + + // During the quiesce period the old connection still delivers inbound messages + conn1.FromServer <- `old inbound` + assert.Equal(t, `old inbound`, string(<-wsc.Receive())) + + // The old connection closes at the end of the quiesce period + waitDone(t, conn1) + + // After the cycle, sends go to the new connection, and inbound flows on it + err = wsc.Send(context.Background(), []byte(`to new conn`)) + assert.NoError(t, err) + expectMsg(t, conn2, `to new conn`) + conn2.FromServer <- `new inbound` + assert.Equal(t, `new inbound`, string(<-wsc.Receive())) + + // A second cycle proves the timer restarts after the quiesce completes + conn3 := nextConn(t, connections) + select { + case <-conn1.Done: + // at most two connections ever - the first is long gone before the third arrives + default: + t.Fatal("conn1 still live when conn3 established") + } + expectMsg(t, conn2, `unsubscribe`) + expectMsg(t, conn3, `subscribe`) + waitDone(t, conn2) + + // Close also runs the pre-disconnect handler, on the final connection + wsc.Close() + expectMsg(t, conn3, `unsubscribe`) + _, ok := <-wsc.Receive() + assert.False(t, ok) + waitDone(t, conn3) +} + +func TestWSConnectionCycleReceiveExt(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := cycleTestConfig(url) + wsConfig.HeartbeatInterval = 25 * time.Millisecond + wsConfig.ReceiveExt = true + afterConnect := func(ctx context.Context, w WSClient) error { + return w.Send(ctx, []byte(`subscribe`)) + } + + wsc, err := New(context.Background(), wsConfig, nil, afterConnect) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + expectMsg(t, conn1, `subscribe`) + + // No pre-disconnect handler configured - the cycle works fine without one + conn2 := nextConn(t, connections) + expectMsg(t, conn2, `subscribe`) + + // The old connection still delivers via the ext receive channel during quiesce + conn1.FromServer <- `old inbound` + payload := <-wsc.ReceiveExt() + b, err := io.ReadAll(payload.Reader) + assert.NoError(t, err) + assert.Equal(t, `old inbound`, string(b)) + payload.Processed() + + waitDone(t, conn1) + + conn2.FromServer <- `new inbound` + payload = <-wsc.ReceiveExt() + b, err = io.ReadAll(payload.Reader) + assert.NoError(t, err) + assert.Equal(t, `new inbound`, string(b)) + payload.Processed() + + wsc.Close() + _, ok := <-wsc.ReceiveExt() + assert.False(t, ok) +} + +func TestWSCycleDialFailureKeepsOldConnection(t *testing.T) { + connections, url, rejectNext, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := cycleTestConfig(url) + wsConfig.ConnectionCycleInterval = 50 * time.Millisecond + + wsc, err := New(context.Background(), wsConfig, nil, nil) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + + // Make the cycle's dial attempts fail - the old connection must remain fully active + rejectNext(true) + time.Sleep(150 * time.Millisecond) // well past the cycle boundary + + err = wsc.Send(context.Background(), []byte(`still on old`)) + assert.NoError(t, err) + expectMsg(t, conn1, `still on old`) + conn1.FromServer <- `old still receiving` + assert.Equal(t, `old still receiving`, string(<-wsc.Receive())) + + // Allow the dial to succeed - the cycle completes + rejectNext(false) + conn2 := nextConn(t, connections) + waitDone(t, conn1) + + err = wsc.Send(context.Background(), []byte(`on new`)) + assert.NoError(t, err) + expectMsg(t, conn2, `on new`) + + wsc.Close() + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCycleAfterConnectFailureRetries(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + var connCount, preDisconnectCount atomic.Int32 + wsConfig := cycleTestConfig(url) + wsConfig.PreDisconnectHandler = func(ctx context.Context, w WSClient) error { + preDisconnectCount.Add(1) + return nil + } + afterConnect := func(ctx context.Context, w WSClient) error { + if connCount.Add(1) == 2 { + return fmt.Errorf("pop") // fail the first attempt of the first cycle + } + return w.Send(ctx, []byte(`subscribe`)) + } + + wsc, err := New(context.Background(), wsConfig, nil, afterConnect) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + expectMsg(t, conn1, `subscribe`) + + // afterConnect fails on this connection, so it is torn down and the dial retried + conn2 := nextConn(t, connections) + waitDone(t, conn2) + + // The retry succeeds, and the cycle completes + conn3 := nextConn(t, connections) + expectMsg(t, conn3, `subscribe`) + waitDone(t, conn1) + + // The pre-disconnect handler ran exactly once, despite the retry + assert.Equal(t, int32(1), preDisconnectCount.Load()) + + wsc.Close() + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCyclePreDisconnectErrorNonFatal(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := cycleTestConfig(url) + wsConfig.PreDisconnectHandler = func(ctx context.Context, w WSClient) error { + return fmt.Errorf("pop") + } + afterConnect := func(ctx context.Context, w WSClient) error { + return w.Send(ctx, []byte(`subscribe`)) + } + + wsc, err := New(context.Background(), wsConfig, nil, afterConnect) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + expectMsg(t, conn1, `subscribe`) + + // The cycle completes regardless of the pre-disconnect handler failing + conn2 := nextConn(t, connections) + expectMsg(t, conn2, `subscribe`) + waitDone(t, conn1) + + err = wsc.Send(context.Background(), []byte(`on new`)) + assert.NoError(t, err) + expectMsg(t, conn2, `on new`) + + wsc.Close() + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCycleCloseMidQuiesce(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := cycleTestConfig(url) + wsConfig.ConnectionCycleInterval = 25 * time.Millisecond + wsConfig.ConnectionCycleQuiesceTime = 1 * time.Minute // Close must not wait for this + + wsc, err := New(context.Background(), wsConfig, nil, nil) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + c1 := wsc.(*wsClient).currentConn() + + _ = nextConn(t, connections) + waitForSwitch(t, wsc, c1) // the old connection is now quiescing + + wsc.Close() + + // Both connections close promptly, and the receive channel closes + waitDone(t, conn1) + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCycleCloseBeforeSwitch(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + var wsc WSClient + var connCount atomic.Int32 + afterConnect := func(ctx context.Context, w WSClient) error { + if connCount.Add(1) == 2 { + wsc.Close() // close while the cycle is mid-flight, before the switch + } + return nil + } + + wsConfig := cycleTestConfig(url) + wsConfig.ConnectionCycleInterval = 25 * time.Millisecond + + wsc, err := New(context.Background(), wsConfig, nil, afterConnect) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + // Both the original connection, and the orphaned new one, are cleaned up + conn1 := nextConn(t, connections) + conn2 := nextConn(t, connections) + waitDone(t, conn1) + waitDone(t, conn2) + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCycleOldConnFailsDuringQuiesce(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := cycleTestConfig(url) + wsConfig.ConnectionCycleInterval = 25 * time.Millisecond + wsConfig.ConnectionCycleQuiesceTime = 1 * time.Minute + afterConnect := func(ctx context.Context, w WSClient) error { + return w.Send(ctx, []byte(`subscribe`)) + } + + wsc, err := New(context.Background(), wsConfig, nil, afterConnect) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + expectMsg(t, conn1, `subscribe`) + c1 := wsc.(*wsClient).currentConn() + + conn2 := nextConn(t, connections) + expectMsg(t, conn2, `subscribe`) + waitForSwitch(t, wsc, c1) // the old connection is now quiescing + + // The server kills the old connection during the (very long) quiesce - the cycle + // completes early rather than waiting for the quiesce timer, proven by the next + // cycle establishing a third connection well within the quiesce time + conn1.CloseConn() + waitDone(t, conn1) + conn3 := nextConn(t, connections) + expectMsg(t, conn3, `subscribe`) + + wsc.Close() + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCycleTimerResetOnErrorReconnect(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + var preDisconnectCount atomic.Int32 + wsConfig := cycleTestConfig(url) + wsConfig.ConnectionCycleInterval = 300 * time.Millisecond + wsConfig.ConnectionCycleQuiesceTime = 25 * time.Millisecond + wsConfig.PreDisconnectHandler = func(ctx context.Context, w WSClient) error { + preDisconnectCount.Add(1) + return w.Send(ctx, []byte(`unsubscribe`)) + } + afterConnect := func(ctx context.Context, w WSClient) error { + return w.Send(ctx, []byte(`subscribe`)) + } + + wsc, err := New(context.Background(), wsConfig, nil, afterConnect) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + expectMsg(t, conn1, `subscribe`) + + // Kill the connection - an error-driven reconnect, which does NOT run the + // pre-disconnect handler, and resets the cycle timer + conn1.CloseConn() + conn2 := nextConn(t, connections) + expectMsg(t, conn2, `subscribe`) + assert.Equal(t, int32(0), preDisconnectCount.Load()) + + // The next planned cycle from the reconnected connection does run it + conn3 := nextConn(t, connections) + expectMsg(t, conn2, `unsubscribe`) + expectMsg(t, conn3, `subscribe`) + assert.Equal(t, int32(1), preDisconnectCount.Load()) + + wsc.Close() + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCycleCloseDuringDialRetry(t *testing.T) { + connections, url, rejectNext, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := cycleTestConfig(url) + wsConfig.ConnectionCycleInterval = 25 * time.Millisecond + + wsc, err := New(context.Background(), wsConfig, nil, nil) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + + // The cycle fires and goes into dial retry - then we close mid-retry + rejectNext(true) + time.Sleep(100 * time.Millisecond) + wsc.Close() + + waitDone(t, conn1) + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCycleQuiesceZero(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := cycleTestConfig(url) + wsConfig.ConnectionCycleInterval = 50 * time.Millisecond + wsConfig.ConnectionCycleQuiesceTime = 0 // immediate close after the switch + afterConnect := func(ctx context.Context, w WSClient) error { + return w.Send(ctx, []byte(`subscribe`)) + } + + wsc, err := New(context.Background(), wsConfig, nil, afterConnect) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + expectMsg(t, conn1, `subscribe`) + + conn2 := nextConn(t, connections) + expectMsg(t, conn2, `subscribe`) + waitDone(t, conn1) + + err = wsc.Send(context.Background(), []byte(`on new`)) + assert.NoError(t, err) + expectMsg(t, conn2, `on new`) + + wsc.Close() + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCycleInactiveWithDisableReconnect(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(func(req *http.Request) { + assert.Equal(t, "/", req.URL.Path) + }) + defer done() + + // A plain HTTP request fails the websocket upgrade, and is ignored by the test server + res, err := http.Get(strings.Replace(url, "ws://", "http://", 1)) + assert.NoError(t, err) + res.Body.Close() + + wsConfig := cycleTestConfig(url) + wsConfig.ConnectionCycleInterval = 20 * time.Millisecond + wsConfig.DisableReconnect = true // warns, and cycling is inactive + + wsc, err := New(context.Background(), wsConfig, nil, nil) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + select { + case <-connections: + t.Fatal("unexpected second connection with cycling inactive") + case <-time.After(150 * time.Millisecond): + } + + wsc.Close() + waitDone(t, conn1) + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSReconnectAfterConnectFailureCleansUp(t *testing.T) { + // Not a cycling test (cycling disabled) - covers the teardown of a connection whose + // afterConnect fails on an error-driven reconnect, before dialing again + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + var connCount atomic.Int32 + afterConnect := func(ctx context.Context, w WSClient) error { + if connCount.Add(1) == 2 { + return fmt.Errorf("pop") + } + return w.Send(ctx, []byte(`subscribe`)) + } + + wsConfig := cycleTestConfig(url) + wsConfig.ConnectionCycleInterval = 0 // disabled + + wsc, err := New(context.Background(), wsConfig, nil, afterConnect) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + expectMsg(t, conn1, `subscribe`) + + // Kill the connection - the reconnect's afterConnect fails, so that connection is + // torn down and the client reconnects again + conn1.CloseConn() + conn2 := nextConn(t, connections) + waitDone(t, conn2) + + conn3 := nextConn(t, connections) + expectMsg(t, conn3, `subscribe`) + + err = wsc.Send(context.Background(), []byte(`hello`)) + assert.NoError(t, err) + expectMsg(t, conn3, `hello`) + + wsc.Close() + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSClosePreDisconnect(t *testing.T) { + // The pre-disconnect handler runs on Close() too - even with cycling disabled - so + // all pre-close processing lives in the one handler + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := &WSConfig{ + HTTPURL: url, + PreDisconnectHandler: func(ctx context.Context, w WSClient) error { + return w.Send(ctx, []byte(`unsubscribe`)) + }, + } + + wsc, err := New(context.Background(), wsConfig, nil, nil) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + + wsc.Close() + expectMsg(t, conn1, `unsubscribe`) + waitDone(t, conn1) + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCloseFromPreDisconnectNoDeadlock(t *testing.T) { + connections, url, _, done := NewTestWSServerMulti(nil) + defer done() + + wsConfig := &WSConfig{HTTPURL: url} + wsConfig.PreDisconnectHandler = func(ctx context.Context, w WSClient) error { + w.Close() // re-entrant Close proceeds with the close - no recursion, no deadlock + return w.Send(ctx, []byte(`too late`)) // fails, as the client is now closing + } + + wsc, err := New(context.Background(), wsConfig, nil, nil) + assert.NoError(t, err) + err = wsc.Connect() + assert.NoError(t, err) + + conn1 := nextConn(t, connections) + + wsc.Close() + waitDone(t, conn1) + _, ok := <-wsc.Receive() + assert.False(t, ok) +} + +func TestWSCloseNeverConnectedSkipsPreDisconnect(t *testing.T) { + var preDisconnectCount atomic.Int32 + wsConfig := &WSConfig{ + HTTPURL: "http://localhost:12345", + PreDisconnectHandler: func(ctx context.Context, w WSClient) error { + preDisconnectCount.Add(1) + return nil + }, + } + + wsc, err := New(context.Background(), wsConfig, nil, nil) + assert.NoError(t, err) + + // Never connected - there is no connection to run the handler against + wsc.Close() + assert.Equal(t, int32(0), preDisconnectCount.Load()) +} + +func TestWSSendRetriesOnConnectionCycle(t *testing.T) { + w := &wsClient{ + ctx: context.Background(), + send: make(chan []byte), + closing: make(chan struct{}), + } + c1 := &wsConnection{w: w, sendDone: make(chan []byte)} + c2 := &wsConnection{w: w, sendDone: make(chan []byte)} + w.current = c1 + + sendComplete := make(chan error, 1) + go func() { + sendComplete <- w.Send(context.Background(), []byte(`hello`)) + }() + time.Sleep(20 * time.Millisecond) // let Send block in its select + + // Simulate a connection cycle - the new connection takes over the shared send + // channel, and the old connection's sendLoop exits. Send must retry against the + // new connection rather than failing. + w.stateMux.Lock() + w.current = c2 + w.stateMux.Unlock() + close(c1.sendDone) + + assert.Equal(t, `hello`, string(<-w.send)) + assert.NoError(t, <-sendComplete) +} + +func TestConnBoundClientSendErrors(t *testing.T) { + w := &wsClient{ + ctx: context.Background(), + closing: make(chan struct{}), + } + c := &wsConnection{w: w, send: make(chan *trackedSend), sendDone: make(chan []byte)} + + // Rate limiter failure + w.rateLimiter = rate.NewLimiter(rate.Limit(1), 0) + err := w.boundTo(c).Send(context.Background(), []byte(`a`)) + assert.Regexp(t, "burst", err) + w.rateLimiter = nil + + // Context cancelled + cancelled, cancel := context.WithCancel(context.Background()) + cancel() + err = w.boundTo(c).Send(cancelled, []byte(`a`)) + assert.Regexp(t, "FF00146", err) + + // Connection retired (its sendLoop exited) + close(c.sendDone) + err = w.boundTo(c).Send(context.Background(), []byte(`a`)) + assert.Regexp(t, "FF00147", err) + + // Client closing + c2 := &wsConnection{w: w, send: make(chan *trackedSend), sendDone: make(chan []byte)} + close(w.closing) + err = w.boundTo(c2).Send(context.Background(), []byte(`a`)) + assert.Regexp(t, "FF00147", err) +} + +func TestConnBoundClientSendWriteFailure(t *testing.T) { + _, url, _, done := NewTestWSServerMulti(nil) + defer done() + + wsconn, _, err := websocket.DefaultDialer.Dial(url, nil) + assert.NoError(t, err) + wsconn.Close() // the write will fail + + w := &wsClient{ + ctx: context.Background(), + closing: make(chan struct{}), + } + c := w.newConnection(wsconn) + go c.sendLoop() + defer close(c.receiverDone) + + err = w.boundTo(c).Send(context.Background(), []byte(`fails to write`)) + assert.Regexp(t, "FF00147", err) +} + +func TestConnBoundClientSendWriteTimeout(t *testing.T) { + w := &wsClient{ + ctx: context.Background(), + closing: make(chan struct{}), + } + c := &wsConnection{w: w, send: make(chan *trackedSend), sendDone: make(chan []byte)} + go func() { <-c.send }() // accepts the handoff, but the write never completes + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + err := w.boundTo(c).Send(ctx, []byte(`never written`)) + assert.Regexp(t, "FF00146", err) +} diff --git a/pkg/wsclient/wsclient_test.go b/pkg/wsclient/wsclient_test.go index a659fab..a6ec2e2 100644 --- a/pkg/wsclient/wsclient_test.go +++ b/pkg/wsclient/wsclient_test.go @@ -22,6 +22,7 @@ import ( "crypto/x509" "fmt" "io" + "net" "net/http" "net/http/httptest" "os" @@ -41,6 +42,17 @@ func generateConfig() *WSConfig { return &WSConfig{} } +// lastPingCompleted reads the heartbeat state of the current connection race-safely +func lastPingCompleted(wsc WSClient) time.Time { + c := wsc.(*wsClient).currentConn() + if c == nil { + return time.Time{} + } + c.heartbeatMux.Lock() + defer c.heartbeatMux.Unlock() + return c.lastPingCompleted +} + func TestWSClientE2ETLS(t *testing.T) { publicKeyFile, privateKeyFile := GenerateTLSCertficates(t) @@ -122,7 +134,7 @@ func TestWSClientE2ETLS(t *testing.T) { // Check heartbeating works beforePing := time.Now() - for wsc.(*wsClient).lastPingCompleted.Before(beforePing) { + for lastPingCompleted(wsc).Before(beforePing) { time.Sleep(10 * time.Millisecond) } @@ -185,7 +197,7 @@ func TestWSClientE2EBG(t *testing.T) { // Check heartbeating works beforePing := time.Now() - for wsc.(*wsClient).lastPingCompleted.Before(beforePing) { + for lastPingCompleted(wsc).Before(beforePing) { time.Sleep(10 * time.Millisecond) } @@ -253,7 +265,7 @@ func TestWSClientE2EReceiveExt(t *testing.T) { // Check heartbeating works beforePing := time.Now() - for wsc.(*wsClient).lastPingCompleted.Before(beforePing) { + for lastPingCompleted(wsc).Before(beforePing) { time.Sleep(10 * time.Millisecond) } @@ -368,6 +380,17 @@ func TestWSClosedWhileConnecting(t *testing.T) { assert.Regexp(t, "FF00147", err) } +func TestWSClientNetDialer(t *testing.T) { + wsConfig := generateConfig() + wsConfig.HTTPURL = "http://test:12345" + wsConfig.NetDialer = &net.Dialer{} + + wsc, err := NewWithConfig(context.Background(), wsConfig) + assert.NoError(t, err) + defer wsc.Close() + assert.NotNil(t, wsc.(*wsClient).wsdialer.NetDialContext) +} + func TestWSClientBadWSURL(t *testing.T) { wsConfig := generateConfig() wsConfig.WebSocketURL = ":::" @@ -489,10 +512,14 @@ func TestWSSendCanceledContext(t *testing.T) { func TestWSSenderClosed(t *testing.T) { w := &wsClient{ - send: make(chan []byte), + send: make(chan []byte), + } + c := &wsConnection{ + w: w, sendDone: make(chan []byte), } - close(w.sendDone) + w.current = c + close(c.sendDone) err := w.Send(context.Background(), []byte(`sent after close`)) assert.Regexp(t, "FF00147", err) @@ -520,17 +547,20 @@ func TestWSReadLoopSendFailure(t *testing.T) { assert.NoError(t, err) <-toServer w := &wsClient{ - ctx: context.Background(), + ctx: context.Background(), + } + c := &wsConnection{ + w: w, + conn: wsconn, sendDone: make(chan []byte, 1), - wsconn: wsconn, } // Queue a message for the receiver, then immediately close the sender channel fromServer <- `some data from server` - close(w.sendDone) + close(c.sendDone) // Ensure the readLoop exits immediately - w.readLoop() + c.readLoop() // Try reconnect, should fail here _, _, err = websocket.DefaultDialer.Dial(url, nil) @@ -550,17 +580,20 @@ func TestWSReadLoopExtSendFailure(t *testing.T) { <-toServer w := &wsClient{ ctx: context.Background(), - sendDone: make(chan []byte, 1), - wsconn: wsconn, useReceiveExt: true, } + c := &wsConnection{ + w: w, + conn: wsconn, + sendDone: make(chan []byte, 1), + } // Queue a message for the receiver, then immediately close the sender channel fromServer <- `some data from server` - close(w.sendDone) + close(c.sendDone) // Ensure the readLoop exits immediately - w.readLoopExt() + c.readLoopExt() // Try reconnect, should fail here _, _, err = websocket.DefaultDialer.Dial(url, nil) @@ -580,11 +613,14 @@ func TestWSReadLoopExtProcessedFailure(t *testing.T) { <-toServer w := &wsClient{ ctx: context.Background(), - sendDone: make(chan []byte, 1), - wsconn: wsconn, receiveExt: make(chan *WSPayload), useReceiveExt: true, } + c := &wsConnection{ + w: w, + conn: wsconn, + sendDone: make(chan []byte, 1), + } // Queue a message for the receiver, then immediately close the sender channel fromServer <- `some data from server` @@ -592,12 +628,12 @@ func TestWSReadLoopExtProcessedFailure(t *testing.T) { // Ensure the readLoop exits immediately readLoopDone := make(chan struct{}) go func() { - w.readLoopExt() + c.readLoopExt() close(readLoopDone) }() _ = <-w.receiveExt // we don't bother to ack this, making the client wait indefinitely - close(w.sendDone) + close(c.sendDone) <-readLoopDone // Try reconnect, should fail here @@ -622,8 +658,10 @@ func TestWSReconnectFail(t *testing.T) { receiveExt: make(chan *WSPayload), send: make(chan []byte), closing: make(chan struct{}), - wsconn: wsconn, } + c := w.newConnection(wsconn) + close(c.promoted) + w.current = c close(w.send) // will mean sender exits immediately w.receiveReconnectLoop() @@ -645,9 +683,11 @@ func TestWSDisableReconnect(t *testing.T) { receiveExt: make(chan *WSPayload), send: make(chan []byte), closing: make(chan struct{}), - wsconn: wsconn, disableReconnect: true, } + c := w.newConnection(wsconn) + close(c.promoted) + w.current = c w.receiveReconnectLoop() } @@ -666,12 +706,13 @@ func TestWSSendFail(t *testing.T) { receiveExt: make(chan *WSPayload), send: make(chan []byte, 1), closing: make(chan struct{}), - sendDone: make(chan []byte, 1), - wsconn: wsconn, } + c := w.newConnection(wsconn) + close(c.promoted) + w.current = c w.send <- []byte(`wakes sender`) - w.sendLoop(make(chan struct{})) - <-w.sendDone + c.sendLoop() + <-c.sendDone } func TestWSSendInstructClose(t *testing.T) { @@ -688,13 +729,11 @@ func TestWSSendInstructClose(t *testing.T) { receiveExt: make(chan *WSPayload), send: make(chan []byte, 1), closing: make(chan struct{}), - sendDone: make(chan []byte, 1), - wsconn: wsconn, } - receiverClosed := make(chan struct{}) - close(receiverClosed) - w.sendLoop(receiverClosed) - <-w.sendDone + c := w.newConnection(wsconn) + close(c.receiverDone) + c.sendLoop() + <-c.sendDone } func TestHeartbeatTimedout(t *testing.T) { @@ -702,12 +741,19 @@ func TestHeartbeatTimedout(t *testing.T) { now := time.Now() w := &wsClient{ ctx: context.Background(), - sendDone: make(chan []byte), heartbeatInterval: 1 * time.Microsecond, - activePingSent: &now, + } + c := &wsConnection{ + w: w, + send: make(chan *trackedSend), + sendDone: make(chan []byte), + receiverDone: make(chan struct{}), + promoted: make(chan struct{}), + demoted: make(chan struct{}), + activePingSent: &now, } - w.sendLoop(make(chan struct{})) + c.sendLoop() } @@ -723,16 +769,24 @@ func TestHeartbeatSendFailed(t *testing.T) { err = wsc.Connect() assert.NoError(t, err) - // Close and use the underlying wsconn to drive a failure to send a heartbeat - wsc.(*wsClient).wsconn.Close() + // Close and use the underlying connection to drive a failure to send a heartbeat + wsconn := wsc.(*wsClient).currentConn().conn + wsconn.Close() w := &wsClient{ ctx: context.Background(), - sendDone: make(chan []byte), heartbeatInterval: 1 * time.Microsecond, - wsconn: wsc.(*wsClient).wsconn, + } + c := &wsConnection{ + w: w, + conn: wsconn, + send: make(chan *trackedSend), + sendDone: make(chan []byte), + receiverDone: make(chan struct{}), + promoted: make(chan struct{}), + demoted: make(chan struct{}), } - w.sendLoop(make(chan struct{})) + c.sendLoop() } @@ -824,7 +878,7 @@ func TestRequestWithRateLimiter(t *testing.T) { for i := 0; i < expectedNumberOfRequest; i++ { go func() { // Send some data back - err = wsc.Send(context.Background(), []byte(`some data to server`)) + err := wsc.Send(context.Background(), []byte(`some data to server`)) assert.NoError(t, err) // Check the sevrer got it @@ -893,7 +947,7 @@ func TestRequestWithRateLimiterHighBurst(t *testing.T) { for i := 0; i < expectedNumberOfRequest; i++ { go func() { // Send some data back - err = wsc.Send(context.Background(), []byte(`some data to server`)) + err := wsc.Send(context.Background(), []byte(`some data to server`)) assert.NoError(t, err) // Check the sevrer got it diff --git a/pkg/wsclient/wsconfig.go b/pkg/wsclient/wsconfig.go index 41974f8..243a341 100644 --- a/pkg/wsclient/wsconfig.go +++ b/pkg/wsclient/wsconfig.go @@ -18,6 +18,7 @@ package wsclient import ( "context" + "net" "time" "github.com/hyperledger-firefly/common/pkg/config" @@ -28,11 +29,12 @@ import ( ) const ( - defaultInitialConnectAttempts = 5 - defaultBufferSize = "16Kb" - defaultHeartbeatInterval = "30s" // up to a minute to detect a dead connection - defaultConnectionTimeout = 45 * time.Second // 45 seconds - the built in default for gorilla/websocket - defaultRetryBackoffFactor = 2.0 + defaultInitialConnectAttempts = 5 + defaultBufferSize = "16Kb" + defaultHeartbeatInterval = "30s" // up to a minute to detect a dead connection + defaultConnectionTimeout = 45 * time.Second // 45 seconds - the built in default for gorilla/websocket + defaultRetryBackoffFactor = 2.0 + defaultConnectionCycleQuiesceTime = "5s" // only used when connection cycling is enabled ) const ( @@ -54,6 +56,10 @@ const ( WSConfigKeyHeartbeatInterval = "ws.heartbeatInterval" // WSConnectionTimeout is the amount of time to wait while attempting to establish a connection (or automatic reconnection) WSConfigKeyConnectionTimeout = "ws.connectionTimeout" + // WSConfigKeyConnectionCycleInterval when non-zero enables proactive cycling of the connection on this interval - the new connection is fully established before the old one is quiesced and closed + WSConfigKeyConnectionCycleInterval = "ws.connectionCycleInterval" + // WSConfigKeyConnectionCycleQuiesceTime is how long the old connection continues to deliver inbound messages after a connection cycle, before it is closed + WSConfigKeyConnectionCycleQuiesceTime = "ws.connectionCycleQuiesceTime" // WSConfigDelayFactor the exponential backoff factor for delay WSConfigDelayFactor = "retry.factor" ) @@ -77,6 +83,8 @@ func InitConfig(conf config.Section) { conf.AddKnownKey(WSConfigURL) conf.AddKnownKey(WSConfigKeyHeartbeatInterval, defaultHeartbeatInterval) conf.AddKnownKey(WSConfigKeyConnectionTimeout, defaultConnectionTimeout) + conf.AddKnownKey(WSConfigKeyConnectionCycleInterval) // no default - connection cycling is disabled unless set + conf.AddKnownKey(WSConfigKeyConnectionCycleQuiesceTime, defaultConnectionCycleQuiesceTime) conf.AddKnownKey(WSConfigDelayFactor, defaultRetryBackoffFactor) InitConfigWrap(conf) } @@ -102,6 +110,9 @@ func GenerateConfig(ctx context.Context, conf config.Section) (*WSConfig, error) AuthPassword: conf.GetString(ffresty.HTTPConfigAuthPassword), HeartbeatInterval: conf.GetDuration(WSConfigKeyHeartbeatInterval), ConnectionTimeout: conf.GetDuration(WSConfigKeyConnectionTimeout), + + ConnectionCycleInterval: conf.GetDuration(WSConfigKeyConnectionCycleInterval), + ConnectionCycleQuiesceTime: conf.GetDuration(WSConfigKeyConnectionCycleQuiesceTime), } tlsSection := conf.SubSection("tls") tlsClientConfig, err := fftls.ConstructTLSConfig(ctx, tlsSection, fftls.ClientType) @@ -113,18 +124,12 @@ func GenerateConfig(ctx context.Context, conf config.Section) (*WSConfig, error) // Build the underlying TCP dialer with the custom DNS resolver and SSRF egress guard, // from the same "net" subsection that ffresty.InitConfig set up on this config tree. + var netDialer *net.Dialer + resolver := ffdns.NewResolver(conf.SubSection("dns")) netCfg, err := ffnet.GenerateConfig(conf.SubSection("net")) - if err != nil { - return nil, err - } - - dnsCfg, err := ffdns.GenerateConfig(conf.SubSection("dns")) - if err != nil { - return nil, err + if err == nil { + netDialer, err = ffnet.NewDialer(ctx, netCfg, resolver) } - resolver := ffdns.NewResolverWithConfig(dnsCfg) - - netDialer, err := ffnet.NewDialer(ctx, netCfg, resolver) if err != nil { return nil, err } diff --git a/pkg/wsclient/wsconfig_test.go b/pkg/wsclient/wsconfig_test.go index e0d4aa4..6725c8e 100644 --- a/pkg/wsclient/wsconfig_test.go +++ b/pkg/wsclient/wsconfig_test.go @@ -104,6 +104,15 @@ func TestWSConfigNetDialerCustom(t *testing.T) { assert.Error(t, wsConfig.NetDialer.Control("tcp", "169.254.169.254:80", nil)) } +func TestWSConfigNetDialerFail(t *testing.T) { + resetConf() + utConf.SubSection("net").Set(ffnet.NetCIDRDenylist, []string{"not-a-cidr"}) + + ctx := context.Background() + _, err := GenerateConfig(ctx, utConf) + assert.Regexp(t, "FF00260", err) +} + func TestWSConfigTLSGenerationFail(t *testing.T) { resetConf() diff --git a/pkg/wsclient/wsconnection.go b/pkg/wsclient/wsconnection.go new file mode 100644 index 0000000..b25bb86 --- /dev/null +++ b/pkg/wsclient/wsconnection.go @@ -0,0 +1,323 @@ +// Copyright © 2026 Kaleido, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package wsclient + +import ( + "context" + "sync" + "sync/atomic" + "time" + + "github.com/gorilla/websocket" + "github.com/hyperledger-firefly/common/pkg/i18n" + "github.com/hyperledger-firefly/common/pkg/log" +) + +// wsConnection holds all the state for a single underlying WebSocket connection. +// Normally a wsClient has exactly one of these at a time, but during a planned +// connection cycle (connectionCycleInterval) two coexist briefly - the old +// connection quiescing while the new connection takes over. +type wsConnection struct { + w *wsClient + conn *websocket.Conn + + send chan *trackedSend // connection-bound sends (used by the hook facade) + sendDone chan []byte // closed by this connection's sendLoop on exit + receiverDone chan struct{} // closed by teardown, telling the sendLoop to exit + readDone chan struct{} // closed when the read loop goroutine returns + promoted chan struct{} // closed when this connection becomes the consumer of the shared send channel + demoted chan struct{} // closed when this connection stops being the current one (connection cycling) + senderStarted bool // only accessed on the orchestrating goroutine (receiveReconnectLoop) + readerStarted bool // only accessed on the orchestrating goroutine (receiveReconnectLoop) + closeOnce sync.Once + teardownOnce sync.Once + + // The preDisconnect handler runs at most once per connection - claimed either by a + // connection cycle retiring this connection, or by Close() (whichever happens first) + preDisconnected atomic.Bool + + // Heartbeat state is per-connection, so that during a connection cycle a pong on the + // old connection cannot affect the new connection's ping bookkeeping or read deadline + heartbeatMux sync.Mutex + activePingSent *time.Time + lastPingCompleted time.Time + quiesceDeadline time.Time // set when the connection is demoted during a connection cycle +} + +// trackedSend is a send on this connection, with a success/failure notification, +// so the caller can block until the send is either on the wire or confirmed as failed. +type trackedSend struct { + message []byte + sent chan bool // buffered(1) - receives true if the write succeeded +} + +func (w *wsClient) newConnection(conn *websocket.Conn) *wsConnection { + c := &wsConnection{ + w: w, + conn: conn, + send: make(chan *trackedSend), + sendDone: make(chan []byte, 1), + receiverDone: make(chan struct{}), + readDone: make(chan struct{}), + promoted: make(chan struct{}), + demoted: make(chan struct{}), + } + c.pongReceivedOrReset(false) + conn.SetPongHandler(c.pongHandler) + return c +} + +func (c *wsConnection) startSender() { + c.senderStarted = true + go c.sendLoop() +} + +func (c *wsConnection) startReader() { + c.readerStarted = true + go func() { + defer close(c.readDone) + if c.w.useReceiveExt { + c.readLoopExt() + } else { + c.readLoop() + } + }() +} + +// claimPreDisconnect returns true for exactly one caller per connection, electing it to +// run the preDisconnect handler for this connection +func (c *wsConnection) claimPreDisconnect() bool { + return c.preDisconnected.CompareAndSwap(false, true) +} + +// closeConn closes the underlying connection exactly once (teardown, quiesce completion, +// and Close() can all race to do this) +func (c *wsConnection) closeConn() { + c.closeOnce.Do(func() { + if err := c.conn.Close(); err != nil { + log.L(c.w.ctx).Warnf("WS %s ignoring websocket connection close error: %v", c.w.url, err) + } + }) +} + +// teardownConnection fully stops a connection exactly once, closing the underlying +// connection and blocking until any started send/read loops have exited - so after +// this returns nothing can be delivered from this connection to the receive channel. +func (w *wsClient) teardownConnection(c *wsConnection) { + c.teardownOnce.Do(func() { + close(c.receiverDone) // tells the sendLoop to exit + c.closeConn() // unblocks a reader blocked in ReadMessage/NextReader + if c.readerStarted { + // A reader blocked delivering to the receive channel is unblocked by the close + // of sendDone (from the sendLoop exiting on receiverDone above) + <-c.readDone + } + if c.senderStarted { + <-c.sendDone + } + }) +} + +func (c *wsConnection) readLoop() { + w := c.w + l := log.L(w.ctx) + for { + mt, message, err := c.conn.ReadMessage() + if err != nil { + // We treat this as informational, as it's normal for the client to disconnect here + l.Infof("WS %s closed: %s", w.url, err) + return + } + + // Pass the message to the consumer + l.Tracef("WS %s read (mt=%d): %s", w.url, mt, message) + select { + case <-c.sendDone: + l.Debugf("WS %s closing reader after send error", w.url) + return + case w.receive <- message: + } + } +} + +func (c *wsConnection) readLoopExt() { + w := c.w + l := log.L(w.ctx) + for { + // We set a deadline for twice the heartbeat interval - note we bump this on pong + if deadline, hasDeadline := c.nextReadDeadline(); hasDeadline { + _ = c.conn.SetReadDeadline(deadline) + } + + mt, r, err := c.conn.NextReader() + if err != nil { + // We treat this as informational, as it's normal for the client to disconnect here + l.Infof("WS %s closed: %s", w.url, err) + return + } + + // Pass the message to the consumer + l.Tracef("WS %s read (mt=%d)", w.url, mt) + payload := NewWSPayload(mt, r) + select { + case <-c.sendDone: + l.Debugf("WS %s closing reader after send error (waiting for data)", w.url) + return + case w.receiveExt <- payload: + } + select { + case <-payload.processed: + // It's the callers responsibility to ensure they call done on this before we can get the next payload + case <-c.sendDone: + l.Debugf("WS %s closing reader after send error (waiting for processing of data by client)", w.url) + return + } + } +} + +func (c *wsConnection) sendLoop() { + w := c.w + l := log.L(w.ctx) + defer close(c.sendDone) + + var sharedSend chan []byte // nil (never selected) until this connection is promoted + promoted := c.promoted + demoted := c.demoted + isDemoted := false + disconnecting := false + for !disconnecting { + timeoutContext, timeoutCancel := c.heartbeatTimeout(w.ctx, isDemoted) + + select { + case message := <-sharedSend: + disconnecting = c.writeText(message) + case bs := <-c.send: + disconnecting = c.writeText(bs.message) + bs.sent <- !disconnecting + case <-promoted: + // We are now the current connection - consume the shared send channel + sharedSend = w.send + promoted = nil + case <-demoted: + // A connection cycle has replaced us - stop consuming the shared send + // channel, and stop heartbeating (we just quiesce until closed) + sharedSend = nil + demoted = nil + isDemoted = true + case <-timeoutContext.Done(): + if err := c.heartbeatCheck(); err != nil { + l.Errorf("WS %s closing: %s", w.url, err) + disconnecting = true + } else { + l.Debugf("WS %s send heartbeat ping", w.url) + if err := c.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil { + l.Errorf("WS %s heartbeat send failed: %s", w.url, err) + disconnecting = true + } + } + case <-c.receiverDone: + l.Debugf("WS %s send loop exiting", w.url) + disconnecting = true + } + + timeoutCancel() + } +} + +func (c *wsConnection) writeText(message []byte) (disconnecting bool) { + l := log.L(c.w.ctx) + l.Tracef("WS sending: %s", message) + if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil { + l.Errorf("WS %s send failed: %s", c.w.url, err) + return true + } + return false +} + +func (c *wsConnection) pongHandler(_ string) error { + c.pongReceivedOrReset(true) + return nil +} + +func (c *wsConnection) pongReceivedOrReset(isPong bool) { + c.heartbeatMux.Lock() + if isPong && c.activePingSent != nil { + log.L(c.w.ctx).Debugf("WS %s heartbeat completed (pong) after %.2fms", c.w.url, float64(time.Since(*c.activePingSent))/float64(time.Millisecond)) + } + c.lastPingCompleted = time.Now() // in new connection case we still want to consider now the time we completed the ping + c.activePingSent = nil + c.heartbeatMux.Unlock() + + // We set a deadline for twice the heartbeat interval + if deadline, hasDeadline := c.nextReadDeadline(); hasDeadline { + _ = c.conn.SetReadDeadline(deadline) + } +} + +// nextReadDeadline returns the deadline to use for the next read (heartbeat or regular). +// Caller must set it onto the c.conn to activate it. +func (c *wsConnection) nextReadDeadline() (deadline time.Time, hasDeadline bool) { + c.heartbeatMux.Lock() + defer c.heartbeatMux.Unlock() + + if !c.quiesceDeadline.IsZero() { + return c.quiesceDeadline, true // quiesce deadline wins over heartbeat + } + if c.w.heartbeatInterval > 0 { + return time.Now().Add(2 * c.w.heartbeatInterval), true + } + return time.Time{}, false +} + +// setQuiesceDeadline calculates and applies the quiesce deadline for future reads +func (c *wsConnection) setQuiesceDeadline(quiesceTime time.Duration) { + c.heartbeatMux.Lock() + c.quiesceDeadline = time.Now().Add(quiesceTime + 500*time.Millisecond) + c.heartbeatMux.Unlock() + + if deadline, hasDeadline := c.nextReadDeadline(); hasDeadline { + _ = c.conn.SetReadDeadline(deadline) + } +} + +func (c *wsConnection) heartbeatCheck() error { + c.heartbeatMux.Lock() + defer c.heartbeatMux.Unlock() + + if c.activePingSent != nil { + return i18n.NewError(c.w.ctx, i18n.MsgWSHeartbeatTimeout, float64(time.Since(*c.activePingSent))/float64(time.Millisecond)) + } + log.L(c.w.ctx).Debugf("WS %s heartbeat timer popped (ping) after %.2fms", c.w.url, float64(time.Since(c.lastPingCompleted))/float64(time.Millisecond)) + now := time.Now() + c.activePingSent = &now + return nil +} + +func (c *wsConnection) heartbeatTimeout(ctx context.Context, demoted bool) (context.Context, context.CancelFunc) { + if c.w.heartbeatInterval > 0 && !demoted { + c.heartbeatMux.Lock() + baseTime := c.lastPingCompleted + if c.activePingSent != nil { + // We're waiting for a pong + baseTime = *c.activePingSent + } + waitTime := c.w.heartbeatInterval - time.Since(baseTime) // if negative, will pop immediately + c.heartbeatMux.Unlock() + return context.WithTimeout(ctx, waitTime) + } + return context.WithCancel(ctx) +} diff --git a/pkg/wsclient/wstestserver.go b/pkg/wsclient/wstestserver.go index da8f1e9..e7a204c 100644 --- a/pkg/wsclient/wstestserver.go +++ b/pkg/wsclient/wstestserver.go @@ -141,6 +141,88 @@ func NewTestTLSWSServer(testReq func(req *http.Request), publicKeyFile *os.File, }, nil } +// TestWSConnection is a single server-side connection accepted by NewTestWSServerMulti +type TestWSConnection struct { + ToServer chan string // messages the server received on this connection + FromServer chan string // push messages here to send them to the client on this connection + Done chan struct{} // closed when this connection has closed + CloseConn func() // server-side close of this connection +} + +// NewTestWSServerMulti creates a test server that accepts an unlimited sequence of connections, +// emitting each accepted connection on the connections channel - allowing tests of reconnect +// and connection cycling, where more than one connection can be active at the same time +func NewTestWSServerMulti(testReq func(req *http.Request)) (connections chan *TestWSConnection, url string, rejectNext func(bool), done func()) { + upgrader := &websocket.Upgrader{WriteBufferSize: 1024, ReadBufferSize: 1024} + connections = make(chan *TestWSConnection, 16) + mu := sync.Mutex{} + rejecting := false + var live []*websocket.Conn + svr := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) { + if testReq != nil { + testReq(req) + } + mu.Lock() + reject := rejecting + mu.Unlock() + if reject { + res.WriteHeader(500) + return + } + ws, err := upgrader.Upgrade(res, req, http.Header{}) + if err != nil { + return + } + mu.Lock() + live = append(live, ws) + mu.Unlock() + c := &TestWSConnection{ + ToServer: make(chan string, 16), + FromServer: make(chan string, 16), + Done: make(chan struct{}), + CloseConn: func() { _ = ws.Close() }, + } + go func() { + defer close(c.Done) + for { + _, data, err := ws.ReadMessage() + if err != nil { + return + } + c.ToServer <- string(data) + } + }() + go func() { + for { + select { + case data := <-c.FromServer: + _ = ws.WriteMessage(websocket.TextMessage, []byte(data)) + case <-c.Done: + return + } + } + }() + connections <- c + })) + return connections, + fmt.Sprintf("ws://%s", svr.Listener.Addr()), + func(reject bool) { + mu.Lock() + rejecting = reject + mu.Unlock() + }, + func() { + mu.Lock() + conns := live + live = nil + mu.Unlock() + for _, ws := range conns { + _ = ws.Close() + } + svr.Close() + } +} + // NewTestWSServer creates a little test server for packages (including wsclient itself) to use in unit tests func NewTestWSServer(testReq func(req *http.Request)) (toServer, fromServer chan string, url string, done func()) { upgrader := &websocket.Upgrader{WriteBufferSize: 1024, ReadBufferSize: 1024} @@ -149,13 +231,16 @@ func NewTestWSServer(testReq func(req *http.Request)) (toServer, fromServer chan sendDone := make(chan struct{}) receiveDone := make(chan struct{}) connected := false + mu := sync.Mutex{} svr := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) { + mu.Lock() if testReq != nil { testReq(req) } if connected { // test server only handles one open connection, as it only has one set of channels res.WriteHeader(409) + mu.Unlock() return } ws, _ := upgrader.Upgrade(res, req, http.Header{}) @@ -177,11 +262,15 @@ func NewTestWSServer(testReq func(req *http.Request)) (toServer, fromServer chan } }() connected = true + mu.Unlock() })) return toServer, fromServer, fmt.Sprintf("ws://%s", svr.Listener.Addr()), func() { close(fromServer) svr.Close() - if connected { + mu.Lock() + wasConnected := connected + mu.Unlock() + if wasConnected { <-sendDone <-receiveDone }