github.com/lzy4123/fabric@v2.1.1+incompatible/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  	"sync"
    11  	"time"
    12  
    13  	"github.com/pkg/errors"
    14  )
    15  
    16  const (
    17  	subscriptionBuffSize = 50
    18  )
    19  
    20  // PubSub defines a struct that one can use to:
    21  // - publish items to a topic to multiple subscribers
    22  // - and subscribe to items from a topic
    23  // The subscriptions have a TTL and are cleaned when it passes.
    24  type PubSub struct {
    25  	sync.RWMutex
    26  
    27  	// a map from topic to Set of subscriptions
    28  	subscriptions map[string]*Set
    29  }
    30  
    31  // Subscription defines a subscription to a topic
    32  // that can be used to receive publishes on
    33  type Subscription interface {
    34  	// Listen blocks until a publish was made
    35  	// to the subscription, or an error if the
    36  	// subscription's TTL passed
    37  	Listen() (interface{}, error)
    38  }
    39  
    40  type subscription struct {
    41  	top string
    42  	ttl time.Duration
    43  	c   chan interface{}
    44  }
    45  
    46  // Listen blocks until a publish was made
    47  // to the subscription, or an error if the
    48  // subscription's TTL passed
    49  func (s *subscription) Listen() (interface{}, error) {
    50  	select {
    51  	case <-time.After(s.ttl):
    52  		return nil, errors.New("timed out")
    53  	case item := <-s.c:
    54  		return item, nil
    55  	}
    56  }
    57  
    58  // NewPubSub creates a new PubSub with an empty
    59  // set of subscriptions
    60  func NewPubSub() *PubSub {
    61  	return &PubSub{
    62  		subscriptions: make(map[string]*Set),
    63  	}
    64  }
    65  
    66  // Publish publishes an item to all subscribers on the topic
    67  func (ps *PubSub) Publish(topic string, item interface{}) error {
    68  	ps.RLock()
    69  	defer ps.RUnlock()
    70  	s, subscribed := ps.subscriptions[topic]
    71  	if !subscribed {
    72  		return errors.New("no subscribers")
    73  	}
    74  	for _, sub := range s.ToArray() {
    75  		c := sub.(*subscription).c
    76  		select {
    77  		case c <- item:
    78  		default: // Not enough room in buffer, continue in order to not block publisher
    79  		}
    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  }