github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/pubsub/inmemory.go (about)

     1  package pubsub
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"sync"
     7  )
     8  
     9  // InMemoryPubSub is a simple in-memory pubsub implementation used for testing
    10  type InMemoryPubSub[T any] struct {
    11  	subscriber     Subscriber[T]
    12  	subscriberOnce sync.Once
    13  }
    14  
    15  func NewInMemoryPubSub[T any]() *InMemoryPubSub[T] {
    16  	return &InMemoryPubSub[T]{}
    17  }
    18  
    19  func (p *InMemoryPubSub[T]) Publish(ctx context.Context, message T) error {
    20  	if p.subscriber != nil {
    21  		return p.subscriber.Handle(ctx, message)
    22  	}
    23  	return nil
    24  }
    25  
    26  func (p *InMemoryPubSub[T]) Subscribe(ctx context.Context, subscriber Subscriber[T]) error {
    27  	var firstSubscriber bool
    28  	p.subscriberOnce.Do(func() {
    29  		p.subscriber = subscriber
    30  		firstSubscriber = true
    31  	})
    32  	if !firstSubscriber {
    33  		return errors.New("only a single subscriber is allowed. Use ChainedSubscriber to chain multiple subscribers")
    34  	}
    35  	return nil
    36  }
    37  
    38  func (p *InMemoryPubSub[T]) Close(ctx context.Context) error {
    39  	return nil
    40  }
    41  
    42  // InMemorySubscriber is a simple in-memory subscriber implementation used for testing
    43  type InMemorySubscriber[T any] struct {
    44  	events        []T
    45  	badSubscriber bool
    46  }
    47  
    48  func NewInMemorySubscriber[T any]() *InMemorySubscriber[T] {
    49  	return &InMemorySubscriber[T]{
    50  		events: make([]T, 0),
    51  	}
    52  }
    53  
    54  func (s *InMemorySubscriber[T]) Handle(ctx context.Context, message T) error {
    55  	if s.badSubscriber {
    56  		return errors.New("failed to handler message as I am a bad subscriber")
    57  	}
    58  	s.events = append(s.events, message)
    59  	return nil
    60  }
    61  
    62  func (s *InMemorySubscriber[T]) Events() []T {
    63  	res := s.events
    64  	s.events = make([]T, 0)
    65  	return res
    66  }
    67  
    68  // compile-time interface assertions
    69  var _ PubSub[string] = (*InMemoryPubSub[string])(nil)
    70  var _ Subscriber[string] = (*InMemorySubscriber[string])(nil)