v8.run/go/exp@v0.0.26-0.20230226010534-afcdbd3f782d/broadcast/broadcast.go (about)

     1  package broadcast
     2  
     3  import "sync"
     4  
     5  type Topic[T any] struct {
     6  	subscriptions []*Subscription[T]
     7  	rwmu          sync.RWMutex
     8  	pool          sync.Pool
     9  }
    10  
    11  func NewTopic[T any]() *Topic[T] {
    12  	t := &Topic[T]{}
    13  	t.pool.New = func() any {
    14  		s := &Subscription[T]{
    15  			t:        t,
    16  			callback: nil,
    17  		}
    18  		return s
    19  	}
    20  
    21  	return t
    22  }
    23  
    24  type Subscription[T any] struct {
    25  	t        *Topic[T]
    26  	callback func(T)
    27  }
    28  
    29  func (t *Topic[T]) Subscribe(callback func(v T)) *Subscription[T] {
    30  	s := t.pool.Get().(*Subscription[T])
    31  	s.callback = callback
    32  	s.t = t
    33  
    34  	t.rwmu.Lock()
    35  	t.subscriptions = append(t.subscriptions, s)
    36  	t.rwmu.Unlock()
    37  
    38  	return s
    39  }
    40  
    41  func (t *Topic[T]) Unsubscribe(s *Subscription[T]) {
    42  	t.rwmu.Lock()
    43  	for i := range t.subscriptions {
    44  		if t.subscriptions[i] == s {
    45  			t.subscriptions[i] = t.subscriptions[len(t.subscriptions)-1]
    46  			t.subscriptions[len(t.subscriptions)-1] = nil
    47  			t.subscriptions = t.subscriptions[:len(t.subscriptions)-1]
    48  			break
    49  		}
    50  	}
    51  	t.rwmu.Unlock()
    52  
    53  	s.callback = nil
    54  	t.pool.Put(s)
    55  
    56  	return
    57  }
    58  
    59  func (t *Topic[T]) Subscribers() uint64 {
    60  	t.rwmu.RLock()
    61  	s := uint64(len(t.subscriptions))
    62  	t.rwmu.RUnlock()
    63  	return s
    64  }
    65  
    66  func (s *Subscription[T]) Unsubscribe() {
    67  	s.t.Unsubscribe(s)
    68  }
    69  
    70  func (s *Subscription[T]) Topic() *Topic[T] {
    71  	return s.t
    72  }
    73  
    74  func (t *Topic[T]) Broadcast(v T) {
    75  	t.rwmu.RLock()
    76  	for i := range t.subscriptions {
    77  		t.subscriptions[i].callback(v)
    78  	}
    79  	t.rwmu.RUnlock()
    80  }