Stop broadcaster subscribe/unsubscribe from taking the stop signal

pkg/rsnotify/broadcaster/broadcaster.go

@@ -26,6 +26,10 @@ type NotificationBroadcaster struct {
 	subscribe   chan Subscription
 	unsubscribe chan (<-chan listener.Notification)
 	stopSignal  chan bool
+
+	// Closed when broadcast() exits. Callers wait on this rather than
+	// stopSignal so only broadcast() ever receives the stop value.
+	done chan struct{}
 }
 
 func NewNotificationBroadcaster(l listener.Listener, stop chan bool) (*NotificationBroadcaster, error) {
@@ -34,6 +38,7 @@ func NewNotificationBroadcaster(l listener.Listener, stop chan bool) (*Notificat
 		subscribe:   make(chan Subscription),
 		unsubscribe: make(chan (<-chan listener.Notification)),
 		stopSignal:  stop,
+		done:        make(chan struct{}),
 	}
 	var err error
 	b.msgs, b.errs, err = b.listener.Listen()
@@ -55,6 +60,7 @@ func (b *NotificationBroadcaster) IP() string {
 func (b *NotificationBroadcaster) broadcast() {
 	sinks := make([]Subscription, 0)
 	defer close(b.stopSignal)
+	defer close(b.done)
 	for {
 		select {
 		case <-b.stopSignal:
@@ -152,7 +158,7 @@ func (b *NotificationBroadcaster) Subscribe(dataType uint8) <-chan listener.Noti
 		T: dataType,
 	}:
 		return c
-	case <-b.stopSignal:
+	case <-b.done:
 		// Broadcaster has stopped; return nil to indicate no subscription.
 		return nil
 	}
@@ -174,7 +180,7 @@ func (b *NotificationBroadcaster) SubscribeOne(dataType uint8, matcher Matcher)
 		One: matcher,
 	}:
 		return c
-	case <-b.stopSignal:
+	case <-b.done:
 		// Broadcaster has stopped; return nil to indicate no subscription.
 		return nil
 	}
@@ -204,7 +210,7 @@ func (b *NotificationBroadcaster) Unsubscribe(ch <-chan listener.Notification) {
 	go drainer()
 	select {
 	case b.unsubscribe <- ch:
-	case <-b.stopSignal:
+	case <-b.done:
 		// Broadcaster has stopped; all channels were closed by stop().
 		// The drainer will exit when it sees the closed channel.
 	}

pkg/rsnotify/broadcaster/broadcaster_test.go

@@ -317,3 +317,46 @@ func (s *BroadcasterSuite) TestSubscribeAfterStop(c *check.C) {
 	// Unsubscribe(nil) should not block or leak goroutines
 	b.Unsubscribe(ch2)
 }
+
+// TestStopDuringSubscribe verifies that a stop sent while Subscribe and
+// Unsubscribe calls are in flight still reaches the broadcast loop, which
+// acks it by closing the stop channel.
+func (s *BroadcasterSuite) TestStopDuringSubscribe(c *check.C) {
+	for trial := 0; trial < 10; trial++ {
+		l := &FakeListener{
+			items: make(chan listener.Notification),
+			errs:  make(chan error),
+		}
+		stop := make(chan bool)
+		b, err := NewNotificationBroadcaster(l, stop)
+		c.Assert(err, check.IsNil)
+
+		quit := make(chan struct{})
+		wg := &sync.WaitGroup{}
+		for i := 0; i < 50; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				for {
+					select {
+					case <-quit:
+						return
+					default:
+					}
+					b.Unsubscribe(b.Subscribe(1))
+					b.Unsubscribe(b.SubscribeOne(1, func(n listener.Notification) bool { return true }))
+				}
+			}()
+		}
+
+		stop <- true
+		select {
+		case <-stop:
+			// Success - the broadcast loop received the stop and closed the channel.
+		case <-time.After(time.Second):
+			c.Fatalf("trial %d: broadcaster did not ack stop; a subscriber took the stop signal", trial)
+		}
+		close(quit)
+		wg.Wait()
+	}
+}