From a25250d95970cd55898743d32691b7c31a989507 Mon Sep 17 00:00:00 2001
From: diamondburned <datutbrus@gmail.com>
Date: Sun, 1 Nov 2020 10:09:41 -0800
Subject: [PATCH] wsutil: Reset SendLimiter after Dial

---
 utils/wsutil/conn.go |  8 ++++----
 utils/wsutil/ws.go   | 19 +++++++++++--------
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/utils/wsutil/conn.go b/utils/wsutil/conn.go
index 6c216f8..b698915 100644
--- a/utils/wsutil/conn.go
+++ b/utils/wsutil/conn.go
@@ -182,10 +182,10 @@ func startReadLoop(conn *websocket.Conn, eventCh chan<- Event) {
 				return
 			}
 
-			// // Check if the error is a normal one:
-			// if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
-			// 	return
-			// }
+			// Check if the error is a normal one:
+			if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
+				return
+			}
 
 			// Unusual error; log and exit:
 			eventCh <- Event{nil, errors.Wrap(err, "WS error")}
diff --git a/utils/wsutil/ws.go b/utils/wsutil/ws.go
index 0115382..da995d0 100644
--- a/utils/wsutil/ws.go
+++ b/utils/wsutil/ws.go
@@ -41,15 +41,15 @@ type Websocket struct {
 	addr   string
 	closed bool
 
+	sendLimiter *rate.Limiter
+	dialLimiter *rate.Limiter
+
 	// Constants. These must not be changed after the Websocket instance is used
 	// once, as they are not thread-safe.
 
 	// Timeout for connecting and writing to the Websocket, uses default
 	// WSTimeout (global).
 	Timeout time.Duration
-
-	SendLimiter *rate.Limiter
-	DialLimiter *rate.Limiter
 }
 
 // New creates a default Websocket with the given address.
@@ -64,10 +64,10 @@ func NewCustom(conn Connection, addr string) *Websocket {
 		addr:   addr,
 		closed: true,
 
-		Timeout: WSTimeout,
+		sendLimiter: NewSendLimiter(),
+		dialLimiter: NewDialLimiter(),
 
-		SendLimiter: NewSendLimiter(),
-		DialLimiter: NewDialLimiter(),
+		Timeout: WSTimeout,
 	}
 }
 
@@ -80,7 +80,7 @@ func (ws *Websocket) Dial(ctx context.Context) error {
 		ctx = tctx
 	}
 
-	if err := ws.DialLimiter.Wait(ctx); err != nil {
+	if err := ws.dialLimiter.Wait(ctx); err != nil {
 		// Expired, fatal error
 		return errors.Wrap(err, "failed to wait")
 	}
@@ -99,6 +99,9 @@ func (ws *Websocket) Dial(ctx context.Context) error {
 
 	ws.closed = false
 
+	// Reset the send limiter.
+	ws.sendLimiter = NewSendLimiter()
+
 	return nil
 }
 
@@ -125,7 +128,7 @@ func (ws *Websocket) Send(b []byte) error {
 func (ws *Websocket) SendCtx(ctx context.Context, b []byte) error {
 	WSDebug("Waiting for the send rate limiter...")
 
-	if err := ws.SendLimiter.Wait(ctx); err != nil {
+	if err := ws.sendLimiter.Wait(ctx); err != nil {
 		WSDebug("Send rate limiter timed out.")
 		return errors.Wrap(err, "SendLimiter failed")
 	}