From 0aebea60bd1280db05590e10ccd71d852a96b2cf Mon Sep 17 00:00:00 2001
From: "diamondburned (Forefront)" <datutbrus@gmail.com>
Date: Sun, 1 Mar 2020 16:39:40 -0800
Subject: [PATCH] Gateway: Minor bug fixes

---
 gateway/gateway.go    |  2 +-
 gateway/pacemaker.go  | 21 +++++++++++----------
 internal/wsutil/ws.go |  3 +++
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/gateway/gateway.go b/gateway/gateway.go
index bb05e6a..1829cb0 100644
--- a/gateway/gateway.go
+++ b/gateway/gateway.go
@@ -329,7 +329,7 @@ func (g *Gateway) eventLoop() error {
 				return nil
 			}
 
-			return errors.New("Pacemaker died, reconnecting.")
+			return errors.Wrap(err, "Pacemaker died, reconnecting")
 
 		case ev := <-ch:
 			// Check for error
diff --git a/gateway/pacemaker.go b/gateway/pacemaker.go
index bce59b7..ca69be4 100644
--- a/gateway/pacemaker.go
+++ b/gateway/pacemaker.go
@@ -77,21 +77,22 @@ func (p *Pacemaker) start(stop chan struct{}, wg *sync.WaitGroup) error {
 	p.Echo()
 
 	for {
+		if err := p.Pace(); err != nil {
+			return err
+		}
+
+		// Paced, save:
+		atomic.StoreInt64(&p.SentBeat, time.Now().UnixNano())
+
+		if p.Dead() {
+			return ErrDead
+		}
+
 		select {
 		case <-stop:
 			return nil
 
 		case <-tick.C:
-			if err := p.Pace(); err != nil {
-				return err
-			}
-
-			// Paced, save:
-			atomic.StoreInt64(&p.SentBeat, time.Now().UnixNano())
-
-			if p.Dead() {
-				return ErrDead
-			}
 		}
 	}
 }
diff --git a/internal/wsutil/ws.go b/internal/wsutil/ws.go
index 4fb42c9..ab1f3e3 100644
--- a/internal/wsutil/ws.go
+++ b/internal/wsutil/ws.go
@@ -58,6 +58,9 @@ func (ws *Websocket) Dial(ctx context.Context) error {
 		return errors.Wrap(err, "Failed to dial")
 	}
 
+	// Reset the SendLimiter:
+	ws.SendLimiter = NewSendLimiter()
+
 	return nil
 }