diff --git a/session/session.go b/session/session.go index 95607c4..ef9f4a5 100644 --- a/session/session.go +++ b/session/session.go @@ -236,9 +236,8 @@ func (s *Session) Open(ctx context.Context) error { s.state.gateway = g } - ctx, cancel := context.WithCancel(context.Background()) - s.state.ctx = ctx - s.state.cancel = cancel + // Make a context that's stored in state so this can be used throughout. + s.state.ctx, s.state.cancel = context.WithCancel(context.Background()) // TODO: change this to AddSyncHandler. rm := s.AddHandler(evCh) @@ -253,6 +252,10 @@ func (s *Session) Open(ctx context.Context) error { s.close() return ctx.Err() + case <-s.state.ctx.Done(): + s.close() + return s.state.ctx.Err() + case <-s.state.doneCh: // Event loop died. return s.state.gateway.LastError()