From b71d2a9c51a15c007440172b07e294a654315eec Mon Sep 17 00:00:00 2001 From: Laurence Date: Wed, 11 Feb 2026 08:34:12 +0000 Subject: [PATCH] refactor(retry): remove sleep-based waits from websocket and peer test loops - websocket/client: replace connectWithRetry sleep loop with timer-driven retry that can exit immediately on done - peers/monitor/wgtester: refactor TestPeerConnection to a straight happy path and replace timeout sleep with context-aware wait - keep retry/backoff behavior unchanged while improving shutdown/cancel responsiveness --- peers/monitor/wgtester.go | 99 ++++++++++++++++++++------------------- websocket/client.go | 69 +++++++++++++-------------- 2 files changed, 84 insertions(+), 84 deletions(-) diff --git a/peers/monitor/wgtester.go b/peers/monitor/wgtester.go index e9f6f63..62b96ba 100644 --- a/peers/monitor/wgtester.go +++ b/peers/monitor/wgtester.go @@ -182,66 +182,69 @@ func (c *Client) TestPeerConnection(ctx context.Context) (bool, time.Duration) { // Reusable response buffer responseBuffer := make([]byte, packetSize) - // Send multiple attempts as specified + const retryPause = 100 * time.Millisecond + waitForNextAttempt := func(delay time.Duration) bool { + if delay <= 0 { + return true + } + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } + } + for attempt := 0; attempt < c.maxAttempts; attempt++ { select { case <-ctx.Done(): return false, 0 default: - // Add current timestamp to packet - timestamp := time.Now().UnixNano() - binary.BigEndian.PutUint64(packet[5:13], uint64(timestamp)) - - // Lock the connection for the entire send/receive operation - c.connLock.Lock() - - // Check if connection is still valid after acquiring lock - if c.conn == nil { - c.connLock.Unlock() - return false, 0 - } - - _, err := c.conn.Write(packet) - if err != nil { - c.connLock.Unlock() - logger.Info("Error sending packet: %v", err) - continue - } + } - // Set read deadline - c.conn.SetReadDeadline(time.Now().Add(c.timeout)) + timestamp := time.Now().UnixNano() + binary.BigEndian.PutUint64(packet[5:13], uint64(timestamp)) - // Wait for response - n, err := c.conn.Read(responseBuffer) + c.connLock.Lock() + if c.conn == nil { c.connLock.Unlock() + return false, 0 + } - if err != nil { - if netErr, ok := err.(net.Error); ok && netErr.Timeout() { - // Timeout, try next attempt - time.Sleep(100 * time.Millisecond) // Brief pause between attempts - continue + _, err := c.conn.Write(packet) + if err != nil { + c.connLock.Unlock() + logger.Info("Error sending packet: %v", err) + continue + } + c.conn.SetReadDeadline(time.Now().Add(c.timeout)) + n, err := c.conn.Read(responseBuffer) + c.connLock.Unlock() + + if err != nil { + if netErr, ok := err.(net.Error); ok && netErr.Timeout() { + if !waitForNextAttempt(retryPause) { + return false, 0 } + } else { logger.Error("Error reading response: %v", err) - continue } + continue + } - if n != packetSize { - continue // Malformed packet - } - - // Verify response - magic := binary.BigEndian.Uint32(responseBuffer[0:4]) - packetType := responseBuffer[4] - if magic != magicHeader || packetType != packetTypeResponse { - continue // Not our response - } - - // Extract the original timestamp and calculate RTT - sentTimestamp := int64(binary.BigEndian.Uint64(responseBuffer[5:13])) - rtt := time.Duration(time.Now().UnixNano() - sentTimestamp) - - return true, rtt + if n != packetSize { + continue } + magic := binary.BigEndian.Uint32(responseBuffer[0:4]) + packetType := responseBuffer[4] + if magic != magicHeader || packetType != packetTypeResponse { + continue + } + sentTimestamp := int64(binary.BigEndian.Uint64(responseBuffer[5:13])) + rtt := time.Duration(time.Now().UnixNano() - sentTimestamp) + return true, rtt } return false, 0 @@ -294,10 +297,10 @@ func (c *Client) StartMonitor(callback MonitorCallback) error { c.monitorLock.Lock() currentInterval = c.minInterval c.monitorLock.Unlock() - + // Reset backoff state stableCount = 0 - + timer.Reset(currentInterval) logger.Debug("Packet interval updated, reset to %v", currentInterval) case <-timer.C: diff --git a/websocket/client.go b/websocket/client.go index c4e67b0..cdbcad3 100644 --- a/websocket/client.go +++ b/websocket/client.go @@ -92,17 +92,17 @@ type Client struct { configNeedsSave bool // Flag to track if config needs to be saved configVersion int // Latest config version received from server configVersionMux sync.RWMutex - token string // Cached authentication token - exitNodes []ExitNode // Cached exit nodes from token response - tokenMux sync.RWMutex // Protects token and exitNodes - forceNewToken bool // Flag to force fetching a new token on next connection - processingMessage bool // Flag to track if a message is currently being processed - processingMux sync.RWMutex // Protects processingMessage - processingWg sync.WaitGroup // WaitGroup to wait for message processing to complete - getPingData func() map[string]any // Callback to get additional ping data - pingStarted bool // Flag to track if ping monitor has been started - pingStartedMux sync.Mutex // Protects pingStarted - pingDone chan struct{} // Channel to stop the ping monitor independently + token string // Cached authentication token + exitNodes []ExitNode // Cached exit nodes from token response + tokenMux sync.RWMutex // Protects token and exitNodes + forceNewToken bool // Flag to force fetching a new token on next connection + processingMessage bool // Flag to track if a message is currently being processed + processingMux sync.RWMutex // Protects processingMessage + processingWg sync.WaitGroup // WaitGroup to wait for message processing to complete + getPingData func() map[string]any // Callback to get additional ping data + pingStarted bool // Flag to track if ping monitor has been started + pingStartedMux sync.Mutex // Protects pingStarted + pingDone chan struct{} // Channel to stop the ping monitor independently } type ClientOption func(*Client) @@ -462,32 +462,29 @@ func (c *Client) getToken() (string, []ExitNode, error) { } func (c *Client) connectWithRetry() { + timer := time.NewTimer(0) + defer timer.Stop() for { select { case <-c.done: return - default: - err := c.establishConnection() - if err != nil { - // Check if this is an auth error (401/403) - var authErr *AuthError - if errors.As(err, &authErr) { - logger.Error("Authentication failed: %v. Terminating tunnel and retrying...", authErr) - // Trigger auth error callback if set (this should terminate the tunnel) - if c.onAuthError != nil { - c.onAuthError(authErr.StatusCode, authErr.Message) - } - // Continue retrying after auth error - time.Sleep(c.reconnectInterval) - continue - } - // For other errors (5xx, network issues), continue retrying - logger.Error("websocket: Failed to connect: %v. Retrying in %v...", err, c.reconnectInterval) - time.Sleep(c.reconnectInterval) - continue - } + case <-timer.C: + // Timer fired (immediate first time, then after reconnectInterval) + } + err := c.establishConnection() + if err == nil { return } + var authErr *AuthError + if !errors.As(err, &authErr) { + logger.Error("websocket: Failed to connect: %v. Retrying in %v...", err, c.reconnectInterval) + } else { + logger.Error("Authentication failed: %v. Terminating tunnel and retrying...", authErr) + if c.onAuthError != nil { + c.onAuthError(authErr.StatusCode, authErr.Message) + } + } + timer.Reset(c.reconnectInterval) } } @@ -736,15 +733,15 @@ func (c *Client) pingMonitor() { func (c *Client) StartPingMonitor() { c.pingStartedMux.Lock() defer c.pingStartedMux.Unlock() - + if c.pingStarted { return } c.pingStarted = true - + // Create a new pingDone channel for this ping monitor instance c.pingDone = make(chan struct{}) - + // Send an initial ping immediately go func() { c.sendPing() @@ -756,11 +753,11 @@ func (c *Client) StartPingMonitor() { func (c *Client) stopPingMonitor() { c.pingStartedMux.Lock() defer c.pingStartedMux.Unlock() - + if !c.pingStarted { return } - + // Close the pingDone channel to stop the monitor close(c.pingDone) c.pingStarted = false