github.com/olljanat/moby@v1.13.1/pkg/pubsub/publisher.go (about)

     1  package pubsub
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  )
     7  
     8  var wgPool = sync.Pool{New: func() interface{} { return new(sync.WaitGroup) }}
     9  
    10  // NewPublisher creates a new pub/sub publisher to broadcast messages.
    11  // The duration is used as the send timeout as to not block the publisher publishing
    12  // messages to other clients if one client is slow or unresponsive.
    13  // The buffer is used when creating new channels for subscribers.
    14  func NewPublisher(publishTimeout time.Duration, buffer int) *Publisher {
    15  	return &Publisher{
    16  		buffer:      buffer,
    17  		timeout:     publishTimeout,
    18  		subscribers: make(map[subscriber]topicFunc),
    19  	}
    20  }
    21  
    22  type subscriber chan interface{}
    23  type topicFunc func(v interface{}) bool
    24  
    25  // Publisher is basic pub/sub structure. Allows to send events and subscribe
    26  // to them. Can be safely used from multiple goroutines.
    27  type Publisher struct {
    28  	m           sync.RWMutex
    29  	buffer      int
    30  	timeout     time.Duration
    31  	subscribers map[subscriber]topicFunc
    32  }
    33  
    34  // Len returns the number of subscribers for the publisher
    35  func (p *Publisher) Len() int {
    36  	p.m.RLock()
    37  	i := len(p.subscribers)
    38  	p.m.RUnlock()
    39  	return i
    40  }
    41  
    42  // Subscribe adds a new subscriber to the publisher returning the channel.
    43  func (p *Publisher) Subscribe() chan interface{} {
    44  	return p.SubscribeTopic(nil)
    45  }
    46  
    47  // SubscribeTopic adds a new subscriber that filters messages sent by a topic.
    48  func (p *Publisher) SubscribeTopic(topic topicFunc) chan interface{} {
    49  	ch := make(chan interface{}, p.buffer)
    50  	p.m.Lock()
    51  	p.subscribers[ch] = topic
    52  	p.m.Unlock()
    53  	return ch
    54  }
    55  
    56  // Evict removes the specified subscriber from receiving any more messages.
    57  func (p *Publisher) Evict(sub chan interface{}) {
    58  	p.m.Lock()
    59  	delete(p.subscribers, sub)
    60  	close(sub)
    61  	p.m.Unlock()
    62  }
    63  
    64  // Publish sends the data in v to all subscribers currently registered with the publisher.
    65  func (p *Publisher) Publish(v interface{}) {
    66  	p.m.RLock()
    67  	if len(p.subscribers) == 0 {
    68  		p.m.RUnlock()
    69  		return
    70  	}
    71  
    72  	wg := wgPool.Get().(*sync.WaitGroup)
    73  	for sub, topic := range p.subscribers {
    74  		wg.Add(1)
    75  		go p.sendTopic(sub, topic, v, wg)
    76  	}
    77  	wg.Wait()
    78  	wgPool.Put(wg)
    79  	p.m.RUnlock()
    80  }
    81  
    82  // Close closes the channels to all subscribers registered with the publisher.
    83  func (p *Publisher) Close() {
    84  	p.m.Lock()
    85  	for sub := range p.subscribers {
    86  		delete(p.subscribers, sub)
    87  		close(sub)
    88  	}
    89  	p.m.Unlock()
    90  }
    91  
    92  func (p *Publisher) sendTopic(sub subscriber, topic topicFunc, v interface{}, wg *sync.WaitGroup) {
    93  	defer wg.Done()
    94  	if topic != nil && !topic(v) {
    95  		return
    96  	}
    97  
    98  	// send under a select as to not block if the receiver is unavailable
    99  	if p.timeout > 0 {
   100  		select {
   101  		case sub <- v:
   102  		case <-time.After(p.timeout):
   103  		}
   104  		return
   105  	}
   106  
   107  	select {
   108  	case sub <- v:
   109  	default:
   110  	}
   111  }