[client] Fix controller re-connection (#2758)

Rethink the peer reconnection implementation

relay/client/client.go

@@ -142,6 +142,7 @@ type Client struct {
 	muInstanceURL    sync.Mutex
 
 	onDisconnectListener func()
+	onConnectedListener  func()
 	listenerMutex        sync.Mutex
 }
 
@@ -191,6 +192,7 @@ func (c *Client) Connect() error {
 
 	c.wgReadLoop.Add(1)
 	go c.readLoop(c.relayConn)
+	go c.notifyConnected()
 
 	return nil
 }
@@ -238,6 +240,12 @@ func (c *Client) SetOnDisconnectListener(fn func()) {
 	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()
@@ -245,6 +253,12 @@ func (c *Client) HasConns() bool {
 	return len(c.conns) > 0
 }
 
+func (c *Client) Ready() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.serviceIsRunning
+}
+
 // Close closes the connection to the relay server and all connections to other peers.
 func (c *Client) Close() error {
 	return c.close(true)
@@ -363,9 +377,9 @@ func (c *Client) readLoop(relayConn net.Conn) {
 	c.instanceURL = nil
 	c.muInstanceURL.Unlock()
 
-	c.notifyDisconnected()
 	c.wgReadLoop.Done()
 	_ = c.close(false)
+	c.notifyDisconnected()
 }
 
 func (c *Client) handleMsg(msgType messages.MsgType, buf []byte, bufPtr *[]byte, hc *healthcheck.Receiver, internallyStoppedFlag *internalStopFlag) (continueLoop bool) {
@@ -544,6 +558,16 @@ func (c *Client) notifyDisconnected() {
 	go c.onDisconnectListener()
 }
 
+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

@@ -29,6 +29,10 @@ func NewGuard(context context.Context, relayClient *Client) *Guard {
 // OnDisconnected is called when the relay client is disconnected from the relay server. It will trigger the reconnection
 // todo prevent multiple reconnection instances. In the current usage it should not happen, but it is better to prevent
 func (g *Guard) OnDisconnected() {
+	if g.quickReconnect() {
+		return
+	}
+
 	ticker := time.NewTicker(reconnectingTimeout)
 	defer ticker.Stop()
 
@@ -46,3 +50,19 @@ func (g *Guard) OnDisconnected() {
 		}
 	}
 }
+
+func (g *Guard) quickReconnect() bool {
+	ctx, cancel := context.WithTimeout(g.ctx, 1500*time.Millisecond)
+	defer cancel()
+	<-ctx.Done()
+
+	if g.ctx.Err() != nil {
+		return false
+	}
+
+	if err := g.relayClient.Connect(); err != nil {
+		log.Errorf("failed to reconnect to relay server: %s", err)
+		return false
+	}
+	return true
+}

relay/client/manager.go

@@ -65,6 +65,7 @@ type Manager struct {
 	relayClientsMutex sync.RWMutex
 
 	onDisconnectedListeners map[string]*list.List
+	onReconnectedListenerFn func()
 	listenerLock            sync.Mutex
 }
 
@@ -101,6 +102,7 @@ func (m *Manager) Serve() error {
 	m.relayClient = client
 
 	m.reconnectGuard = NewGuard(m.ctx, m.relayClient)
+	m.relayClient.SetOnConnectedListener(m.onServerConnected)
 	m.relayClient.SetOnDisconnectListener(func() {
 		m.onServerDisconnected(client.connectionURL)
 	})
@@ -138,6 +140,18 @@ func (m *Manager) OpenConn(serverAddress, peerKey string) (net.Conn, error) {
 	return netConn, err
 }
 
+// Ready returns true if the home Relay client is connected to the relay server.
+func (m *Manager) Ready() bool {
+	if m.relayClient == nil {
+		return false
+	}
+	return m.relayClient.Ready()
+}
+
+func (m *Manager) SetOnReconnectedListener(f func()) {
+	m.onReconnectedListenerFn = f
+}
+
 // AddCloseListener adds a listener to the given server instance address. The listener will be called if the connection
 // closed.
 func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServerCloseListener) error {
@@ -240,6 +254,13 @@ func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
 	return conn, nil
 }
 
+func (m *Manager) onServerConnected() {
+	if m.onReconnectedListenerFn == nil {
+		return
+	}
+	go m.onReconnectedListenerFn()
+}
+
 func (m *Manager) onServerDisconnected(serverAddress string) {
 	if serverAddress == m.relayClient.connectionURL {
 		go m.reconnectGuard.OnDisconnected()