github.com/xhghs/rclone@v1.51.1-0.20200430155106-e186a28cced8/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  
    12  	"github.com/rclone/rclone/fs"
    13  )
    14  
    15  var (
    16  	fns          = make(map[FnHandle]bool)
    17  	fnsMutex     sync.Mutex
    18  	exitChan     chan os.Signal
    19  	exitOnce     sync.Once
    20  	registerOnce sync.Once
    21  )
    22  
    23  // FnHandle is the type of the handle returned by function `Register`
    24  // that can be used to unregister an at-exit function
    25  type FnHandle *func()
    26  
    27  // Register a function to be called on exit.
    28  // Returns a handle which can be used to unregister the function with `Unregister`.
    29  func Register(fn func()) FnHandle {
    30  	fnsMutex.Lock()
    31  	fns[&fn] = true
    32  	fnsMutex.Unlock()
    33  
    34  	// Run AtExit handlers on exitSignals so everything gets tidied up properly
    35  	registerOnce.Do(func() {
    36  		exitChan = make(chan os.Signal, 1)
    37  		signal.Notify(exitChan, exitSignals...)
    38  		go func() {
    39  			sig := <-exitChan
    40  			if sig == nil {
    41  				return
    42  			}
    43  			fs.Infof(nil, "Signal received: %s", sig)
    44  			Run()
    45  			fs.Infof(nil, "Exiting...")
    46  			os.Exit(0)
    47  		}()
    48  	})
    49  
    50  	return &fn
    51  }
    52  
    53  // Unregister a function using the handle returned by `Register`
    54  func Unregister(handle FnHandle) {
    55  	fnsMutex.Lock()
    56  	defer fnsMutex.Unlock()
    57  	delete(fns, handle)
    58  }
    59  
    60  // IgnoreSignals disables the signal handler and prevents Run from beeing executed automatically
    61  func IgnoreSignals() {
    62  	registerOnce.Do(func() {})
    63  	if exitChan != nil {
    64  		signal.Stop(exitChan)
    65  		close(exitChan)
    66  		exitChan = nil
    67  	}
    68  }
    69  
    70  // Run all the at exit functions if they haven't been run already
    71  func Run() {
    72  	exitOnce.Do(func() {
    73  		fnsMutex.Lock()
    74  		defer fnsMutex.Unlock()
    75  		for fnHandle := range fns {
    76  			(*fnHandle)()
    77  		}
    78  	})
    79  }