github.com/whoyao/protocol@v0.0.0-20230519045905-2d8ace718ca5/utils/event_emitter.go (about)

     1  package utils
     2  
     3  import (
     4  	"container/list"
     5  	"sync"
     6  
     7  	"golang.org/x/exp/maps"
     8  
     9  	"github.com/whoyao/protocol/logger"
    10  )
    11  
    12  const defaultQueueSize = 16
    13  
    14  type EventEmitterParams struct {
    15  	QueueSize int
    16  	Logger    logger.Logger
    17  }
    18  
    19  type EventEmitter[K comparable, V any] struct {
    20  	params    EventEmitterParams
    21  	mu        sync.RWMutex
    22  	observers map[K]*list.List
    23  }
    24  
    25  func NewEventEmitter[K comparable, V any](params EventEmitterParams) *EventEmitter[K, V] {
    26  	return &EventEmitter[K, V]{
    27  		params:    params,
    28  		observers: map[K]*list.List{},
    29  	}
    30  }
    31  
    32  func NewDefaultEventEmitter[K comparable, V any]() *EventEmitter[K, V] {
    33  	return NewEventEmitter[K, V](EventEmitterParams{
    34  		QueueSize: defaultQueueSize,
    35  		Logger:    logger.GetLogger(),
    36  	})
    37  }
    38  
    39  func (e *EventEmitter[K, V]) Emit(k K, v V) {
    40  	e.mu.RLock()
    41  	defer e.mu.RUnlock()
    42  
    43  	l, ok := e.observers[k]
    44  	if !ok {
    45  		return
    46  	}
    47  
    48  	for le := l.Front(); le != nil; le = le.Next() {
    49  		le.Value.(EventObserver[V]).emit(v)
    50  	}
    51  }
    52  
    53  func (e *EventEmitter[K, V]) Observe(k K) EventObserver[V] {
    54  	o := EventObserver[V]{
    55  		logger: e.params.Logger,
    56  		ch:     make(chan V, e.params.QueueSize),
    57  	}
    58  
    59  	e.mu.Lock()
    60  	l, ok := e.observers[k]
    61  	if !ok {
    62  		l = list.New()
    63  		e.observers[k] = l
    64  	}
    65  	le := l.PushBack(o)
    66  	e.mu.Unlock()
    67  
    68  	o.stop = func() { e.stopObserving(k, le) }
    69  
    70  	return o
    71  }
    72  
    73  func (e *EventEmitter[K, V]) ObservedKeys() []K {
    74  	e.mu.Lock()
    75  	defer e.mu.Unlock()
    76  	return maps.Keys(e.observers)
    77  }
    78  
    79  func (e *EventEmitter[K, V]) stopObserving(k K, le *list.Element) {
    80  	e.mu.Lock()
    81  	defer e.mu.Unlock()
    82  
    83  	l, ok := e.observers[k]
    84  	if !ok {
    85  		return
    86  	}
    87  
    88  	l.Remove(le)
    89  	if l.Len() == 0 {
    90  		delete(e.observers, k)
    91  	}
    92  }
    93  
    94  type EventObserver[V any] struct {
    95  	logger logger.Logger
    96  	stop   func()
    97  	ch     chan V
    98  }
    99  
   100  func NewEventObserver[V any](stopFunc func()) (EventObserver[V], func(v V)) {
   101  	o := EventObserver[V]{
   102  		logger: logger.GetLogger(),
   103  		stop:   stopFunc,
   104  		ch:     make(chan V, defaultQueueSize),
   105  	}
   106  	return o, o.emit
   107  }
   108  
   109  func (o EventObserver[V]) emit(v V) {
   110  	select {
   111  	case o.ch <- v:
   112  	default:
   113  		o.logger.Warnw("could not add event to observer queue", nil)
   114  	}
   115  }
   116  
   117  func (o EventObserver[V]) Stop() {
   118  	o.stop()
   119  }
   120  
   121  func (o EventObserver[V]) Events() <-chan V {
   122  	return o.ch
   123  }