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
// at exit.
//
// cancelFunc will only be run once.
func OnError(perr *error, fn func()) func() {
handle := Register(fn)
var once sync.Once
onceFn := func() {
once.Do(fn)
}
handle := Register(onceFn)
return func() {
defer Unregister(handle)
if *perr != nil {
fn()
onceFn()
}
}