main.go raw

   1  // Package interrupt is a library for providing handling for Ctrl-C/Interrupt
   2  // handling and triggering callbacks for such things as closing files, flushing
   3  // buffers, and other elements of graceful shutdowns.
   4  package interrupt
   5  
   6  import (
   7  	"fmt"
   8  	"os"
   9  	"os/signal"
  10  	"runtime"
  11  
  12  	"next.orly.dev/pkg/lol/log"
  13  	"next.orly.dev/pkg/utils/atomic"
  14  	"next.orly.dev/pkg/utils/qu"
  15  )
  16  
  17  // HandlerWithSource is an interrupt handling closure and the source location
  18  // that it was sent from.
  19  type HandlerWithSource struct {
  20  	Source string
  21  	Fn     func()
  22  }
  23  
  24  var (
  25  	// RestartRequested is set true after restart is requested.
  26  	RestartRequested bool // = true
  27  	requested        atomic.Bool
  28  
  29  	// ch is used to receive SIGINT (Ctrl+C) signals.
  30  	ch chan os.Signal
  31  
  32  	// signals is the list of signals that cause the interrupt
  33  	signals = []os.Signal{os.Interrupt}
  34  
  35  	// ShutdownRequestChan is a channel that can receive shutdown requests
  36  	ShutdownRequestChan = qu.T()
  37  
  38  	// addHandlerChan is used to add an interrupt handler to the list of
  39  	// handlers to be invoked on SIGINT (Ctrl+C) signals.
  40  	addHandlerChan = make(chan HandlerWithSource)
  41  
  42  	// HandlersDone is closed after all interrupt handlers run the first time an
  43  	// interrupt is signaled.
  44  	HandlersDone = make(qu.C)
  45  
  46  	interruptCallbacks       []func()
  47  	interruptCallbackSources []string
  48  )
  49  
  50  // Listener listens for interrupt signals, registers interrupt callbacks, and
  51  // responds to custom shutdown signals as required
  52  func Listener() {
  53  	invokeCallbacks := func() {
  54  		// run handlers in LIFO order.
  55  		for i := range interruptCallbacks {
  56  			idx := len(interruptCallbacks) - 1 - i
  57  			log.T.F(
  58  				"running callback %d from %s", idx,
  59  				interruptCallbackSources[idx],
  60  			)
  61  			interruptCallbacks[idx]()
  62  		}
  63  		log.D.Ln("interrupt handlers finished")
  64  		HandlersDone.Q()
  65  		if RestartRequested {
  66  			Restart()
  67  		} else {
  68  			os.Exit(0)
  69  		}
  70  	}
  71  out:
  72  	for {
  73  		select {
  74  		case _ = <-ch:
  75  			fmt.Fprintf(os.Stderr, "\r")
  76  			requested.Store(true)
  77  			invokeCallbacks()
  78  			break out
  79  
  80  		case <-ShutdownRequestChan.Wait():
  81  			log.W.Ln("received shutdown request - shutting down...")
  82  			requested.Store(true)
  83  			invokeCallbacks()
  84  			break out
  85  
  86  		case handler := <-addHandlerChan:
  87  			interruptCallbacks = append(interruptCallbacks, handler.Fn)
  88  			interruptCallbackSources = append(
  89  				interruptCallbackSources,
  90  				handler.Source,
  91  			)
  92  
  93  		case <-HandlersDone.Wait():
  94  			break out
  95  		}
  96  	}
  97  }
  98  
  99  // AddHandler adds a handler to call when a SIGINT (Ctrl+C) is received.
 100  func AddHandler(handler func()) {
 101  	// Create the channel and start the main interrupt handler which invokes all
 102  	// other callbacks and exits if not already done.
 103  	_, loc, line, _ := runtime.Caller(1)
 104  	msg := fmt.Sprintf("%s:%d", loc, line)
 105  	if ch == nil {
 106  		ch = make(chan os.Signal)
 107  		signal.Notify(ch, signals...)
 108  		go Listener()
 109  	}
 110  	addHandlerChan <- HandlerWithSource{
 111  		msg, handler,
 112  	}
 113  }
 114  
 115  // Request programmatically requests a shutdown
 116  func Request() {
 117  	_, f, l, _ := runtime.Caller(1)
 118  	log.D.Ln("interrupt requested", f, l, requested.Load())
 119  	if requested.Load() {
 120  		log.D.Ln("requested again")
 121  		return
 122  	}
 123  	requested.Store(true)
 124  	ShutdownRequestChan.Q()
 125  	var ok bool
 126  	select {
 127  	case _, ok = <-ShutdownRequestChan:
 128  	default:
 129  	}
 130  	if ok {
 131  		close(ShutdownRequestChan)
 132  	}
 133  }
 134  
 135  // GoroutineDump returns a string with the current goroutine dump in order to
 136  // show what's going on in case of timeout.
 137  func GoroutineDump() string {
 138  	buf := make([]byte, 1<<18)
 139  	n := runtime.Stack(buf, true)
 140  	return string(buf[:n])
 141  }
 142  
 143  // RequestRestart sets the reset flag and requests a restart
 144  func RequestRestart() {
 145  	RestartRequested = true
 146  	log.D.Ln("requesting restart")
 147  	Request()
 148  }
 149  
 150  // Requested returns true if an interrupt has been requested
 151  func Requested() bool {
 152  	return requested.Load()
 153  }
 154