github.com/fawick/restic@v0.1.1-0.20171126184616-c02923fbfc79/cmd/restic/cleanup.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"os/signal"
     7  	"sync"
     8  	"syscall"
     9  
    10  	"github.com/restic/restic/internal/debug"
    11  )
    12  
    13  var cleanupHandlers struct {
    14  	sync.Mutex
    15  	list []func() error
    16  	done bool
    17  	ch   chan os.Signal
    18  }
    19  
    20  var stderr = os.Stderr
    21  
    22  func init() {
    23  	cleanupHandlers.ch = make(chan os.Signal)
    24  	go CleanupHandler(cleanupHandlers.ch)
    25  	InstallSignalHandler()
    26  }
    27  
    28  // InstallSignalHandler listens for SIGINT and SIGPIPE, and triggers the cleanup handlers.
    29  func InstallSignalHandler() {
    30  	signal.Notify(cleanupHandlers.ch, syscall.SIGINT)
    31  	signal.Notify(cleanupHandlers.ch, syscall.SIGPIPE)
    32  }
    33  
    34  // SuspendSignalHandler removes the signal handler for SIGINT and SIGPIPE.
    35  func SuspendSignalHandler() {
    36  	signal.Reset(syscall.SIGINT)
    37  	signal.Reset(syscall.SIGPIPE)
    38  }
    39  
    40  // AddCleanupHandler adds the function f to the list of cleanup handlers so
    41  // that it is executed when all the cleanup handlers are run, e.g. when SIGINT
    42  // is received.
    43  func AddCleanupHandler(f func() error) {
    44  	cleanupHandlers.Lock()
    45  	defer cleanupHandlers.Unlock()
    46  
    47  	// reset the done flag for integration tests
    48  	cleanupHandlers.done = false
    49  
    50  	cleanupHandlers.list = append(cleanupHandlers.list, f)
    51  }
    52  
    53  // RunCleanupHandlers runs all registered cleanup handlers
    54  func RunCleanupHandlers() {
    55  	cleanupHandlers.Lock()
    56  	defer cleanupHandlers.Unlock()
    57  
    58  	if cleanupHandlers.done {
    59  		return
    60  	}
    61  	cleanupHandlers.done = true
    62  
    63  	for _, f := range cleanupHandlers.list {
    64  		err := f()
    65  		if err != nil {
    66  			fmt.Fprintf(stderr, "error in cleanup handler: %v\n", err)
    67  		}
    68  	}
    69  	cleanupHandlers.list = nil
    70  }
    71  
    72  // CleanupHandler handles the SIGINT and SIGPIPE signals.
    73  func CleanupHandler(c <-chan os.Signal) {
    74  	for s := range c {
    75  		debug.Log("signal %v received, cleaning up", s)
    76  		fmt.Fprintf(stderr, "%ssignal %v received, cleaning up\n", ClearLine(), s)
    77  
    78  		code := 0
    79  		if s != syscall.SIGINT {
    80  			code = 1
    81  		}
    82  
    83  		Exit(code)
    84  	}
    85  }
    86  
    87  // Exit runs the cleanup handlers and then terminates the process with the
    88  // given exit code.
    89  func Exit(code int) {
    90  	RunCleanupHandlers()
    91  	os.Exit(code)
    92  }