From 46b4286c087953f6ab094cdb9ba5bf4066c565a6 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Mon, 8 Oct 2018 17:05:28 +0200
Subject: [PATCH] rpc: fix subscription corner case and speed up tests

Notifier tracks whether subscription are 'active'. A subscription
becomes active when the subscription ID has been sent to the client. If
the client sends notifications in the request handler before the
subscription becomes active they are dropped. The tests tried to work
around this problem by always waiting 5s before sending the first
notification.

Fix it by buffering notifications until the subscription becomes active.
This speeds up all subscription tests.

Also fix TestSubscriptionMultipleNamespaces to wait for three messages
instead of six. The test now finishes just after all notifications have
been received and doesn't hit the 30s timeout anymore.
---
 rpc/subscription.go      | 34 ++++++++++++++++++++++++----------
 rpc/subscription_test.go | 29 +++++++++++------------------
 2 files changed, 35 insertions(+), 28 deletions(-)

diff --git a/rpc/subscription.go b/rpc/subscription.go
index 6ce7befa1d..6bbb6f75d2 100644
--- a/rpc/subscription.go
+++ b/rpc/subscription.go
@@ -52,9 +52,10 @@ type notifierKey struct{}
 // Server callbacks use the notifier to send notifications.
 type Notifier struct {
 	codec    ServerCodec
-	subMu    sync.RWMutex // guards active and inactive maps
+	subMu    sync.Mutex
 	active   map[ID]*Subscription
 	inactive map[ID]*Subscription
+	buffer   map[ID][]interface{} // unsent notifications of inactive subscriptions
 }
 
 // newNotifier creates a new notifier that can be used to send subscription
@@ -64,6 +65,7 @@ func newNotifier(codec ServerCodec) *Notifier {
 		codec:    codec,
 		active:   make(map[ID]*Subscription),
 		inactive: make(map[ID]*Subscription),
+		buffer:   make(map[ID][]interface{}),
 	}
 }
 
@@ -88,20 +90,26 @@ func (n *Notifier) CreateSubscription() *Subscription {
 // Notify sends a notification to the client with the given data as payload.
 // If an error occurs the RPC connection is closed and the error is returned.
 func (n *Notifier) Notify(id ID, data interface{}) error {
-	n.subMu.RLock()
-	defer n.subMu.RUnlock()
+	n.subMu.Lock()
+	defer n.subMu.Unlock()
 
-	sub, active := n.active[id]
-	if active {
-		notification := n.codec.CreateNotification(string(id), sub.namespace, data)
-		if err := n.codec.Write(notification); err != nil {
-			n.codec.Close()
-			return err
-		}
+	if sub, active := n.active[id]; active {
+		n.send(sub, data)
+	} else {
+		n.buffer[id] = append(n.buffer[id], data)
 	}
 	return nil
 }
 
+func (n *Notifier) send(sub *Subscription, data interface{}) error {
+	notification := n.codec.CreateNotification(string(sub.ID), sub.namespace, data)
+	err := n.codec.Write(notification)
+	if err != nil {
+		n.codec.Close()
+	}
+	return err
+}
+
 // Closed returns a channel that is closed when the RPC connection is closed.
 func (n *Notifier) Closed() <-chan interface{} {
 	return n.codec.Closed()
@@ -127,9 +135,15 @@ func (n *Notifier) unsubscribe(id ID) error {
 func (n *Notifier) activate(id ID, namespace string) {
 	n.subMu.Lock()
 	defer n.subMu.Unlock()
+
 	if sub, found := n.inactive[id]; found {
 		sub.namespace = namespace
 		n.active[id] = sub
 		delete(n.inactive, id)
+		// Send buffered notifications.
+		for _, data := range n.buffer[id] {
+			n.send(sub, data)
+		}
+		delete(n.buffer, id)
 	}
 }
diff --git a/rpc/subscription_test.go b/rpc/subscription_test.go
index 0ba177e63b..1c38600578 100644
--- a/rpc/subscription_test.go
+++ b/rpc/subscription_test.go
@@ -65,7 +65,6 @@ func (s *NotificationTestService) SomeSubscription(ctx context.Context, n, val i
 		// test expects n events, if we begin sending event immediately some events
 		// will probably be dropped since the subscription ID might not be send to
 		// the client.
-		time.Sleep(5 * time.Second)
 		for i := 0; i < n; i++ {
 			if err := notifier.Notify(subscription.ID, val+i); err != nil {
 				return
@@ -287,32 +286,26 @@ func TestSubscriptionMultipleNamespaces(t *testing.T) {
 	timeout := time.After(30 * time.Second)
 	subids := make(map[string]string, 2*len(namespaces))
 	count := make(map[string]int, 2*len(namespaces))
-
-	for {
-		done := true
-		for id := range count {
-			if count, found := count[id]; !found || count < (2*n) {
+	allReceived := func() bool {
+		done := len(count) == len(namespaces)
+		for _, c := range count {
+			if c < n {
 				done = false
 			}
 		}
+		return done
+	}
 
-		if done && len(count) == len(namespaces) {
-			break
-		}
-
+	for !allReceived() {
 		select {
-		case err := <-errors:
-			t.Fatal(err)
 		case suc := <-successes: // subscription created
 			subids[namespaces[int(suc.Id.(float64))]] = suc.Result.(string)
+		case notification := <-notifications:
+			count[notification.Params.Subscription]++
+		case err := <-errors:
+			t.Fatal(err)
 		case failure := <-failures:
 			t.Errorf("received error: %v", failure.Error)
-		case notification := <-notifications:
-			if cnt, found := count[notification.Params.Subscription]; found {
-				count[notification.Params.Subscription] = cnt + 1
-			} else {
-				count[notification.Params.Subscription] = 1
-			}
 		case <-timeout:
 			for _, namespace := range namespaces {
 				subid, found := subids[namespace]