github.com/hernad/nomad@v1.6.112/helper/broker/notify.go (about)

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