github.com/andy2046/gopie@v0.7.0/pkg/pubsub/pubsub.go (about)

     1  // Package pubsub provides a pubsub implementation.
     2  package pubsub
     3  
     4  import (
     5  	"context"
     6  	"crypto/rand"
     7  	"errors"
     8  	"fmt"
     9  	"log"
    10  	"sync"
    11  )
    12  
    13  type (
    14  	// PubSub is a Pub/Sub instance for a single project.
    15  	PubSub struct {
    16  		projectID string
    17  		topics    map[string]*Topic
    18  		mu        sync.RWMutex
    19  	}
    20  
    21  	// Topic represents a PubSub topic.
    22  	Topic struct {
    23  		// The identifier for the topic,
    24  		// in the format "projects/<projid>/topics/<name>".
    25  		name string
    26  
    27  		inbox chan Message
    28  
    29  		stopped bool
    30  
    31  		pubSub *PubSub
    32  
    33  		subscriptions map[string]*Subscription
    34  
    35  		// Errors is the error output channel back to the user. You MUST read from this
    36  		// channel or the Publish will deadlock when the channel is full.
    37  		Errors chan PublishError
    38  
    39  		numGoroutines int
    40  
    41  		wgPublish sync.WaitGroup
    42  
    43  		once sync.Once
    44  
    45  		mu sync.RWMutex
    46  
    47  		wg sync.WaitGroup
    48  	}
    49  
    50  	// Subscription represents a PubSub subscription.
    51  	Subscription struct {
    52  		// The identifier for the subscription,
    53  		// in the format "projects/<projid>/topics/<name>/subscriptions/<name>".
    54  		name string
    55  
    56  		inbox chan Message
    57  
    58  		topic *Topic
    59  
    60  		stopped bool
    61  
    62  		mu sync.RWMutex
    63  
    64  		wg sync.WaitGroup
    65  
    66  		done chan struct{}
    67  
    68  		// numGoroutines is the number of goroutines it will spawn to pull msg concurrently.
    69  		numGoroutines int
    70  	}
    71  
    72  	// Message represents a Pub/Sub message.
    73  	Message struct {
    74  		ID   string
    75  		Data []byte
    76  	}
    77  
    78  	// PublishError is the error generated when it fails to publish a message.
    79  	PublishError struct {
    80  		Msg *Message
    81  		Err error
    82  	}
    83  )
    84  
    85  func (pe PublishError) Error() string {
    86  	return fmt.Sprintf("failed to publish message %s -> %s", pe.Msg.ID, pe.Err)
    87  }
    88  
    89  // UUID generates uuid.
    90  func UUID() string {
    91  	b := make([]byte, 16)
    92  	rand.Read(b)
    93  	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
    94  }
    95  
    96  // New creates a new PubSub.
    97  func New(project string) *PubSub {
    98  	return &PubSub{
    99  		projectID: project,
   100  		topics:    make(map[string]*Topic),
   101  	}
   102  }
   103  
   104  // Name returns the full name for the PubSub.
   105  func (p *PubSub) Name() string {
   106  	return fmt.Sprintf("projects/%s", p.projectID)
   107  }
   108  
   109  // NewTopic creates a new Topic with the given name,
   110  // size is the channel buffer size for topic message chan,
   111  // numGoroutines is the number of goroutines it will spawn to push msg concurrently.
   112  func (p *PubSub) NewTopic(name string, size int, numGoroutines int) (*Topic, error) {
   113  	p.mu.RLock()
   114  	if _, ok := p.topics[name]; ok {
   115  		p.mu.RUnlock()
   116  		return nil, errors.New("duplicated topic name")
   117  	}
   118  	p.mu.RUnlock()
   119  	p.mu.Lock()
   120  	t := &Topic{
   121  		name:          name,
   122  		pubSub:        p,
   123  		subscriptions: make(map[string]*Subscription),
   124  		inbox:         make(chan Message, size),
   125  		Errors:        make(chan PublishError, size),
   126  		numGoroutines: numGoroutines,
   127  	}
   128  	p.topics[name] = t
   129  	p.mu.Unlock()
   130  	return t, nil
   131  }
   132  
   133  // Topic returns the topic by name.
   134  func (p *PubSub) Topic(name string) *Topic {
   135  	p.mu.RLock()
   136  	defer p.mu.RUnlock()
   137  	if t, ok := p.topics[name]; ok {
   138  		return t
   139  	}
   140  	return nil
   141  }
   142  
   143  // Topics list all the topics in the PubSub.
   144  func (p *PubSub) Topics() []string {
   145  	p.mu.RLock()
   146  	defer p.mu.RUnlock()
   147  	ts := make([]string, 0, len(p.topics))
   148  	for k := range p.topics {
   149  		ts = append(ts, k)
   150  	}
   151  	return ts
   152  }
   153  
   154  // Publish publishes msg to the topic asynchronously.
   155  func (t *Topic) Publish(ctx context.Context, msg *Message) error {
   156  	t.mu.RLock()
   157  	defer t.mu.RUnlock()
   158  	if t.stopped {
   159  		return errors.New("topic stopped")
   160  	}
   161  	t.wgPublish.Add(1)
   162  	go func() {
   163  		select {
   164  		case <-ctx.Done():
   165  			t.Errors <- PublishError{
   166  				msg,
   167  				ctx.Err(),
   168  			}
   169  		case t.inbox <- *msg:
   170  		}
   171  		t.wgPublish.Done()
   172  	}()
   173  	return nil
   174  }
   175  
   176  // Delete removes itself from PubSuband stop it.
   177  func (t *Topic) Delete() {
   178  	t.pubSub.mu.Lock()
   179  	for k := range t.pubSub.topics {
   180  		if k == t.name {
   181  			delete(t.pubSub.topics, k)
   182  		}
   183  	}
   184  	t.pubSub.mu.Unlock()
   185  	t.Stop()
   186  }
   187  
   188  // Name returns the full name for the topic.
   189  func (t *Topic) Name() string {
   190  	return fmt.Sprintf("projects/%s/topics/%s", t.pubSub.projectID, t.name)
   191  }
   192  
   193  // Stop stops the topic.
   194  func (t *Topic) Stop() {
   195  	t.mu.Lock()
   196  	t.stopped = true
   197  	t.wgPublish.Wait()
   198  	close(t.inbox)
   199  	t.mu.Unlock()
   200  	t.wg.Wait()
   201  
   202  	for _, v := range t.subscriptions {
   203  		go func(s *Subscription) {
   204  			s.mu.Lock()
   205  			if s.stopped {
   206  				s.mu.Unlock()
   207  				return
   208  			}
   209  			s.stopped = true
   210  			close(s.inbox)
   211  			s.wg.Wait()
   212  			close(s.done)
   213  			s.mu.Unlock()
   214  		}(v)
   215  	}
   216  }
   217  
   218  func (t *Topic) start() {
   219  	for {
   220  		m, open := <-t.inbox
   221  		if !open {
   222  			log.Printf("topic %s inbox closed, exit", t.Name())
   223  			t.wg.Done()
   224  			return
   225  		}
   226  		t.mu.RLock()
   227  		subs := make(map[string]*Subscription, len(t.subscriptions))
   228  		for k, v := range t.subscriptions {
   229  			subs[k] = v
   230  		}
   231  		t.mu.RUnlock()
   232  
   233  		for _, v := range subs {
   234  			go func(s *Subscription) {
   235  				s.mu.RLock()
   236  				if s.stopped {
   237  					s.mu.RUnlock()
   238  					return
   239  				}
   240  				s.inbox <- m
   241  				s.mu.RUnlock()
   242  			}(v)
   243  		}
   244  	}
   245  }
   246  
   247  // Subscriptions list all the subscriptions to this topic.
   248  func (t *Topic) Subscriptions() []string {
   249  	t.mu.RLock()
   250  	defer t.mu.RUnlock()
   251  	sub := make([]string, 0, len(t.subscriptions))
   252  	for k := range t.subscriptions {
   253  		sub = append(sub, k)
   254  	}
   255  	return sub
   256  }
   257  
   258  // NewSubscription creates a new Subscription to this topic,
   259  // numGoroutines is the number of goroutines it will spawn to pull msg concurrently.
   260  func (t *Topic) NewSubscription(numGoroutines int) (*Subscription, error) {
   261  	t.once.Do(func() {
   262  		for range make([]struct{}, t.numGoroutines) {
   263  			t.wg.Add(1)
   264  			go t.start()
   265  		}
   266  	})
   267  	t.mu.Lock()
   268  	n := fmt.Sprintf("%s-sub-%s", t.name, UUID())
   269  	s := &Subscription{
   270  		name:          n,
   271  		inbox:         make(chan Message, 10*numGoroutines),
   272  		topic:         t,
   273  		done:          make(chan struct{}),
   274  		numGoroutines: numGoroutines,
   275  	}
   276  	t.subscriptions[n] = s
   277  	t.mu.Unlock()
   278  	return s, nil
   279  }
   280  
   281  // Subscription returns the subscription by name..
   282  func (t *Topic) Subscription(name string) *Subscription {
   283  	t.mu.RLock()
   284  	defer t.mu.RUnlock()
   285  	if s, ok := t.subscriptions[name]; ok {
   286  		return s
   287  	}
   288  	return nil
   289  }
   290  
   291  // Receive receives message for this subscription.
   292  func (s *Subscription) Receive(f func(*Message)) {
   293  	for range make([]struct{}, s.numGoroutines) {
   294  		s.wg.Add(1)
   295  		d, i := s.done, s.inbox
   296  		go func() {
   297  			for {
   298  				select {
   299  				case <-d:
   300  					d = nil
   301  					return
   302  				case m, open := <-i:
   303  					if !open {
   304  						i = nil
   305  						s.wg.Done()
   306  						break
   307  					}
   308  					f(&m)
   309  				}
   310  			}
   311  		}()
   312  	}
   313  }
   314  
   315  // Delete unsubscribes itself from topic.
   316  func (s *Subscription) Delete() {
   317  	s.topic.mu.Lock()
   318  	for k := range s.topic.subscriptions {
   319  		if k == s.name {
   320  			delete(s.topic.subscriptions, k)
   321  		}
   322  	}
   323  	s.topic.mu.Unlock()
   324  	s.mu.Lock()
   325  	if s.stopped {
   326  		s.mu.Unlock()
   327  		return
   328  	}
   329  	s.stopped = true
   330  	close(s.inbox)
   331  	s.wg.Wait()
   332  	close(s.done)
   333  	s.mu.Unlock()
   334  }