1 // Copyright 2024 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 // Package exithook provides limited support for on-exit cleanup. 6 // 7 // CAREFUL! The expectation is that Add should only be called 8 // from a safe context (e.g. not an error/panic path or signal 9 // handler, preemption enabled, allocation allowed, write barriers 10 // allowed, etc), and that the exit function F will be invoked under 11 // similar circumstances. That is the say, we are expecting that F 12 // uses normal / high-level Go code as opposed to one of the more 13 // restricted dialects used for the trickier parts of the runtime. 14 package exithook 15 16 import ( 17 "internal/runtime/atomic" 18 _ "unsafe" // for linkname 19 ) 20 21 // A Hook is a function to be run at program termination 22 // (when someone invokes os.Exit, or when main.main returns). 23 // Hooks are run in reverse order of registration: 24 // the first hook added is the last one run. 25 type Hook struct { 26 F func() // func to run 27 RunOnFailure bool // whether to run on non-zero exit code 28 } 29 30 var ( 31 locked atomic.Int32 32 runGoid atomic.Uint64 33 hooks []Hook 34 running bool 35 36 // runtime sets these for us 37 Gosched func() 38 Goid func() uint64 39 Throw func(string) 40 ) 41 42 // Add adds a new exit hook. 43 func Add(h Hook) { 44 for !locked.CompareAndSwap(0, 1) { 45 Gosched() 46 } 47 hooks = append(hooks, h) 48 locked.Store(0) 49 } 50 51 // Run runs the exit hooks. 52 // 53 // If an exit hook panics, Run will throw with the panic on the stack. 54 // If an exit hook invokes exit in the same goroutine, the goroutine will throw. 55 // If an exit hook invokes exit in another goroutine, that exit will block. 56 func Run(code int) { 57 for !locked.CompareAndSwap(0, 1) { 58 if Goid() == runGoid.Load() { 59 Throw("exit hook invoked exit") 60 } 61 Gosched() 62 } 63 defer locked.Store(0) 64 runGoid.Store(Goid()) 65 defer runGoid.Store(0) 66 67 defer func() { 68 if e := recover(); e != nil { 69 Throw("exit hook invoked panic") 70 } 71 }() 72 73 for len(hooks) > 0 { 74 h := hooks[len(hooks)-1] 75 hooks = hooks[:len(hooks)-1] 76 if code != 0 && !h.RunOnFailure { 77 continue 78 } 79 h.F() 80 } 81 } 82 83 type exitError string 84 85 func (e exitError) Error() string { return string(e) } 86