[relay-client] Fix reconnection

relay/client/client.go

@@ -141,7 +141,6 @@ type Client struct {
 	muInstanceURL    sync.Mutex
 
 	onDisconnectListener func(string)
-	onConnectedListener  func()
 	listenerMutex        sync.Mutex
 }
 
@@ -190,7 +189,6 @@ func (c *Client) Connect() error {
 
 	c.wgReadLoop.Add(1)
 	go c.readLoop(c.relayConn)
-	go c.notifyConnected()
 
 	return nil
 }
@@ -238,12 +236,6 @@ func (c *Client) SetOnDisconnectListener(fn func(string)) {
 	c.onDisconnectListener = fn
 }
 
-func (c *Client) SetOnConnectedListener(fn func()) {
-	c.listenerMutex.Lock()
-	defer c.listenerMutex.Unlock()
-	c.onConnectedListener = fn
-}
-
 // HasConns returns true if there are connections.
 func (c *Client) HasConns() bool {
 	c.mu.Lock()
@@ -559,16 +551,6 @@ func (c *Client) notifyDisconnected() {
 	go c.onDisconnectListener(c.connectionURL)
 }
 
-func (c *Client) notifyConnected() {
-	c.listenerMutex.Lock()
-	defer c.listenerMutex.Unlock()
-
-	if c.onConnectedListener == nil {
-		return
-	}
-	go c.onConnectedListener()
-}
-
 func (c *Client) writeCloseMsg() {
 	msg := messages.MarshalCloseMsg()
 	_, err := c.relayConn.Write(msg)

relay/client/guard.go

@@ -14,15 +14,17 @@ var (
 
 // Guard manage the reconnection tries to the Relay server in case of disconnection event.
 type Guard struct {
-	// OnNewRelayClient is a channel that is used to notify the relay client about a new relay client instance.
+	// OnNewRelayClient is a channel that is used to notify the relay manager about a new relay client instance.
 	OnNewRelayClient chan *Client
+	OnReconnected    chan struct{}
 	serverPicker     *ServerPicker
 }
 
 // NewGuard creates a new guard for the relay client.
 func NewGuard(sp *ServerPicker) *Guard {
 	g := &Guard{
 		OnNewRelayClient: make(chan *Client, 1),
+		OnReconnected:    make(chan struct{}, 1),
 		serverPicker:     sp,
 	}
 	return g
@@ -39,14 +41,13 @@ func NewGuard(sp *ServerPicker) *Guard {
 // - relayClient: The relay client instance that was disconnected.
 // todo prevent multiple reconnection instances. In the current usage it should not happen, but it is better to prevent
 func (g *Guard) StartReconnectTrys(ctx context.Context, relayClient *Client) {
-	if relayClient == nil {
-		goto RETRY
-	}
-	if g.isServerURLStillValid(relayClient) && g.quickReconnect(ctx, relayClient) {
+	// try to reconnect to the same server
+	if ok := g.tryToQuickReconnect(ctx, relayClient); ok {
+		g.notifyReconnected()
 		return
 	}
 
-RETRY:
+	// start a ticker to pick a new server
 	ticker := exponentTicker(ctx)
 	defer ticker.Stop()
 
@@ -64,28 +65,19 @@ RETRY:
 	}
 }
 
-func (g *Guard) retry(ctx context.Context) error {
-	log.Infof("try to pick up a new Relay server")
-	relayClient, err := g.serverPicker.PickServer(ctx)
-	if err != nil {
-		return err
+func (g *Guard) tryToQuickReconnect(parentCtx context.Context, rc *Client) bool {
+	if rc == nil {
+		return false
 	}
 
-	// prevent to work with a deprecated Relay client instance
-	g.drainRelayClientChan()
-
-	g.OnNewRelayClient <- relayClient
-	return nil
-}
-
-func (g *Guard) quickReconnect(parentCtx context.Context, rc *Client) bool {
-	ctx, cancel := context.WithTimeout(parentCtx, 1500*time.Millisecond)
-	defer cancel()
-	<-ctx.Done()
+	if !g.isServerURLStillValid(rc) {
+		return false
+	}
 
-	if parentCtx.Err() != nil {
+	if cancelled := waiteBeforeRetry(parentCtx); !cancelled {
 		return false
 	}
+
 	log.Infof("try to reconnect to Relay server: %s", rc.connectionURL)
 
 	if err := rc.Connect(); err != nil {
@@ -95,6 +87,20 @@ func (g *Guard) quickReconnect(parentCtx context.Context, rc *Client) bool {
 	return true
 }
 
+func (g *Guard) retry(ctx context.Context) error {
+	log.Infof("try to pick up a new Relay server")
+	relayClient, err := g.serverPicker.PickServer(ctx)
+	if err != nil {
+		return err
+	}
+
+	// prevent to work with a deprecated Relay client instance
+	g.drainRelayClientChan()
+
+	g.OnNewRelayClient <- relayClient
+	return nil
+}
+
 func (g *Guard) drainRelayClientChan() {
 	select {
 	case <-g.OnNewRelayClient:
@@ -111,6 +117,13 @@ func (g *Guard) isServerURLStillValid(rc *Client) bool {
 	return false
 }
 
+func (g *Guard) notifyReconnected() {
+	select {
+	case g.OnReconnected <- struct{}{}:
+	default:
+	}
+}
+
 func exponentTicker(ctx context.Context) *backoff.Ticker {
 	bo := backoff.WithContext(&backoff.ExponentialBackOff{
 		InitialInterval: 2 * time.Second,
@@ -121,3 +134,15 @@ func exponentTicker(ctx context.Context) *backoff.Ticker {
 
 	return backoff.NewTicker(bo)
 }
+
+func waiteBeforeRetry(ctx context.Context) bool {
+	timer := time.NewTimer(1500 * time.Millisecond)
+	defer timer.Stop()
+
+	select {
+	case <-timer.C:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}

relay/client/manager.go

@@ -165,6 +165,9 @@ func (m *Manager) Ready() bool {
 }
 
 func (m *Manager) SetOnReconnectedListener(f func()) {
+	m.listenerLock.Lock()
+	defer m.listenerLock.Unlock()
+
 	m.onReconnectedListenerFn = f
 }
 
@@ -284,6 +287,9 @@ func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
 }
 
 func (m *Manager) onServerConnected() {
+	m.listenerLock.Lock()
+	defer m.listenerLock.Unlock()
+
 	if m.onReconnectedListenerFn == nil {
 		return
 	}
@@ -304,8 +310,11 @@ func (m *Manager) onServerDisconnected(serverAddress string) {
 func (m *Manager) listenGuardEvent(ctx context.Context) {
 	for {
 		select {
+		case <-m.reconnectGuard.OnReconnected:
+			m.onServerConnected()
 		case rc := <-m.reconnectGuard.OnNewRelayClient:
 			m.storeClient(rc)
+			m.onServerConnected()
 		case <-ctx.Done():
 			return
 		}
@@ -317,7 +326,6 @@ func (m *Manager) storeClient(client *Client) {
 	defer m.relayClientMu.Unlock()
 
 	m.relayClient = client
-	m.relayClient.SetOnConnectedListener(m.onServerConnected)
 	m.relayClient.SetOnDisconnectListener(m.onServerDisconnected)
 }