github.com/tenywen/fabric@v1.0.0-beta.0.20170620030522-a5b1ed380643/gossip/util/pubsub.go (about)

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