github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/helper/broker/notify.go (about)

     1  package broker
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/hashicorp/nomad/helper"
     7  )
     8  
     9  // GenericNotifier allows a process to send updates to many subscribers in an
    10  // easy manner.
    11  type GenericNotifier struct {
    12  
    13  	// publishCh is the channel used to receive the update which will be sent
    14  	// to all subscribers.
    15  	publishCh chan interface{}
    16  
    17  	// subscribeCh and unsubscribeCh are the channels used to modify the
    18  	// subscription membership mapping.
    19  	subscribeCh   chan chan interface{}
    20  	unsubscribeCh chan chan interface{}
    21  }
    22  
    23  // NewGenericNotifier returns a generic notifier which can be used by a process
    24  // to notify many subscribers when a specific update is triggered.
    25  func NewGenericNotifier() *GenericNotifier {
    26  	return &GenericNotifier{
    27  		publishCh:     make(chan interface{}, 1),
    28  		subscribeCh:   make(chan chan interface{}, 1),
    29  		unsubscribeCh: make(chan chan interface{}, 1),
    30  	}
    31  }
    32  
    33  // Notify allows the implementer to notify all subscribers with a specific
    34  // update. There is no guarantee the order in which subscribers receive the
    35  // message which is sent linearly.
    36  func (g *GenericNotifier) Notify(msg interface{}) {
    37  	select {
    38  	case g.publishCh <- msg:
    39  	default:
    40  	}
    41  }
    42  
    43  // Run is a long-lived process which handles updating subscribers as well as
    44  // ensuring any update is sent to them. The passed stopCh is used to coordinate
    45  // shutdown.
    46  func (g *GenericNotifier) Run(stopCh <-chan struct{}) {
    47  
    48  	// Store our subscribers inline with a map. This map can only be accessed
    49  	// via a single channel update at a time, meaning we can manage without
    50  	// using a lock.
    51  	subscribers := map[chan interface{}]struct{}{}
    52  
    53  	for {
    54  		select {
    55  		case <-stopCh:
    56  			return
    57  		case msgCh := <-g.subscribeCh:
    58  			subscribers[msgCh] = struct{}{}
    59  		case msgCh := <-g.unsubscribeCh:
    60  			delete(subscribers, msgCh)
    61  		case update := <-g.publishCh:
    62  			for subscriberCh := range subscribers {
    63  
    64  				// The subscribers channels are buffered, but ensure we don't
    65  				// block the whole process on this.
    66  				select {
    67  				case subscriberCh <- update:
    68  				default:
    69  				}
    70  			}
    71  		}
    72  	}
    73  }
    74  
    75  // WaitForChange allows a subscriber to wait until there is a notification
    76  // change, or the timeout is reached. The function will block until one
    77  // condition is met.
    78  func (g *GenericNotifier) WaitForChange(timeout time.Duration) interface{} {
    79  
    80  	// Create a channel and subscribe to any update. This channel is buffered
    81  	// to ensure we do not block the main broker process.
    82  	updateCh := make(chan interface{}, 1)
    83  	g.subscribeCh <- updateCh
    84  
    85  	// Create a timeout timer and use the helper to ensure this routine doesn't
    86  	// panic and making the stop call clear.
    87  	timeoutTimer, timeoutStop := helper.NewSafeTimer(timeout)
    88  
    89  	// Defer a function which performs all the required cleanup of the
    90  	// subscriber once it has been notified of a change, or reached its wait
    91  	// timeout.
    92  	defer func() {
    93  		g.unsubscribeCh <- updateCh
    94  		close(updateCh)
    95  		timeoutStop()
    96  	}()
    97  
    98  	// Enter the main loop which listens for an update or timeout and returns
    99  	// this information to the subscriber.
   100  	select {
   101  	case <-timeoutTimer.C:
   102  		return "wait timed out after " + timeout.String()
   103  	case update := <-updateCh:
   104  		return update
   105  	}
   106  }