hooks.go

     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
View as plain text