├── go.mod ├── examples ├── panic │ └── main.go ├── crash │ └── main.go └── simple │ └── main.go ├── go.sum ├── LICENSE ├── README.md ├── atexit_test.go └── atexit.go /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/patarapolw/atexit 2 | 3 | go 1.17 4 | 5 | require github.com/stretchr/testify v1.7.0 6 | 7 | require ( 8 | github.com/davecgh/go-spew v1.1.0 // indirect 9 | github.com/pmezard/go-difflib v1.0.0 // indirect 10 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect 11 | ) 12 | -------------------------------------------------------------------------------- /examples/panic/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "time" 6 | 7 | "github.com/patarapolw/atexit" 8 | ) 9 | 10 | func handler() { 11 | fmt.Println("atexit triggered") 12 | } 13 | 14 | func main() { 15 | atexit.Listen() 16 | defer atexit.ListenPanic() 17 | 18 | atexit.Register(handler) 19 | 20 | go func() { 21 | defer atexit.ListenPanic() 22 | 23 | time.Sleep(1 * time.Second) 24 | panic("panic") 25 | }() 26 | 27 | time.Sleep(1 * time.Minute) 28 | atexit.Exit(0) 29 | } 30 | -------------------------------------------------------------------------------- /examples/crash/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "time" 6 | 7 | "github.com/patarapolw/atexit" 8 | ) 9 | 10 | func handler() { 11 | fmt.Println("atexit triggered") 12 | } 13 | 14 | func main() { 15 | atexit.Listen() 16 | defer atexit.ListenPanic() 17 | 18 | atexit.Register(handler) 19 | 20 | go func() { 21 | defer atexit.ListenPanic() 22 | 23 | time.Sleep(1 * time.Second) 24 | s := []int{} 25 | s[0] = s[0] + 1 26 | }() 27 | 28 | time.Sleep(1 * time.Minute) 29 | atexit.Exit(0) 30 | } 31 | -------------------------------------------------------------------------------- /examples/simple/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "time" 6 | 7 | "github.com/patarapolw/atexit" 8 | ) 9 | 10 | func handler() { 11 | fmt.Println("atexit triggered") 12 | } 13 | 14 | func main() { 15 | atexit.Listen() // Await for SIGINT, SIGTERM, whatever. Also works in Windows 16 | defer atexit.ListenPanic() // Listen for panic and crashes 17 | 18 | atexit.Register(handler) 19 | time.Sleep(1 * time.Minute) 20 | 21 | atexit.Exit(0) // This also needs to be called at the end of main function, if you want atexit to be executed on normal exit. 22 | } 23 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= 2 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 3 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 4 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 5 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 6 | github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= 7 | github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 8 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 9 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 10 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= 11 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 12 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Pacharapol Withayasakpunt 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # atexit 2 | 3 | Simple `atexit` implementation for [Go](https://golang.org). 4 | 5 | Note that you *have* to call `atexit.Exit` and not `os.Exit` to terminate your 6 | program (that is, if you want the `atexit` handlers to execute). 7 | 8 | ## Example usage 9 | 10 | ```go 11 | package main 12 | 13 | import ( 14 | "fmt" 15 | "time" 16 | 17 | "github.com/patarapolw/atexit" 18 | ) 19 | 20 | func handler() { 21 | fmt.Println("atexit triggered") 22 | } 23 | 24 | func main() { 25 | atexit.Listen() // Await for SIGINT, SIGTERM, whatever. Also works in Windows 26 | defer atexit.ListenPanic() // Listen for panic and crashes 27 | 28 | atexit.Register(handler) 29 | 30 | go func() { 31 | defer atexit.ListenPanic() 32 | 33 | time.Sleep(1 * time.Second) 34 | panic("panic") 35 | }() 36 | time.Sleep(1 * time.Minute) 37 | 38 | atexit.Exit(0) // This also needs to be called at the end of main function, if you want atexit to be executed on normal exit. 39 | } 40 | ``` 41 | 42 | ## Caveats 43 | 44 | - `os.Exit` doesn't call atexit. Use `atexit.Exit` instead. 45 | - `atexit.Exit(0)` needs to be called at the end of main function, if you want atexit to be executed on normal exit. 46 | - `log.Fatal*` also don't call atexit. Use `atexit.Fatal*` instead. 47 | - `panic` can be listened, but do call `defer atexit.ListenPanic()` at the beginning of main function. 48 | - If `panic` is inside Goroutine, call `defer atexit.ListenPanic()` at the beginning of main function, too. 49 | 50 | ## Install 51 | 52 | go get github.com/patarapolw/atexit 53 | -------------------------------------------------------------------------------- /atexit_test.go: -------------------------------------------------------------------------------- 1 | package atexit 2 | 3 | import ( 4 | "fmt" 5 | "io/ioutil" 6 | "os/exec" 7 | "path" 8 | "path/filepath" 9 | "testing" 10 | "time" 11 | 12 | "github.com/stretchr/testify/require" 13 | ) 14 | 15 | func TestRegister(t *testing.T) { 16 | require := require.New(t) 17 | current := len(handlers) 18 | Register(func() {}) 19 | require.Equal(current+1, len(handlers), "register") 20 | } 21 | 22 | func TestCancel(t *testing.T) { 23 | require := require.New(t) 24 | 25 | id := Register(func() {}) 26 | require.NoError(id.Cancel()) 27 | _, ok := handlers[id] 28 | require.False(ok, "cancel") 29 | } 30 | 31 | func TestHandler(t *testing.T) { 32 | require := require.New(t) 33 | root, err := ioutil.TempDir("", "atexit-test") 34 | require.NoError(err, "TempDir") 35 | 36 | progFile := path.Join(root, "main.go") 37 | err = ioutil.WriteFile(progFile, []byte(testprog), 0666) 38 | require.NoError(err, "prog") 39 | 40 | here, err := filepath.Abs(".") 41 | require.NoError(err, "abs .") 42 | 43 | mod := fmt.Sprintf(modTmpl, here) 44 | modFile := path.Join(root, "go.mod") 45 | err = ioutil.WriteFile(modFile, []byte(mod), 0666) 46 | require.NoError(err, "mod") 47 | outFile := path.Join(root, "main.out") 48 | 49 | arg := time.Now().UTC().String() 50 | err = exec.Command("go", "run", progFile, outFile, arg).Run() 51 | require.Error(err, "run") 52 | 53 | data, err := ioutil.ReadFile(outFile) 54 | require.NoError(err, "read out") 55 | require.Equal(arg, string(data), "output") 56 | } 57 | 58 | var ( 59 | testprog = ` 60 | // Test program for atexit, gets output file and data as arguments and writes 61 | // data to output file in atexit handler. 62 | package main 63 | 64 | import ( 65 | "flag" 66 | "fmt" 67 | "io/ioutil" 68 | 69 | "github.com/patarapolw/atexit" 70 | ) 71 | 72 | var outfile = "" 73 | var data = "" 74 | 75 | func handler() { 76 | ioutil.WriteFile(outfile, []byte(data), 0666) 77 | } 78 | 79 | func badHandler() { 80 | n := 0 81 | fmt.Println(1/n) 82 | } 83 | 84 | func unusedHandler() { 85 | ioutil.WriteFile(outfile, []byte("\nunused"), 0666) 86 | } 87 | 88 | func main() { 89 | flag.Parse() 90 | outfile = flag.Arg(0) 91 | data = flag.Arg(1) 92 | 93 | atexit.Register(handler) 94 | id := atexit.Register(unusedHandler) 95 | atexit.Register(badHandler) 96 | id.Cancel() 97 | atexit.Exit(1) 98 | } 99 | ` 100 | modTmpl = ` 101 | module testexit 102 | 103 | go 1.17 104 | 105 | replace github.com/patarapolw/atexit => %s 106 | ` 107 | ) 108 | -------------------------------------------------------------------------------- /atexit.go: -------------------------------------------------------------------------------- 1 | /*Package atexit lets you define handlers when the program exits. 2 | 3 | Add handlers using Register. 4 | You must call atexit.Exit to get the handler invoked (and then terminate the program). 5 | 6 | This package also provides replacements to log.Fatal, log.Fatalf and log.Fatalln. 7 | 8 | Example: 9 | 10 | package main 11 | 12 | import ( 13 | "fmt" 14 | "time" 15 | 16 | "github.com/patarapolw/atexit" 17 | ) 18 | 19 | func handler() { 20 | fmt.Println("atexit triggered") 21 | } 22 | 23 | func main() { 24 | atexit.Listen() // Await for SIGINT, SIGTERM, whatever. Also works in Windows 25 | defer atexit.ListenPanic() // Listen for panic and crashes 26 | 27 | atexit.Register(handler) 28 | time.Sleep(1 * time.Minute) 29 | 30 | atexit.Exit(0) // This also needs to be called at the end of main function, if you want atexit to be executed on normal exit. 31 | } 32 | */ 33 | package atexit 34 | 35 | import ( 36 | "fmt" 37 | "log" 38 | "os" 39 | "os/signal" 40 | "runtime/debug" 41 | "sync" 42 | "syscall" 43 | ) 44 | 45 | const ( 46 | // Version is package version 47 | Version = "0.4.0" 48 | ) 49 | 50 | var ( 51 | handlers = make(map[HandlerID]func()) 52 | nextHandlerID uint 53 | handlersLock sync.RWMutex // protects the above two 54 | 55 | once sync.Once 56 | ) 57 | 58 | type HandlerID uint 59 | 60 | // Cancel cancels the handler associated with id 61 | func (id HandlerID) Cancel() error { 62 | handlersLock.Lock() 63 | defer handlersLock.Unlock() 64 | 65 | _, ok := handlers[id] 66 | if !ok { 67 | return fmt.Errorf("handler %d not found", id) 68 | } 69 | 70 | delete(handlers, id) 71 | return nil 72 | } 73 | 74 | // Exit runs all the atexit handlers and then terminates the program using 75 | // os.Exit(code) 76 | func Exit(code int) { 77 | RunHandlers() 78 | os.Exit(code) 79 | } 80 | 81 | // Fatal runs all the atexit handler then calls log.Fatal (which will terminate 82 | // the program) 83 | func Fatal(v ...interface{}) { 84 | RunHandlers() 85 | log.Fatal(v...) 86 | } 87 | 88 | // Fatalf runs all the atexit handler then calls log.Fatalf (which will 89 | // terminate the program) 90 | func Fatalf(format string, v ...interface{}) { 91 | RunHandlers() 92 | log.Fatalf(format, v...) 93 | } 94 | 95 | // Fatalln runs all the atexit handler then calls log.Fatalln (which will 96 | // terminate the program) 97 | func Fatalln(v ...interface{}) { 98 | RunHandlers() 99 | log.Fatalln(v...) 100 | } 101 | 102 | // Register adds a handler, call atexit.Exit to invoke all handlers. 103 | func Register(handler func()) HandlerID { 104 | handlersLock.Lock() 105 | defer handlersLock.Unlock() 106 | 107 | nextHandlerID++ 108 | id := HandlerID(nextHandlerID) 109 | handlers[id] = handler 110 | return id 111 | } 112 | 113 | func runHandler(handler func()) { 114 | defer func() { 115 | if err := recover(); err != nil { 116 | fmt.Fprintln(os.Stderr, "error: atexit handler error:", err) 117 | } 118 | }() 119 | 120 | handler() 121 | } 122 | 123 | func executeHandlers() { 124 | handlersLock.RLock() 125 | defer handlersLock.RUnlock() 126 | for _, handler := range handlers { 127 | runHandler(handler) 128 | } 129 | } 130 | 131 | func RunHandlers() { 132 | once.Do(executeHandlers) 133 | } 134 | 135 | func Listen(signals ...os.Signal) chan os.Signal { 136 | c := make(chan os.Signal, 1) 137 | 138 | if len(signals) == 0 { 139 | signal.Notify(c, // https://www.gnu.org/software/libc/manual/html_node/Termination-Signals.html 140 | syscall.SIGTERM, // "the normal way to politely ask a program to terminate" 141 | syscall.SIGINT, // Ctrl+C 142 | syscall.SIGQUIT, // Ctrl-\ 143 | syscall.SIGHUP, // "terminal is disconnected" 144 | syscall.SIGKILL, // "always fatal", "SIGKILL and SIGSTOP may not be caught by a program" 145 | ) 146 | } else { 147 | signal.Notify(c, signals...) 148 | } 149 | 150 | go func() { 151 | s := <-c 152 | if s.String() == "SIGTERM" { 153 | Exit(0) 154 | } 155 | Exit(1) 156 | }() 157 | 158 | return c 159 | } 160 | 161 | func ListenPanic() { 162 | if err := recover(); err != nil { 163 | Fatal(err, "\n", string(debug.Stack())) 164 | } 165 | } 166 | --------------------------------------------------------------------------------