diff --git a/README.md b/README.md index b3aa528..3e5e36e 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,14 @@ # grawt graceful terminations for go applications + +## example usage + +``` +var waiter = grawt.NewWaiter() + +waiter.AddCloseHandler(func() { + nacl.FinalizeStan() + }, false) + +waiter.Wait() +``` \ No newline at end of file diff --git a/close_handler.go b/close_handler.go new file mode 100644 index 0000000..83698ca --- /dev/null +++ b/close_handler.go @@ -0,0 +1,17 @@ +package grawt + +type CloseHandler struct { + waiter *Waiter + Quit chan bool + active bool + autoDone bool + handlerFunc *func() +} + +func (ch *CloseHandler) Halt(err error) { + ch.waiter.Halt(err) +} + +func (ch *CloseHandler) Done() { + ch.waiter.terminateHandler(ch, true) +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..a308930 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module github.com/jar3b/grawt + +require github.com/sirupsen/logrus v1.4.1 diff --git a/waiter.go b/waiter.go new file mode 100644 index 0000000..6215136 --- /dev/null +++ b/waiter.go @@ -0,0 +1,82 @@ +package grawt + +import ( + log "github.com/sirupsen/logrus" + "os" + "os/signal" + "sync" + "syscall" +) + +type Waiter struct { + waitGroup sync.WaitGroup + closeHandlers []*CloseHandler +} + +func (w *Waiter) addHandler(f *func(), autoDone bool) *CloseHandler { + ch := CloseHandler{ + w, + make(chan bool, 1), + true, + autoDone, + f, + } + w.waitGroup.Add(1) + w.closeHandlers = append(w.closeHandlers, &ch) + + return &ch +} + +func (w *Waiter) terminateHandler(h *CloseHandler, forceWaitGroupDone bool) { + if h.handlerFunc != nil && *h.handlerFunc != nil { + (*h.handlerFunc)() + } + h.Quit <- true + if h.autoDone || forceWaitGroupDone { + w.waitGroup.Done() + } + h.active = false +} + +func (w *Waiter) AddCloseHandler(f func(), waitForChannel bool) *CloseHandler { + return w.addHandler(&f, !waitForChannel) +} + +func (w *Waiter) Halt(err error) { + for _, h := range w.closeHandlers { + if h.active { + w.terminateHandler(h, false) + } + } + if err != nil { + log.Errorf("Program was terminated with error: %v", err) + } else { + log.Info("Program was terminated gracefully.") + } +} + +func (w *Waiter) Wait() { + log.Info("Waiting...") + w.waitGroup.Wait() +} + +func (w *Waiter) onSignal(sig os.Signal) { + log.Infof("Received signal '%s'! Exiting...", sig.String()) + w.Halt(nil) +} + +func NewWaiter() *Waiter { + w := Waiter{ + sync.WaitGroup{}, + make([]*CloseHandler, 0), + } + sigs := make(chan os.Signal, 1) + signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) + + go func() { + sig := <-sigs + w.onSignal(sig) + }() + + return &w +}