github.com/uriddle/docker@v0.0.0-20210926094723-4072e6aeb013/pkg/pubsub/publisher.go (about)

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