diff --git a/close_handler.go b/close_handler.go index 98f7257..22d4471 100644 --- a/close_handler.go +++ b/close_handler.go @@ -1,11 +1,17 @@ package grawt +import ( + "sync" +) + type CloseHandler struct { waiter *Waiter Quit chan struct{} active bool autoDone bool + wgDone bool handlerFunc *func() + mu *sync.Mutex } func (ch *CloseHandler) Halt(err error) { diff --git a/waiter.go b/waiter.go index 230941b..3772189 100644 --- a/waiter.go +++ b/waiter.go @@ -22,7 +22,9 @@ func (w *Waiter) addHandler(f *func(), autoDone bool) *CloseHandler { Quit: make(chan struct{}, 1), active: true, autoDone: autoDone, + wgDone: false, handlerFunc: f, + mu: &sync.Mutex{}, } w.Lock() @@ -34,18 +36,23 @@ func (w *Waiter) addHandler(f *func(), autoDone bool) *CloseHandler { } func (w *Waiter) terminateHandler(h *CloseHandler, forceWaitGroupDone bool) { + h.mu.Lock() + defer h.mu.Unlock() if !h.active { + if !h.wgDone { + w.waitGroup.Done() + } return } + if h.handlerFunc != nil && *h.handlerFunc != nil { (*h.handlerFunc)() } - if h.active { - close(h.Quit) - } + close(h.Quit) if h.autoDone || forceWaitGroupDone { w.waitGroup.Done() + h.wgDone = true } h.active = false @@ -87,6 +94,7 @@ func (w *Waiter) Wait() { w.blockingMode = true log.Info("Waiting...") w.waitGroup.Wait() + log.Info("Terminated") } func (w *Waiter) onSignal(sig os.Signal) {