{Voice,}Gateway: Fixed various race conditions

This commit fixes race conditions in both package voice, package
voicegateway and package gateway.

Originally, several race conditions exist when both the user's and the
pacemaker's goroutines both want to do several things to the websocket
connection. For example, the user's goroutine could be writing, and the
pacemaker's goroutine could trigger a reconnection. This is racey.

This issue is partially fixed by removing the pacer loop from package
heart and combining the ticker into the event (pacemaker) loop itself.

Technically, a race condition could still be triggered with care, but
the API itself never guaranteed any of those. As events are handled
using an internal loop into a channel, a race condition will not be
triggered just by handling events and writing to the websocket.

Author: diamondburned

gateway/gateway.go

@@ -161,40 +161,38 @@ func (g *Gateway) AddIntent(i Intents) {
 }
 
 // Close closes the underlying Websocket connection.
-func (g *Gateway) Close() error {
+func (g *Gateway) Close() (err error) {
 	wsutil.WSDebug("Trying to close.")
 
 	// Check if the WS is already closed:
-	if g.waitGroup == nil && g.PacerLoop.Stopped() {
+	if g.PacerLoop.Stopped() {
 		wsutil.WSDebug("Gateway is already closed.")
-
-		g.AfterClose(nil)
-		return nil
+		return err
 	}
 
+	// Trigger the close callback on exit.
+	defer func() { g.AfterClose(err) }()
+
 	// If the pacemaker is running:
 	if !g.PacerLoop.Stopped() {
 		wsutil.WSDebug("Stopping pacemaker...")
 
-		// Stop the pacemaker and the event handler
+		// Stop the pacemaker and the event handler.
 		g.PacerLoop.Stop()
 
 		wsutil.WSDebug("Stopped pacemaker.")
 	}
 
+	wsutil.WSDebug("Closing the websocket...")
+	err = g.WS.Close()
+
 	wsutil.WSDebug("Waiting for WaitGroup to be done.")
 
 	// This should work, since Pacemaker should signal its loop to stop, which
 	// would also exit our event loop. Both would be 2.
 	g.waitGroup.Wait()
 
-	// Mark g.waitGroup as empty:
-	g.waitGroup = nil
-
 	wsutil.WSDebug("WaitGroup is done. Closing the websocket.")
-
-	err := g.WS.Close()
-	g.AfterClose(err)
 	return err
 }
 

gateway/op.go

@@ -46,15 +46,17 @@ func (g *Gateway) HandleOP(op *wsutil.OP) error {
 		defer cancel()
 
 		// Server requesting a heartbeat.
-		return g.PacerLoop.Pace(ctx)
+		if err := g.PacerLoop.Pace(ctx); err != nil {
+			return wsutil.ErrBrokenConnection(errors.Wrap(err, "failed to pace"))
+		}
 
 	case ReconnectOP:
 		// Server requests to reconnect, die and retry.
 		wsutil.WSDebug("ReconnectOP received.")
 
 		// Exit with the ReconnectOP error to force the heartbeat event loop to
 		// reconnect synchronously. Not really a fatal error.
-		return ErrReconnectRequest
+		return wsutil.ErrBrokenConnection(ErrReconnectRequest)
 
 	case InvalidSessionOP:
 		// Discord expects us to sleep for no reason
@@ -66,7 +68,7 @@ func (g *Gateway) HandleOP(op *wsutil.OP) error {
 		// Invalid session, try and Identify.
 		if err := g.IdentifyCtx(ctx); err != nil {
 			// Can't identify, reconnect.
-			go g.Reconnect()
+			return wsutil.ErrBrokenConnection(ErrReconnectRequest)
 		}
 
 		return nil

go.sum

@@ -16,3 +16,4 @@ golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1 h1:NusfzzA6yGQ+ua51ck7E3omNUX/JuqbFSaRGqU8CcLI=
 golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e h1:EHBhcS0mlXEAVwNyO2dLfjToGsyY4j24pTs2ScHnX7s=

internal/heart/heart.go

@@ -3,7 +3,6 @@ package heart
 
 import (
 	"context"
-	"sync"
 	"sync/atomic"
 	"time"
 
@@ -36,23 +35,30 @@ type Pacemaker struct {
 	// Heartrate is the received duration between heartbeats.
 	Heartrate time.Duration
 
+	ticker time.Ticker
+	Ticks  <-chan time.Time
+
 	// Time in nanoseconds, guarded by atomic read/writes.
 	SentBeat AtomicTime
 	EchoBeat AtomicTime
 
 	// Any callback that returns an error will stop the pacer.
-	Pace func(context.Context) error
-
-	stop  chan struct{}
-	once  sync.Once
-	death chan error
+	Pacer func(context.Context) error
 }
 
-func NewPacemaker(heartrate time.Duration, pacer func(context.Context) error) *Pacemaker {
-	return &Pacemaker{
+func NewPacemaker(heartrate time.Duration, pacer func(context.Context) error) Pacemaker {
+	p := Pacemaker{
 		Heartrate: heartrate,
-		Pace:      pacer,
+		Pacer:     pacer,
+		ticker:    *time.NewTicker(heartrate),
 	}
+	p.Ticks = p.ticker.C
+	// Reset states to its old position.
+	now := time.Now()
+	p.EchoBeat.Set(now)
+	p.SentBeat.Set(now)
+
+	return p
 }
 
 func (p *Pacemaker) Echo() {
@@ -62,14 +68,6 @@ func (p *Pacemaker) Echo() {
 
 // Dead, if true, will have Pace return an ErrDead.
 func (p *Pacemaker) Dead() bool {
-	/* Deprecated
-	if p.LastBeat[0].IsZero() || p.LastBeat[1].IsZero() {
-		return false
-	}
-
-	return p.LastBeat[0].Sub(p.LastBeat[1]) > p.Heartrate*2
-	*/
-
 	var (
 		echo = p.EchoBeat.Get()
 		sent = p.SentBeat.Get()
@@ -84,75 +82,84 @@ func (p *Pacemaker) Dead() bool {
 
 // Stop stops the pacemaker, or it does nothing if the pacemaker is not started.
 func (p *Pacemaker) Stop() {
-	Debug("(*Pacemaker).stop is trying sync.Once.")
-
-	p.once.Do(func() {
-		Debug("(*Pacemaker).stop closed the channel.")
-		close(p.stop)
-	})
+	p.ticker.Stop()
 }
 
 // pace sends a heartbeat with the appropriate timeout for the context.
-func (p *Pacemaker) pace() error {
+func (p *Pacemaker) Pace() error {
 	ctx, cancel := context.WithTimeout(context.Background(), p.Heartrate)
 	defer cancel()
 
-	return p.Pace(ctx)
+	return p.PaceCtx(ctx)
 }
 
-func (p *Pacemaker) start() error {
-	// Reset states to its old position.
-	p.EchoBeat.Set(time.Time{})
-	p.SentBeat.Set(time.Time{})
-
-	// Create a new ticker.
-	tick := time.NewTicker(p.Heartrate)
-	defer tick.Stop()
-
-	// Echo at least once
-	p.Echo()
-
-	for {
-		if err := p.pace(); err != nil {
-			return errors.Wrap(err, "failed to pace")
-		}
-
-		// Paced, save:
-		p.SentBeat.Set(time.Now())
-
-		if p.Dead() {
-			return ErrDead
-		}
-
-		select {
-		case <-p.stop:
-			return nil
-
-		case <-tick.C:
-		}
+func (p *Pacemaker) PaceCtx(ctx context.Context) error {
+	if err := p.Pacer(ctx); err != nil {
+		return err
 	}
-}
 
-// StartAsync starts the pacemaker asynchronously. The WaitGroup is optional.
-func (p *Pacemaker) StartAsync(wg *sync.WaitGroup) (death chan error) {
-	p.death = make(chan error)
-	p.stop = make(chan struct{})
-	p.once = sync.Once{}
+	p.SentBeat.Set(time.Now())
 
-	if wg != nil {
-		wg.Add(1)
+	if p.Dead() {
+		return ErrDead
 	}
 
-	go func() {
-		p.death <- p.start()
-		// Debug.
-		Debug("Pacemaker returned.")
-
-		// Mark the pacemaker loop as done.
-		if wg != nil {
-			wg.Done()
-		}
-	}()
-
-	return p.death
+	return nil
 }
+
+// func (p *Pacemaker) start() error {
+// 	// Reset states to its old position.
+// 	p.EchoBeat.Set(time.Time{})
+// 	p.SentBeat.Set(time.Time{})
+
+// 	// Create a new ticker.
+// 	tick := time.NewTicker(p.Heartrate)
+// 	defer tick.Stop()
+
+// 	// Echo at least once
+// 	p.Echo()
+
+// 	for {
+// 		if err := p.pace(); err != nil {
+// 			return errors.Wrap(err, "failed to pace")
+// 		}
+
+// 		// Paced, save:
+// 		p.SentBeat.Set(time.Now())
+
+// 		if p.Dead() {
+// 			return ErrDead
+// 		}
+
+// 		select {
+// 		case <-p.stop:
+// 			return nil
+
+// 		case <-tick.C:
+// 		}
+// 	}
+// }
+
+// // StartAsync starts the pacemaker asynchronously. The WaitGroup is optional.
+// func (p *Pacemaker) StartAsync(wg *sync.WaitGroup) (death chan error) {
+// 	p.death = make(chan error)
+// 	p.stop = make(chan struct{})
+// 	p.once = sync.Once{}
+
+// 	if wg != nil {
+// 		wg.Add(1)
+// 	}
+
+// 	go func() {
+// 		p.death <- p.start()
+// 		// Debug.
+// 		Debug("Pacemaker returned.")
+
+// 		// Mark the pacemaker loop as done.
+// 		if wg != nil {
+// 			wg.Done()
+// 		}
+// 	}()
+
+// 	return p.death
+// }