You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
441 lines
12 KiB
441 lines
12 KiB
// Package pubsub implements a pub-sub model with a single publisher (Server)
|
|
// and multiple subscribers (clients).
|
|
//
|
|
// Though you can have multiple publishers by sharing a pointer to a server or
|
|
// by giving the same channel to each publisher and publishing messages from
|
|
// that channel (fan-in).
|
|
//
|
|
// Clients subscribe for messages, which could be of any type, using a query.
|
|
// When some message is published, we match it with all queries. If there is a
|
|
// match, this message will be pushed to all clients, subscribed to that query.
|
|
// See query subpackage for our implementation.
|
|
//
|
|
// Due to the blocking send implementation, a single subscriber can freeze an
|
|
// entire server by not reading messages before it unsubscribes. To avoid such
|
|
// scenario, subscribers must either:
|
|
//
|
|
// a) make sure they continue to read from the out channel until
|
|
// Unsubscribe(All) is called
|
|
//
|
|
// s.Subscribe(ctx, sub, qry, out)
|
|
// go func() {
|
|
// for msg := range out {
|
|
// // handle msg
|
|
// // will exit automatically when out is closed by Unsubscribe(All)
|
|
// }
|
|
// }()
|
|
// s.UnsubscribeAll(ctx, sub)
|
|
//
|
|
// b) drain the out channel before calling Unsubscribe(All)
|
|
//
|
|
// s.Subscribe(ctx, sub, qry, out)
|
|
// defer func() {
|
|
// // drain out to make sure we don't block
|
|
// LOOP:
|
|
// for {
|
|
// select {
|
|
// case <-out:
|
|
// default:
|
|
// break LOOP
|
|
// }
|
|
// }
|
|
// s.UnsubscribeAll(ctx, sub)
|
|
// }()
|
|
// for msg := range out {
|
|
// // handle msg
|
|
// if err != nil {
|
|
// return err
|
|
// }
|
|
// }
|
|
//
|
|
package pubsub
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
|
|
cmn "github.com/tendermint/tendermint/libs/common"
|
|
)
|
|
|
|
type operation int
|
|
|
|
const (
|
|
sub operation = iota
|
|
pub
|
|
unsub
|
|
shutdown
|
|
)
|
|
|
|
var (
|
|
// ErrSubscriptionNotFound is returned when a client tries to unsubscribe
|
|
// from not existing subscription.
|
|
ErrSubscriptionNotFound = errors.New("subscription not found")
|
|
|
|
// ErrAlreadySubscribed is returned when a client tries to subscribe twice or
|
|
// more using the same query.
|
|
ErrAlreadySubscribed = errors.New("already subscribed")
|
|
)
|
|
|
|
type cmd struct {
|
|
op operation
|
|
query Query
|
|
ch chan<- interface{}
|
|
clientID string
|
|
msg interface{}
|
|
tags TagMap
|
|
}
|
|
|
|
// Query defines an interface for a query to be used for subscribing.
|
|
type Query interface {
|
|
Matches(tags TagMap) bool
|
|
String() string
|
|
}
|
|
|
|
// Server allows clients to subscribe/unsubscribe for messages, publishing
|
|
// messages with or without tags, and manages internal state.
|
|
type Server struct {
|
|
cmn.BaseService
|
|
|
|
cmds chan cmd
|
|
cmdsCap int
|
|
|
|
mtx sync.RWMutex
|
|
subscriptions map[string]map[string]struct{} // subscriber -> query (string) -> empty struct
|
|
}
|
|
|
|
// Option sets a parameter for the server.
|
|
type Option func(*Server)
|
|
|
|
// TagMap is used to associate tags to a message.
|
|
// They can be queried by subscribers to choose messages they will received.
|
|
type TagMap interface {
|
|
// Get returns the value for a key, or nil if no value is present.
|
|
// The ok result indicates whether value was found in the tags.
|
|
Get(key string) (value string, ok bool)
|
|
// Len returns the number of tags.
|
|
Len() int
|
|
}
|
|
|
|
type tagMap map[string]string
|
|
|
|
var _ TagMap = (*tagMap)(nil)
|
|
|
|
// NewTagMap constructs a new immutable tag set from a map.
|
|
func NewTagMap(data map[string]string) TagMap {
|
|
return tagMap(data)
|
|
}
|
|
|
|
// Get returns the value for a key, or nil if no value is present.
|
|
// The ok result indicates whether value was found in the tags.
|
|
func (ts tagMap) Get(key string) (value string, ok bool) {
|
|
value, ok = ts[key]
|
|
return
|
|
}
|
|
|
|
// Len returns the number of tags.
|
|
func (ts tagMap) Len() int {
|
|
return len(ts)
|
|
}
|
|
|
|
// NewServer returns a new server. See the commentary on the Option functions
|
|
// 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{
|
|
subscriptions: make(map[string]map[string]struct{}),
|
|
}
|
|
s.BaseService = *cmn.NewBaseService(nil, "PubSub", s)
|
|
|
|
for _, option := range options {
|
|
option(s)
|
|
}
|
|
|
|
// if BufferCapacity option was not set, the channel is unbuffered
|
|
s.cmds = make(chan cmd, s.cmdsCap)
|
|
|
|
return s
|
|
}
|
|
|
|
// BufferCapacity allows you to specify capacity for the internal server's
|
|
// queue. Since the server, given Y subscribers, could only process X messages,
|
|
// this option could be used to survive spikes (e.g. high amount of
|
|
// transactions during peak hours).
|
|
func BufferCapacity(cap int) Option {
|
|
return func(s *Server) {
|
|
if cap > 0 {
|
|
s.cmdsCap = cap
|
|
}
|
|
}
|
|
}
|
|
|
|
// BufferCapacity returns capacity of the internal server's queue.
|
|
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. 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 ErrAlreadySubscribed
|
|
}
|
|
|
|
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()
|
|
case <-s.Quit():
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// 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()]
|
|
}
|
|
s.mtx.RUnlock()
|
|
if !ok {
|
|
return ErrSubscriptionNotFound
|
|
}
|
|
|
|
select {
|
|
case s.cmds <- cmd{op: unsub, clientID: clientID, query: query}:
|
|
s.mtx.Lock()
|
|
delete(clientSubscriptions, query.String())
|
|
if len(clientSubscriptions) == 0 {
|
|
delete(s.subscriptions, clientID)
|
|
}
|
|
s.mtx.Unlock()
|
|
return nil
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-s.Quit():
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// UnsubscribeAll removes all client subscriptions. An error will be returned
|
|
// 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 ErrSubscriptionNotFound
|
|
}
|
|
|
|
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()
|
|
case <-s.Quit():
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// Publish publishes the given message. An error will be returned to the caller
|
|
// if the context is canceled.
|
|
func (s *Server) Publish(ctx context.Context, msg interface{}) error {
|
|
return s.PublishWithTags(ctx, msg, NewTagMap(make(map[string]string)))
|
|
}
|
|
|
|
// PublishWithTags publishes the given message with the set of tags. The set is
|
|
// matched with clients queries. If there is a match, the message is sent to
|
|
// the client.
|
|
func (s *Server) PublishWithTags(ctx context.Context, msg interface{}, tags TagMap) error {
|
|
select {
|
|
case s.cmds <- cmd{op: pub, msg: msg, tags: tags}:
|
|
return nil
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-s.Quit():
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// OnStop implements Service.OnStop by shutting down the server.
|
|
func (s *Server) OnStop() {
|
|
s.cmds <- cmd{op: shutdown}
|
|
}
|
|
|
|
// NOTE: not goroutine safe
|
|
type state struct {
|
|
// query string -> client -> ch
|
|
queryToChanMap map[string]map[string]chan<- interface{}
|
|
// client -> query string -> struct{}
|
|
clientToQueryMap map[string]map[string]struct{}
|
|
// query string -> queryPlusRefCount
|
|
queries map[string]*queryPlusRefCount
|
|
}
|
|
|
|
// queryPlusRefCount holds a pointer to a query and reference counter. When
|
|
// refCount is zero, query will be removed.
|
|
type queryPlusRefCount struct {
|
|
q Query
|
|
refCount int
|
|
}
|
|
|
|
// OnStart implements Service.OnStart by starting the server.
|
|
func (s *Server) OnStart() error {
|
|
go s.loop(state{
|
|
queryToChanMap: make(map[string]map[string]chan<- interface{}),
|
|
clientToQueryMap: make(map[string]map[string]struct{}),
|
|
queries: make(map[string]*queryPlusRefCount),
|
|
})
|
|
return nil
|
|
}
|
|
|
|
// OnReset implements Service.OnReset
|
|
func (s *Server) OnReset() error {
|
|
return nil
|
|
}
|
|
|
|
func (s *Server) loop(state state) {
|
|
loop:
|
|
for cmd := range s.cmds {
|
|
switch cmd.op {
|
|
case unsub:
|
|
if cmd.query != nil {
|
|
state.remove(cmd.clientID, cmd.query)
|
|
} else {
|
|
state.removeAll(cmd.clientID)
|
|
}
|
|
case shutdown:
|
|
for clientID := range state.clientToQueryMap {
|
|
state.removeAll(clientID)
|
|
}
|
|
break loop
|
|
case sub:
|
|
state.add(cmd.clientID, cmd.query, cmd.ch)
|
|
case pub:
|
|
state.send(cmd.msg, cmd.tags)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (state *state) add(clientID string, q Query, ch chan<- interface{}) {
|
|
qStr := q.String()
|
|
|
|
// initialize clientToChannelMap per query if needed
|
|
if _, ok := state.queryToChanMap[qStr]; !ok {
|
|
state.queryToChanMap[qStr] = make(map[string]chan<- interface{})
|
|
}
|
|
|
|
// create subscription
|
|
state.queryToChanMap[qStr][clientID] = ch
|
|
|
|
// initialize queries if needed
|
|
if _, ok := state.queries[qStr]; !ok {
|
|
state.queries[qStr] = &queryPlusRefCount{q: q, refCount: 0}
|
|
}
|
|
// increment reference counter
|
|
state.queries[qStr].refCount++
|
|
|
|
// add client if needed
|
|
if _, ok := state.clientToQueryMap[clientID]; !ok {
|
|
state.clientToQueryMap[clientID] = make(map[string]struct{})
|
|
}
|
|
state.clientToQueryMap[clientID][qStr] = struct{}{}
|
|
}
|
|
|
|
func (state *state) remove(clientID string, q Query) {
|
|
qStr := q.String()
|
|
|
|
clientToChannelMap, ok := state.queryToChanMap[qStr]
|
|
if !ok {
|
|
return
|
|
}
|
|
|
|
ch, ok := clientToChannelMap[clientID]
|
|
if !ok {
|
|
return
|
|
}
|
|
|
|
close(ch)
|
|
|
|
// remove the query from client map.
|
|
// if client is not subscribed to anything else, remove it.
|
|
delete(state.clientToQueryMap[clientID], qStr)
|
|
if len(state.clientToQueryMap[clientID]) == 0 {
|
|
delete(state.clientToQueryMap, clientID)
|
|
}
|
|
|
|
// remove the client from query map.
|
|
// if query has no other clients subscribed, remove it.
|
|
delete(state.queryToChanMap[qStr], clientID)
|
|
if len(state.queryToChanMap[qStr]) == 0 {
|
|
delete(state.queryToChanMap, 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)
|
|
}
|
|
}
|
|
|
|
func (state *state) removeAll(clientID string) {
|
|
queryMap, ok := state.clientToQueryMap[clientID]
|
|
if !ok {
|
|
return
|
|
}
|
|
|
|
for qStr := range queryMap {
|
|
ch := state.queryToChanMap[qStr][clientID]
|
|
close(ch)
|
|
|
|
// remove the client from query map.
|
|
// if query has no other clients subscribed, remove it.
|
|
delete(state.queryToChanMap[qStr], clientID)
|
|
if len(state.queryToChanMap[qStr]) == 0 {
|
|
delete(state.queryToChanMap, 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)
|
|
}
|
|
}
|
|
|
|
// remove the client.
|
|
delete(state.clientToQueryMap, clientID)
|
|
}
|
|
|
|
func (state *state) send(msg interface{}, tags TagMap) {
|
|
for qStr, clientToChannelMap := range state.queryToChanMap {
|
|
q := state.queries[qStr].q
|
|
if q.Matches(tags) {
|
|
for _, ch := range clientToChannelMap {
|
|
ch <- msg
|
|
}
|
|
}
|
|
}
|
|
}
|