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)