github.com/myafeier/fabric@v1.0.1-0.20170722181825-3a4b1f2bce86/gossip/util/pubsub.go (about)

     1  /*
     2  Copyright IBM Corp. All Rights Reserved.
     3  
     4  SPDX-License-Identifier: Apache-2.0
     5  */
     6  
     7  package util
     8  
     9  import (
    10  	"errors"
    11  	"sync"
    12  	"time"
    13  )
    14  
    15  const (
    16  	subscriptionBuffSize = 5
    17  )
    18  
    19  // PubSub defines a struct that one can use to:
    20  // - publish items to a topic to multiple subscribers
    21  // - and subscribe to items from a topic
    22  // The subscriptions have a TTL and are cleaned when it passes.
    23  type PubSub struct {
    24  	sync.RWMutex
    25  
    26  	// a map from topic to Set of subscriptions
    27  	subscriptions map[string]*Set
    28  }
    29  
    30  // Subscription defines a subscription to a topic
    31  // that can be used to receive publishes on
    32  type Subscription interface {
    33  	// Listen blocks until a publish was made
    34  	// to the subscription, or an error if the
    35  	// subscription's TTL passed
    36  	Listen() (interface{}, error)
    37  }
    38  
    39  type subscription struct {
    40  	top string
    41  	ttl time.Duration
    42  	c   chan interface{}
    43  }
    44  
    45  // Listen blocks until a publish was made
    46  // to the subscription, or an error if the
    47  // subscription's TTL passed
    48  func (s *subscription) Listen() (interface{}, error) {
    49  	select {
    50  	case <-time.After(s.ttl):
    51  		return nil, errors.New("timed out")
    52  	case item := <-s.c:
    53  		return item, nil
    54  	}
    55  }
    56  
    57  // NewPubSub creates a new PubSub with an empty
    58  // set of subscriptions
    59  func NewPubSub() *PubSub {
    60  	return &PubSub{
    61  		subscriptions: make(map[string]*Set),
    62  	}
    63  }
    64  
    65  // Publish publishes an item to all subscribers on the topic
    66  func (ps *PubSub) Publish(topic string, item interface{}) error {
    67  	ps.RLock()
    68  	defer ps.RUnlock()
    69  	s, subscribed := ps.subscriptions[topic]
    70  	if !subscribed {
    71  		return errors.New("no subscribers")
    72  	}
    73  	for _, sub := range s.ToArray() {
    74  		c := sub.(*subscription).c
    75  		// Not enough room in buffer, continue in order to not block publisher
    76  		if len(c) == subscriptionBuffSize {
    77  			continue
    78  		}
    79  		c <- item
    80  	}
    81  	return nil
    82  }
    83  
    84  // Subscribe returns a subscription to a topic that expires when given TTL passes
    85  func (ps *PubSub) Subscribe(topic string, ttl time.Duration) Subscription {
    86  	sub := &subscription{
    87  		top: topic,
    88  		ttl: ttl,
    89  		c:   make(chan interface{}, subscriptionBuffSize),
    90  	}
    91  
    92  	ps.Lock()
    93  	// Add subscription to subscriptions map
    94  	s, exists := ps.subscriptions[topic]
    95  	// If no subscription set for the topic exists, create one
    96  	if !exists {
    97  		s = NewSet()
    98  		ps.subscriptions[topic] = s
    99  	}
   100  	ps.Unlock()
   101  
   102  	// Add the subscription
   103  	s.Add(sub)
   104  
   105  	// When the timeout expires, remove the subscription
   106  	time.AfterFunc(ttl, func() {
   107  		ps.unSubscribe(sub)
   108  	})
   109  	return sub
   110  }
   111  
   112  func (ps *PubSub) unSubscribe(sub *subscription) {
   113  	ps.Lock()
   114  	defer ps.Unlock()
   115  	ps.subscriptions[sub.top].Remove(sub)
   116  	if ps.subscriptions[sub.top].Size() != 0 {
   117  		return
   118  	}
   119  	// Else, this is the last subscription for the topic.
   120  	// Remove the set from the subscriptions map
   121  	delete(ps.subscriptions, sub.top)
   122  
   123  }