fix: potential deadlock in Subscriber (#584)

* fix: potential deadlock in Subscriber

* another try

* new try

* handle more cases

bolt_transport_test.go

@@ -227,7 +227,7 @@ func TestBoltTransportClosed(t *testing.T) {
 
 	assert.Equal(t, transport.Dispatch(&Update{Topics: s.Topics}), ErrClosedTransport)
 
-	_, ok := <-s.disconnected
+	_, ok := <-s.out
 	assert.False(t, ok)
 }
 

local_transport_bench_test.go

@@ -39,7 +39,6 @@ func subBenchLocalTransport(b *testing.B, topics, concurrency, matchPct int, tes
 		}
 	}
 	out := make(chan *Update, 50000)
-	once := &sync.Once{}
 	for i := 0; i < concurrency; i++ {
 		s := NewSubscriber("", zap.NewNop())
 		if i%100 < matchPct {
@@ -48,7 +47,6 @@ func subBenchLocalTransport(b *testing.B, topics, concurrency, matchPct int, tes
 			s.SetTopics(tsNoMatch, nil)
 		}
 		s.out = out
-		s.disconnectedOnce = once
 		tr.AddSubscriber(s)
 	}
 	ctx, done := context.WithCancel(context.Background())

subscribe_test.go

@@ -272,7 +272,7 @@ func TestUnsubscribe(t *testing.T) {
 		hub.SubscribeHandler(httptest.NewRecorder(), req)
 		assert.Equal(t, 0, s.subscribers.Len())
 		s.subscribers.Walk(0, func(s *Subscriber) bool {
-			_, ok := <-s.disconnected
+			_, ok := <-s.out
 			assert.False(t, ok)
 
 			return true

subscriber.go

@@ -26,9 +26,9 @@ type Subscriber struct {
 	PrivateRegexps     []*regexp.Regexp
 	Debug              bool
 
-	disconnectedOnce    *sync.Once
+	disconnected        int32
 	out                 chan *Update
-	disconnected        chan struct{}
+	outMutex            sync.RWMutex
 	responseLastEventID chan string
 	logger              Logger
 	ready               int32
@@ -45,20 +45,16 @@ func NewSubscriber(lastEventID string, logger Logger) *Subscriber {
 		RequestLastEventID:  lastEventID,
 		responseLastEventID: make(chan string, 1),
 		out:                 make(chan *Update, 1000),
-		disconnected:        make(chan struct{}),
 		logger:              logger,
-		disconnectedOnce:    &sync.Once{},
 	}
 
 	return s
 }
 
 // Dispatch an update to the subscriber.
 func (s *Subscriber) Dispatch(u *Update, fromHistory bool) bool {
-	select {
-	case <-s.disconnected:
+	if atomic.LoadInt32(&s.disconnected) > 0 {
 		return false
-	default:
 	}
 
 	if !fromHistory && atomic.LoadInt32(&s.ready) < 1 {
@@ -71,6 +67,13 @@ func (s *Subscriber) Dispatch(u *Update, fromHistory bool) bool {
 		}
 		s.liveMutex.Unlock()
 	}
+
+	s.outMutex.Lock()
+	defer s.outMutex.Unlock()
+	if atomic.LoadInt32(&s.disconnected) > 0 {
+		return false
+	}
+
 	s.out <- u
 
 	return true
@@ -80,6 +83,8 @@ func (s *Subscriber) Dispatch(u *Update, fromHistory bool) bool {
 func (s *Subscriber) Ready() int {
 	s.liveMutex.Lock()
 	defer s.liveMutex.Unlock()
+	s.outMutex.Lock()
+	defer s.outMutex.Unlock()
 
 	n := len(s.liveQueue)
 	for _, u := range s.liveQueue {
@@ -102,10 +107,15 @@ func (s *Subscriber) HistoryDispatched(responseLastEventID string) {
 
 // Disconnect disconnects the subscriber.
 func (s *Subscriber) Disconnect() {
-	s.disconnectedOnce.Do(func() {
-		close(s.disconnected)
-		close(s.out)
-	})
+	if atomic.LoadInt32(&s.disconnected) > 0 {
+		return
+	}
+
+	s.outMutex.Lock()
+	defer s.outMutex.Unlock()
+
+	atomic.StoreInt32(&s.disconnected, 1)
+	close(s.out)
 }
 
 // SetTopics compiles topic selector regexps.

transport_test.go

@@ -62,7 +62,7 @@ func TestLocalTransportClosed(t *testing.T) {
 	assert.Equal(t, transport.AddSubscriber(NewSubscriber("", zap.NewNop())), ErrClosedTransport)
 	assert.Equal(t, transport.Dispatch(&Update{}), ErrClosedTransport)
 
-	_, ok := <-s.disconnected
+	_, ok := <-s.out
 	assert.False(t, ok)
 }