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]