From 1726c45cbf091005805228aa327a440692e6551a Mon Sep 17 00:00:00 2001
From: p4u <pau@dabax.net>
Date: Mon, 5 Oct 2020 10:28:13 +0200
Subject: [PATCH] add a close channel in order to finish the dial loop

before, connect() would keep iterating fater connection close

Signed-off-by: p4u <pau@dabax.net>
---
 recws.go | 68 +++++++++++++++++++++++++++++++-------------------------
 1 file changed, 38 insertions(+), 30 deletions(-)

diff --git a/recws.go b/recws.go
index 1d4b27a..94aec5c 100644
--- a/recws.go
+++ b/recws.go
@@ -54,6 +54,8 @@ type RecConn struct {
 	httpResp    *http.Response
 	dialErr     error
 	dialer      *websocket.Dialer
+	// if set to true, close stops dial reconnection
+	close chan (bool)
 
 	*websocket.Conn
 }
@@ -87,7 +89,7 @@ func (rc *RecConn) Close() {
 		rc.Conn.Close()
 		rc.mu.Unlock()
 	}
-
+	rc.close <- true
 	rc.setIsConnected(false)
 }
 
@@ -290,6 +292,7 @@ func (rc *RecConn) Dial(urlStr string, reqHeader http.Header) {
 		log.Fatalf("Dial: %v", err)
 	}
 
+	rc.close = make(chan bool, 1)
 	// Config
 	rc.setURL(urlStr)
 	rc.setReqHeader(reqHeader)
@@ -394,43 +397,48 @@ func (rc *RecConn) connect() {
 	rand.Seed(time.Now().UTC().UnixNano())
 
 	for {
-		nextItvl := b.Duration()
-		wsConn, httpResp, err := rc.dialer.Dial(rc.url, rc.reqHeader)
-
-		rc.mu.Lock()
-		rc.Conn = wsConn
-		rc.dialErr = err
-		rc.isConnected = err == nil
-		rc.httpResp = httpResp
-		rc.mu.Unlock()
-
-		if err == nil {
-			if !rc.getNonVerbose() {
-				log.Printf("Dial: connection was successfully established with %s\n", rc.url)
-			}
+		select {
+		case <-rc.close:
+			return
+		default:
+			nextItvl := b.Duration()
+			wsConn, httpResp, err := rc.dialer.Dial(rc.url, rc.reqHeader)
+
+			rc.mu.Lock()
+			rc.Conn = wsConn
+			rc.dialErr = err
+			rc.isConnected = err == nil
+			rc.httpResp = httpResp
+			rc.mu.Unlock()
+
+			if err == nil {
+				if !rc.getNonVerbose() {
+					log.Printf("Dial: connection was successfully established with %s\n", rc.url)
+				}
 
-			if rc.hasSubscribeHandler() {
-				if err := rc.SubscribeHandler(); err != nil {
-					log.Fatalf("Dial: connect handler failed with %s", err.Error())
+				if rc.hasSubscribeHandler() {
+					if err := rc.SubscribeHandler(); err != nil {
+						log.Fatalf("Dial: connect handler failed with %s", err.Error())
+					}
+					if !rc.getNonVerbose() {
+						log.Printf("Dial: connect handler was successfully established with %s\n", rc.url)
+					}
 				}
-				if !rc.getNonVerbose() {
-					log.Printf("Dial: connect handler was successfully established with %s\n", rc.url)
+
+				if rc.getKeepAliveTimeout() != 0 {
+					rc.keepAlive()
 				}
+
+				return
 			}
 
-			if rc.getKeepAliveTimeout() != 0 {
-				rc.keepAlive()
+			if !rc.getNonVerbose() {
+				log.Println(err)
+				log.Println("Dial: will try again in", nextItvl, "seconds.")
 			}
-		
-			return
-		}
 
-		if !rc.getNonVerbose() {
-			log.Println(err)
-			log.Println("Dial: will try again in", nextItvl, "seconds.")
+			time.Sleep(nextItvl)
 		}
-
-		time.Sleep(nextItvl)
 	}
 }