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 }