github.com/rclone/rclone@v1.66.1-0.20240517100346-7b89735ae726/lib/atexit/atexit.go (about)

     1  // Package atexit provides handling for functions you want called when
     2  // the program exits unexpectedly due to a signal.
     3  //
     4  // You should also make sure you call Run in the normal exit path.
     5  package atexit
     6  
     7  import (
     8  	"os"
     9  	"os/signal"
    10  	"sync"
    11  	"sync/atomic"
    12  
    13  	"github.com/rclone/rclone/fs"
    14  )
    15  
    16  var (
    17  	fns          = make(map[FnHandle]bool)
    18  	fnsMutex     sync.Mutex
    19  	exitChan     chan os.Signal
    20  	exitOnce     sync.Once
    21  	registerOnce sync.Once
    22  	signalled    atomic.Int32
    23  	runCalled    atomic.Int32
    24  )
    25  
    26  // FnHandle is the type of the handle returned by function `Register`
    27  // that can be used to unregister an at-exit function
    28  type FnHandle *func()
    29  
    30  // Register a function to be called on exit.
    31  // Returns a handle which can be used to unregister the function with `Unregister`.
    32  func Register(fn func()) FnHandle {
    33  	if running() {
    34  		return nil
    35  	}
    36  	fnsMutex.Lock()
    37  	fns[&fn] = true
    38  	fnsMutex.Unlock()
    39  
    40  	// Run AtExit handlers on exitSignals so everything gets tidied up properly
    41  	registerOnce.Do(func() {
    42  		exitChan = make(chan os.Signal, 1)
    43  		signal.Notify(exitChan, exitSignals...)
    44  		go func() {
    45  			sig := <-exitChan
    46  			if sig == nil {
    47  				return
    48  			}
    49  			signal.Stop(exitChan)
    50  			signalled.Store(1)
    51  			fs.Infof(nil, "Signal received: %s", sig)
    52  			Run()
    53  			fs.Infof(nil, "Exiting...")
    54  			os.Exit(exitCode(sig))
    55  		}()
    56  	})
    57  
    58  	return &fn
    59  }
    60  
    61  // Signalled returns true if an exit signal has been received
    62  func Signalled() bool {
    63  	return signalled.Load() != 0
    64  }
    65  
    66  // running returns true if run has been called
    67  func running() bool {
    68  	return runCalled.Load() != 0
    69  }
    70  
    71  // Unregister a function using the handle returned by `Register`
    72  func Unregister(handle FnHandle) {
    73  	if running() {
    74  		return
    75  	}
    76  	fnsMutex.Lock()
    77  	defer fnsMutex.Unlock()
    78  	delete(fns, handle)
    79  }
    80  
    81  // IgnoreSignals disables the signal handler and prevents Run from being executed automatically
    82  func IgnoreSignals() {
    83  	if running() {
    84  		return
    85  	}
    86  	registerOnce.Do(func() {})
    87  	if exitChan != nil {
    88  		signal.Stop(exitChan)
    89  		close(exitChan)
    90  		exitChan = nil
    91  	}
    92  }
    93  
    94  // Run all the at exit functions if they haven't been run already
    95  func Run() {
    96  	runCalled.Store(1)
    97  	// Take the lock here (not inside the exitOnce) so we wait
    98  	// until the exit handlers have run before any calls to Run()
    99  	// return.
   100  	fnsMutex.Lock()
   101  	defer fnsMutex.Unlock()
   102  	exitOnce.Do(func() {
   103  		for fnHandle := range fns {
   104  			(*fnHandle)()
   105  		}
   106  	})
   107  }
   108  
   109  // OnError registers fn with atexit and returns a function which
   110  // runs fn() if *perr != nil and deregisters fn
   111  //
   112  // It should be used in a defer statement normally so
   113  //
   114  //	defer OnError(&err, cancelFunc)()
   115  //
   116  // So cancelFunc will be run if the function exits with an error or
   117  // at exit.
   118  //
   119  // cancelFunc will only be run once.
   120  func OnError(perr *error, fn func()) func() {
   121  	var once sync.Once
   122  	onceFn := func() {
   123  		once.Do(fn)
   124  	}
   125  	handle := Register(onceFn)
   126  	return func() {
   127  		defer Unregister(handle)
   128  		if *perr != nil {
   129  			onceFn()
   130  		}
   131  	}
   132  
   133  }