wsutil: Reset SendLimiter after Dial

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(),
+		Timeout: WSTimeout,
-		DialLimiter: NewDialLimiter(),
 	}
 }
 
@@ -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
 }
 
@@ -123,13 +126,20 @@ func (ws *Websocket) Send(b []byte) error {
 // SendCtx sends b over the Websocket with a deadline. It closes the internal
 // Websocket if the Send method errors out.
 func (ws *Websocket) SendCtx(ctx context.Context, b []byte) error {
-	if err := ws.SendLimiter.Wait(ctx); err != nil {
+	WSDebug("Waiting for the send rate limiter...")
+
+	if err := ws.sendLimiter.Wait(ctx); err != nil {
+		WSDebug("Send rate limiter timed out.")
 		return errors.Wrap(err, "SendLimiter failed")
 	}
 
+	WSDebug("Send is passed the rate limiting. Waiting on mutex.")
+
 	ws.mutex.Lock()
 	defer ws.mutex.Unlock()
 
+	WSDebug("Mutex lock acquired.")
+
 	if ws.closed {
 		return ErrWebsocketClosed
 	}