You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
118 lines
3.3 KiB
118 lines
3.3 KiB
package graceful
|
|
|
|
import (
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"sync"
|
|
)
|
|
|
|
// This is the channel that the connections select on. When it is closed, the
|
|
// connections should gracefully exit.
|
|
var kill = make(chan struct{})
|
|
|
|
// This is the channel that the Wait() function selects on. It should only be
|
|
// closed once all the posthooks have been called.
|
|
var wait = make(chan struct{})
|
|
|
|
// This is the WaitGroup that indicates when all the connections have gracefully
|
|
// shut down.
|
|
var wg sync.WaitGroup
|
|
|
|
// This lock protects the list of pre- and post- hooks below.
|
|
var hookLock sync.Mutex
|
|
var prehooks = make([]func(), 0)
|
|
var posthooks = make([]func(), 0)
|
|
|
|
var sigchan = make(chan os.Signal, 1)
|
|
|
|
func init() {
|
|
AddSignal(os.Interrupt)
|
|
go waitForSignal()
|
|
}
|
|
|
|
// AddSignal adds the given signal to the set of signals that trigger a graceful
|
|
// shutdown. Note that for convenience the default interrupt (SIGINT) handler is
|
|
// installed at package load time, and unless you call ResetSignals() will be
|
|
// listened for in addition to any signals you provide by calling this function.
|
|
func AddSignal(sig ...os.Signal) {
|
|
signal.Notify(sigchan, sig...)
|
|
}
|
|
|
|
// ResetSignals resets the list of signals that trigger a graceful shutdown.
|
|
// Useful if, for instance, you don't want to use the default interrupt (SIGINT)
|
|
// handler. Since we necessarily install the SIGINT handler before you have a
|
|
// chance to call ResetSignals(), there will be a brief window during which the
|
|
// set of signals this package listens for will not be as you intend. Therefore,
|
|
// if you intend on using this function, we encourage you to call it as soon as
|
|
// possible.
|
|
func ResetSignals() {
|
|
signal.Stop(sigchan)
|
|
}
|
|
|
|
type userShutdown struct{}
|
|
|
|
func (u userShutdown) String() string {
|
|
return "application initiated shutdown"
|
|
}
|
|
func (u userShutdown) Signal() {}
|
|
|
|
// Shutdown manually trigger a shutdown from your application. Like Wait(),
|
|
// blocks until all connections have gracefully shut down.
|
|
func Shutdown() {
|
|
sigchan <- userShutdown{}
|
|
<-wait
|
|
}
|
|
|
|
// PreHook registers a function to be called before any of this package's normal
|
|
// shutdown actions. All listeners will be called in the order they were added,
|
|
// from a single goroutine.
|
|
func PreHook(f func()) {
|
|
hookLock.Lock()
|
|
defer hookLock.Unlock()
|
|
|
|
prehooks = append(prehooks, f)
|
|
}
|
|
|
|
// PostHook registers a function to be called after all of this package's normal
|
|
// shutdown actions. All listeners will be called in the order they were added,
|
|
// from a single goroutine, and are guaranteed to be called after all listening
|
|
// connections have been closed, but before Wait() returns.
|
|
//
|
|
// If you've Hijack()ed any connections that must be gracefully shut down in
|
|
// some other way (since this library disowns all hijacked connections), it's
|
|
// reasonable to use a PostHook() to signal and wait for them.
|
|
func PostHook(f func()) {
|
|
hookLock.Lock()
|
|
defer hookLock.Unlock()
|
|
|
|
posthooks = append(posthooks, f)
|
|
}
|
|
|
|
func waitForSignal() {
|
|
sig := <-sigchan
|
|
log.Printf("Received %v, gracefully shutting down!", sig)
|
|
|
|
hookLock.Lock()
|
|
defer hookLock.Unlock()
|
|
|
|
for _, f := range prehooks {
|
|
f()
|
|
}
|
|
|
|
close(kill)
|
|
wg.Wait()
|
|
|
|
for _, f := range posthooks {
|
|
f()
|
|
}
|
|
|
|
close(wait)
|
|
}
|
|
|
|
// Wait for all connections to gracefully shut down. This is commonly called at
|
|
// the bottom of the main() function to prevent the program from exiting
|
|
// prematurely.
|
|
func Wait() {
|
|
<-wait
|
|
}
|