From f5ec3e00d55d7cf06375b3eb8f69efb1d33710f2 Mon Sep 17 00:00:00 2001
From: Armin <me@obrown.io>
Date: Mon, 18 Dec 2017 19:46:04 +0100
Subject: [PATCH] rpc: Closing Notifier must tear down subscriptions

* Fixes #15553
* Add a pipeline on the `Notifier`'s `codec` that propagates
the `close` to `Subscription`s
* Adjust `subscription_test` accordingly since closing the notifier
now implies closing the subscription
  * This should automatically test the initial issue and also
makes the tests more determistic by saving the `select`
---
 rpc/subscription.go      | 22 ++++++++++++++++++++--
 rpc/subscription_test.go | 14 ++++----------
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/rpc/subscription.go b/rpc/subscription.go
index 6ce7befa1d..463f36acac 100644
--- a/rpc/subscription.go
+++ b/rpc/subscription.go
@@ -55,16 +55,34 @@ type Notifier struct {
 	subMu    sync.RWMutex // guards active and inactive maps
 	active   map[ID]*Subscription
 	inactive map[ID]*Subscription
+	closed   chan interface{}
 }
 
 // newNotifier creates a new notifier that can be used to send subscription
 // notifications to the client.
 func newNotifier(codec ServerCodec) *Notifier {
-	return &Notifier{
+	n := &Notifier{
 		codec:    codec,
 		active:   make(map[ID]*Subscription),
 		inactive: make(map[ID]*Subscription),
+		closed:   make(chan interface{}),
 	}
+	in := n.codec.Closed()
+	go func() {
+		for e := range in {
+			n.closed <- e
+		}
+		n.subMu.Lock()
+		defer n.subMu.Unlock()
+		for _, subscription := range n.active {
+			close(subscription.err)
+		}
+		for _, subscription := range n.inactive {
+			close(subscription.err)
+		}
+		close(n.closed)
+	}()
+	return n
 }
 
 // NotifierFromContext returns the Notifier value stored in ctx, if any.
@@ -104,7 +122,7 @@ func (n *Notifier) Notify(id ID, data interface{}) error {
 
 // Closed returns a channel that is closed when the RPC connection is closed.
 func (n *Notifier) Closed() <-chan interface{} {
-	return n.codec.Closed()
+	return n.closed
 }
 
 // unsubscribe a subscription.
diff --git a/rpc/subscription_test.go b/rpc/subscription_test.go
index 0ba177e63b..ee85cfc851 100644
--- a/rpc/subscription_test.go
+++ b/rpc/subscription_test.go
@@ -72,16 +72,10 @@ func (s *NotificationTestService) SomeSubscription(ctx context.Context, n, val i
 			}
 		}
 
-		select {
-		case <-notifier.Closed():
-			s.mu.Lock()
-			s.unsubscribed = true
-			s.mu.Unlock()
-		case <-subscription.Err():
-			s.mu.Lock()
-			s.unsubscribed = true
-			s.mu.Unlock()
-		}
+		<-subscription.Err()
+		s.mu.Lock()
+		s.unsubscribed = true
+		s.mu.Unlock()
 	}()
 
 	return subscription, nil