github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/module/lifecycle/lifecycle.go (about)

     1  package lifecycle
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // LifecycleManager is a support struct for implementing module.ReadyDoneAware
     8  type LifecycleManager struct {
     9  	stateTransition   sync.Mutex    // lock for preventing concurrent state transitions
    10  	started           chan struct{} // used to signal that startup has completed
    11  	stopped           chan struct{} // used to signal that shutdown has completed
    12  	startupCommenced  bool          // indicates whether OnStart() has been invoked
    13  	shutdownCommenced bool          // indicates whether OnStop() has been invoked
    14  	shutdownSignal    chan struct{} // used to signal that shutdown has commenced
    15  }
    16  
    17  func NewLifecycleManager() *LifecycleManager {
    18  	return &LifecycleManager{
    19  		stateTransition:   sync.Mutex{},
    20  		startupCommenced:  false,
    21  		started:           make(chan struct{}),
    22  		shutdownCommenced: false,
    23  		stopped:           make(chan struct{}),
    24  		shutdownSignal:    make(chan struct{}),
    25  	}
    26  }
    27  
    28  // OnStart will commence startup of the LifecycleManager. If OnStop has already been called
    29  // before the first call to OnStart, startup will not be performed. After the first call,
    30  // subsequent calls to OnStart do nothing.
    31  func (lm *LifecycleManager) OnStart(startupFns ...func()) {
    32  	lm.stateTransition.Lock()
    33  	if lm.shutdownCommenced || lm.startupCommenced {
    34  		lm.stateTransition.Unlock()
    35  		return
    36  	}
    37  	lm.startupCommenced = true
    38  	lm.stateTransition.Unlock()
    39  
    40  	go func() {
    41  		for _, fn := range startupFns {
    42  			fn()
    43  		}
    44  		close(lm.started)
    45  	}()
    46  }
    47  
    48  // OnStop will commence shutdown of the LifecycleManager. If the LifecycleManager is still
    49  // starting up, we will wait for startup to complete before shutting down. After the first
    50  // call, subsequent calls to OnStop do nothing.
    51  func (lm *LifecycleManager) OnStop(shutdownFns ...func()) {
    52  	lm.stateTransition.Lock()
    53  	if lm.shutdownCommenced {
    54  		lm.stateTransition.Unlock()
    55  		return
    56  	}
    57  	lm.shutdownCommenced = true
    58  	waitForStartup := lm.startupCommenced
    59  	lm.stateTransition.Unlock()
    60  
    61  	close(lm.shutdownSignal)
    62  	go func() {
    63  		if waitForStartup {
    64  			<-lm.started
    65  			for _, fn := range shutdownFns {
    66  				fn()
    67  			}
    68  		}
    69  		close(lm.stopped)
    70  	}()
    71  }
    72  
    73  // ShutdownSignal returns a channel that is closed when shutdown has commenced.
    74  func (lm *LifecycleManager) ShutdownSignal() <-chan struct{} {
    75  	return lm.shutdownSignal
    76  }
    77  
    78  // Started returns a channel that is closed when startup has completed.
    79  // If the LifecycleManager is stopped before OnStart() is ever called,
    80  // the returned channel will never be closed.
    81  func (lm *LifecycleManager) Started() <-chan struct{} {
    82  	return lm.started
    83  }
    84  
    85  // Stopped returns a channel that is closed when shutdown has completed
    86  func (lm *LifecycleManager) Stopped() <-chan struct{} {
    87  	return lm.stopped
    88  }