[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)