github.com/mailgun/holster/v4@v4.20.0/syncutil/broadcast.go (about)

     1  /*
     2  Copyright 2022 Mailgun Technologies Inc
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8  	http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  package syncutil
    17  
    18  import "sync"
    19  
    20  type Broadcaster interface {
    21  	WaitChan(string) chan struct{}
    22  	Wait(string)
    23  	Broadcast()
    24  	Has(string) bool
    25  	Remove(string)
    26  	Done()
    27  }
    28  
    29  // Broadcasts to goroutines a new event has occurred and any waiting go routines should
    30  // stop waiting and do work. The current implementation is limited to 10,0000 unconsumed
    31  // broadcasts. If the user broadcasts more events than can be consumed calls to broadcast()
    32  // will eventually block until the goroutines can catch up. This ensures goroutines will
    33  // receive at least one event per broadcast() call.
    34  type broadcast struct {
    35  	clients     map[string]chan struct{}
    36  	done        chan struct{}
    37  	mutex       sync.Mutex
    38  	channelSize int
    39  }
    40  
    41  type BroadcasterOption interface {
    42  	Apply(*broadcast)
    43  }
    44  
    45  const DefaultChannelSize = 10000
    46  
    47  func NewBroadcaster(opts ...BroadcasterOption) Broadcaster {
    48  	br := &broadcast{
    49  		clients:     make(map[string]chan struct{}),
    50  		done:        make(chan struct{}),
    51  		channelSize: DefaultChannelSize,
    52  	}
    53  
    54  	for _, opt := range opts {
    55  		opt.Apply(br)
    56  	}
    57  
    58  	return br
    59  }
    60  
    61  // Notify all Waiting goroutines
    62  func (b *broadcast) Broadcast() {
    63  	b.mutex.Lock()
    64  	for _, channel := range b.clients {
    65  		channel <- struct{}{}
    66  	}
    67  	b.mutex.Unlock()
    68  }
    69  
    70  // Cancels any Wait() calls that are currently blocked
    71  func (b *broadcast) Done() {
    72  	close(b.done)
    73  }
    74  
    75  // Blocks until a broadcast is received
    76  func (b *broadcast) Wait(name string) {
    77  	b.mutex.Lock()
    78  	channel, ok := b.clients[name]
    79  	if !ok {
    80  		b.clients[name] = make(chan struct{}, b.channelSize)
    81  		channel = b.clients[name]
    82  	}
    83  	b.mutex.Unlock()
    84  
    85  	// Wait for a new event or done is closed
    86  	select {
    87  	case <-channel:
    88  		return
    89  	case <-b.done:
    90  		return
    91  	}
    92  }
    93  
    94  // Returns a channel the caller can use to wait for a broadcast
    95  func (b *broadcast) WaitChan(name string) chan struct{} {
    96  	b.mutex.Lock()
    97  	channel, ok := b.clients[name]
    98  	if !ok {
    99  		b.clients[name] = make(chan struct{}, b.channelSize)
   100  		channel = b.clients[name]
   101  	}
   102  	b.mutex.Unlock()
   103  	return channel
   104  }
   105  
   106  // Has checks if a client name is registered.
   107  func (b *broadcast) Has(name string) bool {
   108  	b.mutex.Lock()
   109  	defer b.mutex.Unlock()
   110  	_, exists := b.clients[name]
   111  	return exists
   112  }
   113  
   114  // Remove client name previously registered by Wait/WaitChan.
   115  func (b *broadcast) Remove(name string) {
   116  	b.mutex.Lock()
   117  	delete(b.clients, name)
   118  	b.mutex.Unlock()
   119  }
   120  
   121  type withChannelSizeOption struct {
   122  	channelSize int
   123  }
   124  
   125  // WithChannelSize sets the client's broadcast channel size.
   126  func WithChannelSize(channelSize int) BroadcasterOption {
   127  	return &withChannelSizeOption{channelSize: channelSize}
   128  }
   129  
   130  func (o *withChannelSizeOption) Apply(b *broadcast) {
   131  	b.channelSize = o.channelSize
   132  }