Reset lastReceivedTime before connecting websocket to resolve race co…

…ndition
HuobiRDCenter · Sep 11, 2020 · 73530f6 · 73530f6
1 parent dba6b5f
commit 73530f6
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/pkg/client/websocketclientbase/websocketv1clientbase.go b/pkg/client/websocketclientbase/websocketv1clientbase.go
@@ -58,11 +58,20 @@ func (p *WebSocketV1ClientBase) SetHandler(authHandler AuthenticationV1ResponseH
 // Connect to websocket server
 // if autoConnect is true, then the connection can be re-connect if no data received after the pre-defined timeout
 func (p *WebSocketV1ClientBase) Connect(autoConnect bool) error {
+
+	// reset last received time as now
+	p.lastReceivedTime = time.Now()
+
+	// connect to websocket
 	err := p.connectWebSocket()
 	if err != nil {
 		return err
 	}
 
+	// start loop to read and handle message
+	p.startReadLoop()
+
+	// start ticker to manage connection
 	if autoConnect {
 		p.startTicker()
 	}
@@ -108,8 +117,6 @@ func (p *WebSocketV1ClientBase) connectWebSocket() error {
 		return err
 	}
 
-	p.startReadLoop()
-
 	return nil
 }
 
@@ -134,7 +141,6 @@ func (p *WebSocketV1ClientBase) disconnectWebSocket() {
 // initialize a ticker and start a goroutine tickerLoop()
 func (p *WebSocketV1ClientBase) startTicker() {
 	p.ticker = time.NewTicker(TimerIntervalSecond * time.Second)
-	p.lastReceivedTime = time.Now()
 
 	go p.tickerLoop()
 }

diff --git a/pkg/client/websocketclientbase/websocketv2clientbase.go b/pkg/client/websocketclientbase/websocketv2clientbase.go
@@ -58,8 +58,16 @@ func (p *WebSocketV2ClientBase) SetHandler(authHandler AuthenticationV2ResponseH
 // Connect to websocket server
 // if autoConnect is true, then the connection can be re-connect if no data received after the pre-defined timeout
 func (p *WebSocketV2ClientBase) Connect(autoConnect bool) {
+	// reset last received time as now
+	p.lastReceivedTime = time.Now()
+
+	// connect to websocket
 	p.connectWebSocket()
 
+	// start loop to read and handle message
+	p.startReadLoop()
+
+	// start ticker to manage connection
 	if autoConnect {
 		p.startTicker()
 	}
@@ -106,8 +114,6 @@ func (p *WebSocketV2ClientBase) connectWebSocket() {
 	}
 
 	p.Send(auth)
-
-	p.startReadLoop()
 }
 
 // disconnect with server
@@ -132,7 +138,6 @@ func (p *WebSocketV2ClientBase) disconnectWebSocket() {
 // initialize a ticker and start a goroutine tickerLoop()
 func (p *WebSocketV2ClientBase) startTicker() {
 	p.ticker = time.NewTicker(TimerIntervalSecond * time.Second)
-	p.lastReceivedTime = time.Now()
 
 	go p.tickerLoop()
 }