graceful: add {Pre,Post}HookWithSignal functions
Carl Jackson
5 years ago
| 11 | 11 | |
| 12 | 12 | var mu sync.Mutex // protects everything that follows |
| 13 | 13 | var listeners = make([]*listener.T, 0) |
| 14 | var prehooks = make([]func(), 0) | |
| 15 | var posthooks = make([]func(), 0) | |
| 14 | var prehooks = make([]func(os.Signal), 0) | |
| 15 | var posthooks = make([]func(os.Signal), 0) | |
| 16 | 16 | var closing int32 |
| 17 | 17 | var doubleKick, timeout time.Duration |
| 18 | 18 | |
| 39 | 39 | signal.Stop(sigchan) |
| 40 | 40 | } |
| 41 | 41 | |
| 42 | // PreHookWithSignal registers a function to be called before any of this | |
| 43 | // package's normal shutdown actions, which recieves the signal that caused the | |
| 44 | // shutdown (or nil for manual shutdowns). All listeners will be called in the | |
| 45 | // order they were added, from a single goroutine. | |
| 46 | func PreHookWithSignal(f func(os.Signal)) { | |
| 47 | mu.Lock() | |
| 48 | defer mu.Unlock() | |
| 49 | ||
| 50 | prehooks = append(prehooks, f) | |
| 51 | } | |
| 52 | ||
| 42 | 53 | // PreHook registers a function to be called before any of this package's normal |
| 43 | 54 | // shutdown actions. All listeners will be called in the order they were added, |
| 44 | 55 | // from a single goroutine. |
| 45 | 56 | func PreHook(f func()) { |
| 46 | mu.Lock() | |
| 47 | defer mu.Unlock() | |
| 48 | ||
| 49 | prehooks = append(prehooks, f) | |
| 57 | PreHookWithSignal(func(_ os.Signal) { | |
| 58 | f() | |
| 59 | }) | |
| 60 | } | |
| 61 | ||
| 62 | // PostHookWithSignal registers a function to be called after all of this | |
| 63 | // package's normal shutdown actions, which receives the signal that caused the | |
| 64 | // shutdown (or nil for manual shutdowns). All listeners will be called in the | |
| 65 | // order they were added, from a single goroutine, and are guaranteed to be | |
| 66 | // called after all listening connections have been closed, but before Wait() | |
| 67 | // returns. | |
| 68 | // | |
| 69 | // If you've Hijacked any connections that must be gracefully shut down in some | |
| 70 | // other way (since this library disowns all hijacked connections), it's | |
| 71 | // reasonable to use a PostHook to signal and wait for them. | |
| 72 | func PostHookWithSignal(f func(os.Signal)) { | |
| 73 | mu.Lock() | |
| 74 | defer mu.Unlock() | |
| 75 | ||
| 76 | posthooks = append(posthooks, f) | |
| 50 | 77 | } |
| 51 | 78 | |
| 52 | 79 | // PostHook registers a function to be called after all of this package's normal |
| 58 | 85 | // other way (since this library disowns all hijacked connections), it's |
| 59 | 86 | // reasonable to use a PostHook to signal and wait for them. |
| 60 | 87 | func PostHook(f func()) { |
| 61 | mu.Lock() | |
| 62 | defer mu.Unlock() | |
| 63 | ||
| 64 | posthooks = append(posthooks, f) | |
| 88 | PostHookWithSignal(func(_ os.Signal) { | |
| 89 | f() | |
| 90 | }) | |
| 65 | 91 | } |
| 66 | 92 | |
| 67 | 93 | // Shutdown manually triggers a shutdown from your application. Like Wait, |
| 68 | 94 | // blocks until all connections have gracefully shut down. |
| 69 | 95 | func Shutdown() { |
| 70 | shutdown(false) | |
| 96 | shutdown(false, nil) | |
| 71 | 97 | } |
| 72 | 98 | |
| 73 | 99 | // ShutdownNow triggers an immediate shutdown from your application. All |
| 74 | 100 | // connections (not just those that are idle) are immediately closed, even if |
| 75 | 101 | // they are in the middle of serving a request. |
| 76 | 102 | func ShutdownNow() { |
| 77 | shutdown(true) | |
| 103 | shutdown(true, nil) | |
| 78 | 104 | } |
| 79 | 105 | |
| 80 | 106 | // DoubleKickWindow sets the length of the window during which two back-to-back |
| 122 | 148 | func sigLoop() { |
| 123 | 149 | var last time.Time |
| 124 | 150 | for { |
| 125 | <-sigchan | |
| 151 | sig := <-sigchan | |
| 126 | 152 | now := time.Now() |
| 127 | 153 | mu.Lock() |
| 128 | 154 | force := doubleKick != 0 && now.Sub(last) < doubleKick |
| 129 | 155 | if t := timeout; t != 0 && !force { |
| 130 | 156 | go func() { |
| 131 | 157 | time.Sleep(t) |
| 132 | shutdown(true) | |
| 158 | shutdown(true, sig) | |
| 133 | 159 | }() |
| 134 | 160 | } |
| 135 | 161 | mu.Unlock() |
| 136 | go shutdown(force) | |
| 162 | go shutdown(force, sig) | |
| 137 | 163 | last = now |
| 138 | 164 | } |
| 139 | 165 | } |
| 140 | 166 | |
| 141 | 167 | var preOnce, closeOnce, forceOnce, postOnce, notifyOnce sync.Once |
| 142 | 168 | |
| 143 | func shutdown(force bool) { | |
| 169 | func shutdown(force bool, sig os.Signal) { | |
| 144 | 170 | preOnce.Do(func() { |
| 145 | 171 | mu.Lock() |
| 146 | 172 | defer mu.Unlock() |
| 147 | 173 | for _, f := range prehooks { |
| 148 | f() | |
| 174 | f(sig) | |
| 149 | 175 | } |
| 150 | 176 | }) |
| 151 | 177 | |
| 163 | 189 | mu.Lock() |
| 164 | 190 | defer mu.Unlock() |
| 165 | 191 | for _, f := range posthooks { |
| 166 | f() | |
| 192 | f(sig) | |
| 167 | 193 | } |
| 168 | 194 | }) |
| 169 | 195 | |