github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/client/allocrunner/tasklifecycle/gate.go (about)

     1  package tasklifecycle
     2  
     3  const (
     4  	gateClosed = false
     5  	gateOpened = true
     6  )
     7  
     8  // Gate is used by the Coordinator to block or allow tasks from running.
     9  //
    10  // It provides a channel that taskRunners listens on to determine when they are
    11  // allowed to run. The Gate has an infinite loop that is either feeding this
    12  // channel (therefore allowing listeners to proceed) or not doing anything
    13  // (causing listeners to block an wait).
    14  //
    15  // The Coordinator uses the Gate Open() and Close() methods to control this
    16  // producer loop.
    17  type Gate struct {
    18  	sendCh     chan struct{}
    19  	updateCh   chan bool
    20  	shutdownCh <-chan struct{}
    21  }
    22  
    23  // NewGate returns a new Gate that is initially closed. The Gate should not be
    24  // used after the shutdownCh is closed.
    25  func NewGate(shutdownCh <-chan struct{}) *Gate {
    26  	g := &Gate{
    27  		sendCh:     make(chan struct{}),
    28  		updateCh:   make(chan bool),
    29  		shutdownCh: shutdownCh,
    30  	}
    31  	go g.run(gateClosed)
    32  
    33  	return g
    34  }
    35  
    36  // WaitCh returns a channel that the listener must block on before starting its
    37  // task.
    38  //
    39  // Callers must also check the state of the shutdownCh used to create the Gate
    40  // to avoid blocking indefinitely.
    41  func (g *Gate) WaitCh() <-chan struct{} {
    42  	return g.sendCh
    43  }
    44  
    45  // Open is used to allow listeners to proceed.
    46  // If the gate shutdownCh channel is closed, this method is a no-op so callers
    47  // should check its state.
    48  func (g *Gate) Open() {
    49  	select {
    50  	case <-g.shutdownCh:
    51  	case g.updateCh <- gateOpened:
    52  	}
    53  }
    54  
    55  // Close is used to block listeners from proceeding.
    56  // if the gate shutdownch channel is closed, this method is a no-op so callers
    57  // should check its state.
    58  func (g *Gate) Close() {
    59  	select {
    60  	case <-g.shutdownCh:
    61  	case g.updateCh <- gateClosed:
    62  	}
    63  }
    64  
    65  // run starts the infinite loop that feeds the channel if the Gate is opened.
    66  func (g *Gate) run(initState bool) {
    67  	isOpen := initState
    68  	for {
    69  		if isOpen {
    70  			select {
    71  			// Feed channel if the gate is open.
    72  			case g.sendCh <- struct{}{}:
    73  			case <-g.shutdownCh:
    74  				return
    75  			case isOpen = <-g.updateCh:
    76  				continue
    77  			}
    78  		} else {
    79  			select {
    80  			case <-g.shutdownCh:
    81  				return
    82  			case isOpen = <-g.updateCh:
    83  				continue
    84  			}
    85  		}
    86  	}
    87  }