cleaner.go

 1package interrupt
 2
 3import (
 4	"fmt"
 5	"os"
 6	"os/signal"
 7	"sync"
 8	"syscall"
 9)
10
11// CleanerFunc is a function to be executed when an interrupt trigger
12type CleanerFunc func() error
13
14// CancelFunc, if called, will disable the associated cleaner.
15// This allow to create temporary cleaner. Be mindful though to not
16// create too much of them as they are just disabled, not removed from
17// memory.
18type CancelFunc func()
19
20type wrapper struct {
21	f        CleanerFunc
22	disabled bool
23}
24
25var mu sync.Mutex
26var cleaners []*wrapper
27var handlerCreated = false
28
29// RegisterCleaner is responsible for registering a cleaner function.
30// When a function is registered, the Signal watcher is started in a goroutine.
31func RegisterCleaner(cleaner CleanerFunc) CancelFunc {
32	mu.Lock()
33	defer mu.Unlock()
34
35	w := &wrapper{f: cleaner}
36
37	cancel := func() {
38		mu.Lock()
39		defer mu.Unlock()
40		w.disabled = true
41	}
42
43	// prepend to later execute then in reverse order
44	cleaners = append([]*wrapper{w}, cleaners...)
45
46	if handlerCreated {
47		return cancel
48	}
49
50	handlerCreated = true
51	go func() {
52		ch := make(chan os.Signal, 1)
53		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
54		<-ch
55		// Prevent un-terminated ^C character in terminal
56		fmt.Println()
57		errl := clean()
58		for _, err := range errl {
59			_, _ = fmt.Fprintln(os.Stderr, err)
60		}
61		os.Exit(1)
62	}()
63
64	return cancel
65}
66
67// clean invokes all registered cleanup functions, and returns a list of errors, if they exist.
68func clean() (errorList []error) {
69	mu.Lock()
70	defer mu.Unlock()
71
72	for _, cleaner := range cleaners {
73		if cleaner.disabled {
74			continue
75		}
76		err := cleaner.f()
77		if err != nil {
78			errorList = append(errorList, err)
79		}
80	}
81	cleaners = []*wrapper{}
82	return
83}