Merge pull request #96 from tendermint/tm-945

add String method to Query interface

pubsub/pubsub.go

@@ -13,6 +13,8 @@ package pubsub
 
 import (
 	"context"
+	"errors"
+	"sync"
 
 	cmn "github.com/tendermint/tmlibs/common"
 )
@@ -38,6 +40,7 @@ type cmd struct {
 // Query defines an interface for a query to be used for subscribing.
 type Query interface {
 	Matches(tags map[string]interface{}) bool
+	String() string
 }
 
 // Server allows clients to subscribe/unsubscribe for messages, publishing
@@ -47,6 +50,9 @@ type Server struct {
 
 	cmds    chan cmd
 	cmdsCap int
+
+	mtx           sync.RWMutex
+	subscriptions map[string]map[string]struct{} // subscriber -> query -> struct{}
 }
 
 // Option sets a parameter for the server.
@@ -56,7 +62,9 @@ type Option func(*Server)
 // for a detailed description of how to configure buffering. If no options are
 // provided, the resulting server's queue is unbuffered.
 func NewServer(options ...Option) *Server {
-	s := &Server{}
+	s := &Server{
+		subscriptions: make(map[string]map[string]struct{}),
+	}
 	s.BaseService = *cmn.NewBaseService(nil, "PubSub", s)
 
 	for _, option := range options {
@@ -82,17 +90,33 @@ func BufferCapacity(cap int) Option {
 }
 
 // BufferCapacity returns capacity of the internal server's queue.
-func (s Server) BufferCapacity() int {
+func (s *Server) BufferCapacity() int {
 	return s.cmdsCap
 }
 
 // Subscribe creates a subscription for the given client. It accepts a channel
-// on which messages matching the given query can be received. If the
-// subscription already exists, the old channel will be closed. An error will
-// be returned to the caller if the context is canceled.
+// on which messages matching the given query can be received. An error will be
+// returned to the caller if the context is canceled or if subscription already
+// exist for pair clientID and query.
 func (s *Server) Subscribe(ctx context.Context, clientID string, query Query, out chan<- interface{}) error {
+	s.mtx.RLock()
+	clientSubscriptions, ok := s.subscriptions[clientID]
+	if ok {
+		_, ok = clientSubscriptions[query.String()]
+	}
+	s.mtx.RUnlock()
+	if ok {
+		return errors.New("already subscribed")
+	}
+
 	select {
 	case s.cmds <- cmd{op: sub, clientID: clientID, query: query, ch: out}:
+		s.mtx.Lock()
+		if _, ok = s.subscriptions[clientID]; !ok {
+			s.subscriptions[clientID] = make(map[string]struct{})
+		}
+		s.subscriptions[clientID][query.String()] = struct{}{}
+		s.mtx.Unlock()
 		return nil
 	case <-ctx.Done():
 		return ctx.Err()
@@ -100,10 +124,24 @@ 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.
+// 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()]
+	}
+	s.mtx.RUnlock()
+	if !ok {
+		return errors.New("subscription not found")
+	}
+
 	select {
 	case s.cmds <- cmd{op: unsub, clientID: clientID, query: query}:
+		s.mtx.Lock()
+		delete(clientSubscriptions, query.String())
+		s.mtx.Unlock()
 		return nil
 	case <-ctx.Done():
 		return ctx.Err()
@@ -111,10 +149,20 @@ func (s *Server) Unsubscribe(ctx context.Context, clientID string, query Query)
 }
 
 // UnsubscribeAll removes all client subscriptions. An error will be returned
-// to the caller if the context is canceled.
+// to the caller if the context is canceled or if subscription does not exist.
 func (s *Server) UnsubscribeAll(ctx context.Context, clientID string) error {
+	s.mtx.RLock()
+	_, ok := s.subscriptions[clientID]
+	s.mtx.RUnlock()
+	if !ok {
+		return errors.New("subscription not found")
+	}
+
 	select {
 	case s.cmds <- cmd{op: unsub, clientID: clientID}:
+		s.mtx.Lock()
+		delete(s.subscriptions, clientID)
+		s.mtx.Unlock()
 		return nil
 	case <-ctx.Done():
 		return ctx.Err()
@@ -186,13 +234,8 @@ loop:
 
 func (state *state) add(clientID string, q Query, ch chan<- interface{}) {
 	// add query if needed
-	if clientToChannelMap, ok := state.queries[q]; !ok {
+	if _, ok := state.queries[q]; !ok {
 		state.queries[q] = make(map[string]chan<- interface{})
-	} else {
-		// check if already subscribed
-		if oldCh, ok := clientToChannelMap[clientID]; ok {
-			close(oldCh)
-		}
 	}
 
 	// create subscription

pubsub/pubsub_test.go

@@ -86,14 +86,11 @@ func TestClientSubscribesTwice(t *testing.T) {
 
 	ch2 := make(chan interface{}, 1)
 	err = s.Subscribe(ctx, clientID, q, ch2)
-	require.NoError(t, err)
-
-	_, ok := <-ch1
-	assert.False(t, ok)
+	require.Error(t, err)
 
 	err = s.PublishWithTags(ctx, "Spider-Man", map[string]interface{}{"tm.events.type": "NewBlock"})
 	require.NoError(t, err)
-	assertReceive(t, "Spider-Man", ch2)
+	assertReceive(t, "Spider-Man", ch1)
 }
 
 func TestUnsubscribe(t *testing.T) {
@@ -117,6 +114,27 @@ func TestUnsubscribe(t *testing.T) {
 	assert.False(t, ok)
 }
 
+func TestResubscribe(t *testing.T) {
+	s := pubsub.NewServer()
+	s.SetLogger(log.TestingLogger())
+	s.Start()
+	defer s.Stop()
+
+	ctx := context.Background()
+	ch := make(chan interface{})
+	err := s.Subscribe(ctx, clientID, query.Empty{}, ch)
+	require.NoError(t, err)
+	err = s.Unsubscribe(ctx, clientID, query.Empty{})
+	require.NoError(t, err)
+	ch = make(chan interface{})
+	err = s.Subscribe(ctx, clientID, query.Empty{}, ch)
+	require.NoError(t, err)
+
+	err = s.Publish(ctx, "Cable")
+	require.NoError(t, err)
+	assertReceive(t, "Cable", ch)
+}
+
 func TestUnsubscribeAll(t *testing.T) {
 	s := pubsub.NewServer()
 	s.SetLogger(log.TestingLogger())
@@ -125,9 +143,9 @@ func TestUnsubscribeAll(t *testing.T) {
 
 	ctx := context.Background()
 	ch1, ch2 := make(chan interface{}, 1), make(chan interface{}, 1)
-	err := s.Subscribe(ctx, clientID, query.Empty{}, ch1)
+	err := s.Subscribe(ctx, clientID, query.MustParse("tm.events.type='NewBlock'"), ch1)
 	require.NoError(t, err)
-	err = s.Subscribe(ctx, clientID, query.Empty{}, ch2)
+	err = s.Subscribe(ctx, clientID, query.MustParse("tm.events.type='NewBlockHeader'"), ch2)
 	require.NoError(t, err)
 
 	err = s.UnsubscribeAll(ctx, clientID)