6 Commits

3 changed files with 25 additions and 7 deletions

View File

@@ -10,5 +10,6 @@ waiter.AddCloseHandler(func() {
nacl.FinalizeStan() nacl.FinalizeStan()
}, false) }, false)
// blocking wait, if no need to block (with http server, for example), you can omit .Wait() call
waiter.Wait() waiter.Wait()
``` ```

View File

@@ -1,8 +1,11 @@
package grawt package grawt
import "sync"
type CloseHandler struct { type CloseHandler struct {
sync.Mutex
waiter *Waiter waiter *Waiter
Quit chan bool Quit chan struct{}
active bool active bool
autoDone bool autoDone bool
handlerFunc *func() handlerFunc *func()

View File

@@ -9,17 +9,18 @@ import (
) )
type Waiter struct { type Waiter struct {
blockingMode bool
waitGroup sync.WaitGroup waitGroup sync.WaitGroup
closeHandlers []*CloseHandler closeHandlers []*CloseHandler
} }
func (w *Waiter) addHandler(f *func(), autoDone bool) *CloseHandler { func (w *Waiter) addHandler(f *func(), autoDone bool) *CloseHandler {
ch := CloseHandler{ ch := CloseHandler{
w, waiter: w,
make(chan bool, 1), Quit: make(chan struct{}, 1),
true, active: true,
autoDone, autoDone: autoDone,
f, handlerFunc: f,
} }
w.waitGroup.Add(1) w.waitGroup.Add(1)
w.closeHandlers = append(w.closeHandlers, &ch) w.closeHandlers = append(w.closeHandlers, &ch)
@@ -31,11 +32,15 @@ func (w *Waiter) terminateHandler(h *CloseHandler, forceWaitGroupDone bool) {
if h.handlerFunc != nil && *h.handlerFunc != nil { if h.handlerFunc != nil && *h.handlerFunc != nil {
(*h.handlerFunc)() (*h.handlerFunc)()
} }
h.Quit <- true h.Lock()
if h.active {
close(h.Quit)
}
if h.autoDone || forceWaitGroupDone { if h.autoDone || forceWaitGroupDone {
w.waitGroup.Done() w.waitGroup.Done()
} }
h.active = false h.active = false
h.Unlock()
} }
func (w *Waiter) AddCloseHandler(f func(), waitForChannel bool) *CloseHandler { func (w *Waiter) AddCloseHandler(f func(), waitForChannel bool) *CloseHandler {
@@ -53,9 +58,17 @@ func (w *Waiter) Halt(err error) {
} else { } else {
log.Info("Program was terminated gracefully.") log.Info("Program was terminated gracefully.")
} }
if !w.blockingMode {
if err != nil {
os.Exit(1)
} else {
os.Exit(0)
}
}
} }
func (w *Waiter) Wait() { func (w *Waiter) Wait() {
w.blockingMode = true
log.Info("Waiting...") log.Info("Waiting...")
w.waitGroup.Wait() w.waitGroup.Wait()
} }
@@ -67,6 +80,7 @@ func (w *Waiter) onSignal(sig os.Signal) {
func NewWaiter() *Waiter { func NewWaiter() *Waiter {
w := Waiter{ w := Waiter{
false,
sync.WaitGroup{}, sync.WaitGroup{},
make([]*CloseHandler, 0), make([]*CloseHandler, 0),
} }