Gateway: Fixed a race condition on ReconnectOP

gateway/op.go

@@ -30,6 +30,11 @@ const (
 	GuildSubscriptionsOP  OPCode = 14
 )
 
+// ErrReconnectRequest is returned by HandleOP if a ReconnectOP is given. This
+// is used mostly internally to signal the heartbeat loop to reconnect, if
+// needed. It is not a fatal error.
+var ErrReconnectRequest = errors.New("ReconnectOP received")
+
 func (g *Gateway) HandleOP(op *wsutil.OP) error {
 	switch op.Code {
 	case HeartbeatAckOP:
@@ -47,13 +52,9 @@ func (g *Gateway) HandleOP(op *wsutil.OP) error {
 		// Server requests to reconnect, die and retry.
 		wsutil.WSDebug("ReconnectOP received.")
 
-		// We must reconnect in another goroutine, as running Reconnect
+		// Exit with the ReconnectOP error to force the heartbeat event loop to
-		// synchronously would prevent the main event loop from exiting.
+		// reconnect synchronously. Not really a fatal error.
-		go g.Reconnect()
+		return ErrReconnectRequest
-
-		// Gracefully exit with a nil let the event handler take the signal from
-		// the pacemaker.
-		return nil
 
 	case InvalidSessionOP:
 		// Discord expects us to sleep for no reason
@@ -101,7 +102,7 @@ func (g *Gateway) HandleOP(op *wsutil.OP) error {
 			g.SessionID = ev.SessionID
 		}
 
-		// Throw the event into a channel, it's valid now.
+		// Throw the event into a channel; it's valid now.
 		g.Events <- ev
 		return nil
 

session/session.go

@@ -4,6 +4,8 @@
 package session
 
 import (
+	"sync"
+
 	"github.com/pkg/errors"
 
 	"github.com/diamondburned/arikawa/api"
@@ -39,6 +41,7 @@ type Session struct {
 	Ticket string
 
 	hstop chan struct{}
+	wstop sync.Once
 }
 
 func NewWithIntents(token string, intents ...gateway.Intents) (*Session, error) {
@@ -103,9 +106,9 @@ func NewWithGateway(gw *gateway.Gateway) *Session {
 
 func (s *Session) Open() error {
 	// Start the handler beforehand so no events are missed.
-	stop := make(chan struct{})
+	s.hstop = make(chan struct{})
-	s.hstop = stop
+	s.wstop = sync.Once{}
-	go s.startHandler(stop)
+	go s.startHandler()
 
 	// Set the AfterClose's handler.
 	s.Gateway.AfterClose = func(err error) {
@@ -121,10 +124,10 @@ func (s *Session) Open() error {
 	return nil
 }
 
-func (s *Session) startHandler(stop <-chan struct{}) {
+func (s *Session) startHandler() {
 	for {
 		select {
-		case <-stop:
+		case <-s.hstop:
 			return
 		case ev := <-s.Gateway.Events:
 			s.Call(ev)
@@ -134,14 +137,7 @@ func (s *Session) startHandler(stop <-chan struct{}) {
 
 func (s *Session) Close() error {
 	// Stop the event handler
-	s.close()
+	s.wstop.Do(func() { s.hstop <- struct{}{} })
-
 	// Close the websocket
 	return s.Gateway.Close()
 }
-
-func (s *Session) close() {
-	if s.hstop != nil {
-		close(s.hstop)
-	}
-}

utils/wsutil/heart.go

@@ -108,7 +108,7 @@ func (p *PacemakerLoop) startLoop() error {
 
 			// Handle the event
 			if err := p.handler(o); err != nil {
-				p.errorLog(errors.Wrap(err, "handler failed"))
+				return errors.Wrap(err, "handler failed")
 			}
 		}
 	}