add PubSub's context to Subscription

This commit is contained in:
lukesolo 2019-09-30 12:55:37 +03:00
parent 9f04364996
commit 9d03237eec
2 changed files with 6 additions and 1 deletions

View File

@ -696,6 +696,7 @@ func (p *PubSub) SubscribeByTopicDescriptor(td *pb.TopicDescriptor, opts ...SubO
sub := &Subscription{
topic: td.GetName(),
ctx: p.ctx,
ch: make(chan *Message, 32),
peerEvtCh: make(chan PeerEvent, 1),

View File

@ -18,6 +18,7 @@ type Subscription struct {
ch chan *Message
cancelCh chan<- *Subscription
err error
ctx context.Context
peerEvtCh chan PeerEvent
evtLogMx sync.Mutex
@ -49,7 +50,10 @@ func (sub *Subscription) Next(ctx context.Context) (*Message, error) {
}
func (sub *Subscription) Cancel() {
sub.cancelCh <- sub
select {
case sub.cancelCh <- sub:
case <-sub.ctx.Done():
}
}
func (sub *Subscription) close() {