lib/atexit: ensure OnError only calls cancel function once

Before this change the cancelFunc could be called twice, once while
handling the interrupt (CTRL-C) and once while unwinding the stack if
the function happened to finish.

This change ensure the cancelFunc is only called once by wrapping it
in a sync.Once
This commit is contained in:
Nick Craig-Wood 2023-03-22 10:44:47 +00:00
parent ddb3b17e96
commit cdfa0beafb
1 changed files with 8 additions and 2 deletions

View File

@ -115,12 +115,18 @@ func Run() {
// //
// So cancelFunc will be run if the function exits with an error or // So cancelFunc will be run if the function exits with an error or
// at exit. // at exit.
//
// cancelFunc will only be run once.
func OnError(perr *error, fn func()) func() { func OnError(perr *error, fn func()) func() {
handle := Register(fn) var once sync.Once
onceFn := func() {
once.Do(fn)
}
handle := Register(onceFn)
return func() { return func() {
defer Unregister(handle) defer Unregister(handle)
if *perr != nil { if *perr != nil {
fn() onceFn()
} }
} }