Session: Fixed a potential race condition on Close

This commit is contained in:
diamondburned 2020-10-17 03:21:07 -07:00
parent 397d288927
commit 86795e42a6
1 changed files with 9 additions and 13 deletions

View File

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