rpc: add subscription id to events (#6386)

Addresses  #3931
This commit is contained in:
Sam Kleinman
2021-06-15 17:33:47 +00:00
committed by GitHub
parent 8d0c38257e
commit 886519e3ca
12 changed files with 221 additions and 70 deletions
+104 -27
View File
@@ -40,6 +40,7 @@ import (
"fmt"
tmsync "github.com/tendermint/tendermint/internal/libs/sync"
"github.com/tendermint/tendermint/libs/pubsub/query"
"github.com/tendermint/tendermint/libs/service"
)
@@ -73,6 +74,24 @@ type Query interface {
String() string
}
type UnsubscribeArgs struct {
ID string
Subscriber string
Query Query
}
func (args UnsubscribeArgs) Validate() error {
if args.Subscriber == "" {
return errors.New("must specify a subscriber")
}
if args.ID == "" && args.Query == nil {
return fmt.Errorf("subscription is not fully defined [subscriber=%q]", args.Subscriber)
}
return nil
}
type cmd struct {
op operation
@@ -96,8 +115,12 @@ type Server struct {
// check if we have subscription before
// subscribing or unsubscribing
mtx tmsync.RWMutex
subscriptions map[string]map[string]struct{} // subscriber -> query (string) -> empty struct
mtx tmsync.RWMutex
// subscriber -> [query->id (string) OR id->query (string))],
// track connections both by ID (new) and query (legacy) to
// avoid breaking the interface.
subscriptions map[string]map[string]string
}
// Option sets a parameter for the server.
@@ -108,7 +131,7 @@ type Option func(*Server)
// provided, the resulting server's queue is unbuffered.
func NewServer(options ...Option) *Server {
s := &Server{
subscriptions: make(map[string]map[string]struct{}),
subscriptions: make(map[string]map[string]string),
}
s.BaseService = *service.NewBaseService(nil, "PubSub", s)
@@ -186,9 +209,10 @@ func (s *Server) subscribe(ctx context.Context, clientID string, query Query, ou
case s.cmds <- cmd{op: sub, clientID: clientID, query: query, subscription: subscription}:
s.mtx.Lock()
if _, ok = s.subscriptions[clientID]; !ok {
s.subscriptions[clientID] = make(map[string]struct{})
s.subscriptions[clientID] = make(map[string]string)
}
s.subscriptions[clientID][query.String()] = struct{}{}
s.subscriptions[clientID][query.String()] = subscription.id
s.subscriptions[clientID][subscription.id] = query.String()
s.mtx.Unlock()
return subscription, nil
case <-ctx.Done():
@@ -201,23 +225,45 @@ func (s *Server) subscribe(ctx context.Context, clientID string, query Query, ou
// Unsubscribe removes the subscription on the given query. An error will be
// returned to the caller if the context is canceled or if subscription does
// not exist.
func (s *Server) Unsubscribe(ctx context.Context, clientID string, query Query) error {
s.mtx.RLock()
clientSubscriptions, ok := s.subscriptions[clientID]
if ok {
_, ok = clientSubscriptions[query.String()]
func (s *Server) Unsubscribe(ctx context.Context, args UnsubscribeArgs) error {
if err := args.Validate(); err != nil {
return err
}
var qs string
if args.Query != nil {
qs = args.Query.String()
}
s.mtx.RLock()
clientSubscriptions, ok := s.subscriptions[args.Subscriber]
if args.ID != "" {
qs, ok = clientSubscriptions[args.ID]
if ok && args.Query == nil {
var err error
args.Query, err = query.New(qs)
if err != nil {
return err
}
}
} else if qs != "" {
args.ID, ok = clientSubscriptions[qs]
}
s.mtx.RUnlock()
if !ok {
return ErrSubscriptionNotFound
}
select {
case s.cmds <- cmd{op: unsub, clientID: clientID, query: query}:
case s.cmds <- cmd{op: unsub, clientID: args.Subscriber, query: args.Query, subscription: &Subscription{id: args.ID}}:
s.mtx.Lock()
delete(clientSubscriptions, query.String())
delete(clientSubscriptions, args.ID)
delete(clientSubscriptions, qs)
if len(clientSubscriptions) == 0 {
delete(s.subscriptions, clientID)
delete(s.subscriptions, args.Subscriber)
}
s.mtx.Unlock()
return nil
@@ -262,7 +308,7 @@ func (s *Server) NumClients() int {
func (s *Server) NumClientSubscriptions(clientID string) int {
s.mtx.RLock()
defer s.mtx.RUnlock()
return len(s.subscriptions[clientID])
return len(s.subscriptions[clientID]) / 2
}
// Publish publishes the given message. An error will be returned to the caller
@@ -325,7 +371,7 @@ loop:
switch cmd.op {
case unsub:
if cmd.query != nil {
state.remove(cmd.clientID, cmd.query.String(), ErrUnsubscribed)
state.remove(cmd.clientID, cmd.query.String(), cmd.subscription.id, ErrUnsubscribed)
} else {
state.removeClient(cmd.clientID, ErrUnsubscribed)
}
@@ -349,8 +395,14 @@ func (state *state) add(clientID string, q Query, subscription *Subscription) {
if _, ok := state.subscriptions[qStr]; !ok {
state.subscriptions[qStr] = make(map[string]*Subscription)
}
if _, ok := state.subscriptions[subscription.id]; !ok {
state.subscriptions[subscription.id] = make(map[string]*Subscription)
}
// create subscription
state.subscriptions[qStr][clientID] = subscription
state.subscriptions[subscription.id][clientID] = subscription
// initialize query if needed
if _, ok := state.queries[qStr]; !ok {
@@ -360,7 +412,7 @@ func (state *state) add(clientID string, q Query, subscription *Subscription) {
state.queries[qStr].refCount++
}
func (state *state) remove(clientID string, qStr string, reason error) {
func (state *state) remove(clientID string, qStr, id string, reason error) {
clientSubscriptions, ok := state.subscriptions[qStr]
if !ok {
return
@@ -376,37 +428,62 @@ func (state *state) remove(clientID string, qStr string, reason error) {
// remove client from query map.
// if query has no other clients subscribed, remove it.
delete(state.subscriptions[qStr], clientID)
delete(state.subscriptions[id], clientID)
if len(state.subscriptions[qStr]) == 0 {
delete(state.subscriptions, qStr)
}
// decrease ref counter in queries
state.queries[qStr].refCount--
// remove the query if nobody else is using it
if state.queries[qStr].refCount == 0 {
delete(state.queries, qStr)
if ref, ok := state.queries[qStr]; ok {
ref.refCount--
if ref.refCount == 0 {
// remove the query if nobody else is using it
delete(state.queries, qStr)
}
}
}
func (state *state) removeClient(clientID string, reason error) {
seen := map[string]struct{}{}
for qStr, clientSubscriptions := range state.subscriptions {
if _, ok := clientSubscriptions[clientID]; ok {
state.remove(clientID, qStr, reason)
if sub, ok := clientSubscriptions[clientID]; ok {
if _, ok = seen[sub.id]; ok {
// all subscriptions are double indexed by ID and query, only
// process them once.
continue
}
state.remove(clientID, qStr, sub.id, reason)
seen[sub.id] = struct{}{}
}
}
}
func (state *state) removeAll(reason error) {
for qStr, clientSubscriptions := range state.subscriptions {
sub, ok := clientSubscriptions[qStr]
if !ok || ok && sub.id == qStr {
// all subscriptions are double indexed by ID and query, only
// process them once.
continue
}
for clientID := range clientSubscriptions {
state.remove(clientID, qStr, reason)
state.remove(clientID, qStr, sub.id, reason)
}
}
}
func (state *state) send(msg interface{}, events map[string][]string) error {
for qStr, clientSubscriptions := range state.subscriptions {
q := state.queries[qStr].q
if sub, ok := clientSubscriptions[qStr]; ok && sub.id == qStr {
continue
}
var q Query
if qi, ok := state.queries[qStr]; ok {
q = qi.q
} else {
continue
}
match, err := q.Matches(events)
if err != nil {
@@ -417,13 +494,13 @@ func (state *state) send(msg interface{}, events map[string][]string) error {
for clientID, subscription := range clientSubscriptions {
if cap(subscription.out) == 0 {
// block on unbuffered channel
subscription.out <- NewMessage(msg, events)
subscription.out <- NewMessage(subscription.id, msg, events)
} else {
// don't block on buffered channels
select {
case subscription.out <- NewMessage(msg, events):
case subscription.out <- NewMessage(subscription.id, msg, events):
default:
state.remove(clientID, qStr, ErrOutOfCapacity)
state.remove(clientID, qStr, subscription.id, ErrOutOfCapacity)
}
}
}
+10 -4
View File
@@ -291,7 +291,9 @@ func TestUnsubscribe(t *testing.T) {
ctx := context.Background()
subscription, err := s.Subscribe(ctx, clientID, query.MustParse("tm.events.type='NewBlock'"))
require.NoError(t, err)
err = s.Unsubscribe(ctx, clientID, query.MustParse("tm.events.type='NewBlock'"))
err = s.Unsubscribe(ctx, pubsub.UnsubscribeArgs{
Subscriber: clientID,
Query: query.MustParse("tm.events.type='NewBlock'")})
require.NoError(t, err)
err = s.Publish(ctx, "Nick Fury")
@@ -315,10 +317,14 @@ func TestClientUnsubscribesTwice(t *testing.T) {
ctx := context.Background()
_, err = s.Subscribe(ctx, clientID, query.MustParse("tm.events.type='NewBlock'"))
require.NoError(t, err)
err = s.Unsubscribe(ctx, clientID, query.MustParse("tm.events.type='NewBlock'"))
err = s.Unsubscribe(ctx, pubsub.UnsubscribeArgs{
Subscriber: clientID,
Query: query.MustParse("tm.events.type='NewBlock'")})
require.NoError(t, err)
err = s.Unsubscribe(ctx, clientID, query.MustParse("tm.events.type='NewBlock'"))
err = s.Unsubscribe(ctx, pubsub.UnsubscribeArgs{
Subscriber: clientID,
Query: query.MustParse("tm.events.type='NewBlock'")})
assert.Equal(t, pubsub.ErrSubscriptionNotFound, err)
err = s.UnsubscribeAll(ctx, clientID)
assert.Equal(t, pubsub.ErrSubscriptionNotFound, err)
@@ -338,7 +344,7 @@ func TestResubscribe(t *testing.T) {
ctx := context.Background()
_, err = s.Subscribe(ctx, clientID, query.Empty{})
require.NoError(t, err)
err = s.Unsubscribe(ctx, clientID, query.Empty{})
err = s.Unsubscribe(ctx, pubsub.UnsubscribeArgs{Subscriber: clientID, Query: query.Empty{}})
require.NoError(t, err)
subscription, err := s.Subscribe(ctx, clientID, query.Empty{})
require.NoError(t, err)
+31 -10
View File
@@ -2,7 +2,9 @@ package pubsub
import (
"errors"
"fmt"
"github.com/google/uuid"
tmsync "github.com/tendermint/tendermint/internal/libs/sync"
)
@@ -21,6 +23,7 @@ var (
// 2) channel which is closed if a client is too slow or choose to unsubscribe
// 3) err indicating the reason for (2)
type Subscription struct {
id string
out chan Message
canceled chan struct{}
@@ -31,6 +34,7 @@ type Subscription struct {
// NewSubscription returns a new subscription with the given outCapacity.
func NewSubscription(outCapacity int) *Subscription {
return &Subscription{
id: uuid.NewString(),
out: make(chan Message, outCapacity),
canceled: make(chan struct{}),
}
@@ -43,6 +47,8 @@ func (s *Subscription) Out() <-chan Message {
return s.out
}
func (s *Subscription) ID() string { return s.id }
// Canceled returns a channel that's closed when the subscription is
// terminated and supposed to be used in a select statement.
func (s *Subscription) Canceled() <-chan struct{} {
@@ -64,27 +70,42 @@ func (s *Subscription) Err() error {
func (s *Subscription) cancel(err error) {
s.mtx.Lock()
s.err = err
s.mtx.Unlock()
defer s.mtx.Unlock()
defer func() {
perr := recover()
if err == nil && perr != nil {
err = fmt.Errorf("problem closing subscription: %v", perr)
}
}()
if s.err == nil && err != nil {
s.err = err
}
close(s.canceled)
}
// Message glues data and events together.
type Message struct {
subID string
data interface{}
events map[string][]string
}
func NewMessage(data interface{}, events map[string][]string) Message {
return Message{data, events}
func NewMessage(subID string, data interface{}, events map[string][]string) Message {
return Message{
subID: subID,
data: data,
events: events,
}
}
// SubscriptionID returns the unique identifier for the subscription
// that produced this message.
func (msg Message) SubscriptionID() string { return msg.subID }
// Data returns an original data published.
func (msg Message) Data() interface{} {
return msg.data
}
func (msg Message) Data() interface{} { return msg.data }
// Events returns events, which matched the client's query.
func (msg Message) Events() map[string][]string {
return msg.events
}
func (msg Message) Events() map[string][]string { return msg.events }