refactor(retry): remove sleep-based waits from websocket and peer tes…

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:

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