pubsub: fix race on shutdown

Calling `wg.Add` after `wg.Wait` has returned is invalid. This change swaps the
wait group for a plain rwmutex.

(caught with the race detector)

notifications/notifications.go

@@ -20,29 +20,21 @@ type PubSub interface {
 func New() PubSub {
 	return &impl{
 		wrapped: *pubsub.New(bufferSize),
-		cancel:  make(chan struct{}),
 	}
 }
 
 type impl struct {
+	lk      sync.RWMutex
 	wrapped pubsub.PubSub
 
-	// These two fields make up a shutdown "lock".
-	// We need them as calling, e.g., `Unsubscribe` after calling `Shutdown`
-	// blocks forever and fixing this in pubsub would be rather invasive.
-	cancel chan struct{}
-	wg     sync.WaitGroup
+	closed bool
 }
 
 func (ps *impl) Publish(block blocks.Block) {
-	ps.wg.Add(1)
-	defer ps.wg.Done()
-
-	select {
-	case <-ps.cancel:
-		// Already shutdown, bail.
+	ps.lk.RLock()
+	defer ps.lk.RUnlock()
+	if ps.closed {
 		return
-	default:
 	}
 
 	ps.wrapped.Pub(block, block.Cid().KeyString())
@@ -50,12 +42,13 @@ func (ps *impl) Publish(block blocks.Block) {
 
 // Not safe to call more than once.
 func (ps *impl) Shutdown() {
-	// Interrupt in-progress subscriptions.
-	close(ps.cancel)
-	// Wait for them to finish.
-	ps.wg.Wait()
-	// shutdown the pubsub.
+	ps.lk.Lock()
+	defer ps.lk.Unlock()
+	if ps.closed {
+		return
+	}
 	ps.wrapped.Shutdown()
+	ps.closed = true
 }
 
 // Subscribe returns a channel of blocks for the given |keys|. |blockChannel|
@@ -71,32 +64,32 @@ func (ps *impl) Subscribe(ctx context.Context, keys ...cid.Cid) <-chan blocks.Bl
 	}
 
 	// prevent shutdown
-	ps.wg.Add(1)
+	ps.lk.RLock()
+	defer ps.lk.RUnlock()
 
-	// check if shutdown *after* preventing shutdowns.
-	select {
-	case <-ps.cancel:
-		// abort, allow shutdown to continue.
-		ps.wg.Done()
+	if ps.closed {
 		close(blocksCh)
 		return blocksCh
-	default:
 	}
 
 	ps.wrapped.AddSubOnceEach(valuesCh, toStrings(keys)...)
 	go func() {
 		defer func() {
-			ps.wrapped.Unsub(valuesCh)
 			close(blocksCh)
 
-			// Unblock shutdown.
-			ps.wg.Done()
+			ps.lk.RLock()
+			defer ps.lk.RUnlock()
+			if ps.closed {
+				// Don't touch the pubsub instance if we're
+				// already closed.
+				return
+			}
+
+			ps.wrapped.Unsub(valuesCh)
 		}()
 
 		for {
 			select {
-			case <-ps.cancel:
-				return
 			case <-ctx.Done():
 				return
 			case val, ok := <-valuesCh:
@@ -107,9 +100,10 @@ func (ps *impl) Subscribe(ctx context.Context, keys ...cid.Cid) <-chan blocks.Bl
 				if !ok {
 					return
 				}
+				// We could end up blocking here if the client
+				// forgets to cancel the context but that's not
+				// our problem.
 				select {
-				case <-ps.cancel:
-					return
 				case <-ctx.Done():
 					return
 				case blocksCh <- block: // continue

notifications/notifications_test.go

@@ -114,7 +114,7 @@ func TestShutdownBeforeUnsubscribe(t *testing.T) {
 		if ok {
 			t.Fatal("channel should have been closed")
 		}
-	default:
+	case <-time.After(5 * time.Second):
 		t.Fatal("channel should have been closed")
 	}
 }