diff --git a/session/session.go b/session/session.go index 5c3052e..815fcd8 100644 --- a/session/session.go +++ b/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 = stop - go s.startHandler(stop) + s.hstop = make(chan struct{}) + s.wstop = sync.Once{} + 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) - } -}