From a980f1f8d25210e3d069fd3a5104c761f9c8828a Mon Sep 17 00:00:00 2001
From: jar3b <hellotan@live.ru>
Date: Tue, 9 Apr 2019 17:50:55 +0300
Subject: [PATCH] Add blocking mode support

---
 README.md |  2 +-
 waiter.go | 18 +++++++++++++++---
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 3e5e36e..b70478a 100644
--- a/README.md
+++ b/README.md
@@ -10,5 +10,5 @@ waiter.AddCloseHandler(func() {
 		nacl.FinalizeStan()
 	}, false)
 
-waiter.Wait()
+waiter.Wait(true)
 ```
\ No newline at end of file
diff --git a/waiter.go b/waiter.go
index 6215136..d1e5026 100644
--- a/waiter.go
+++ b/waiter.go
@@ -9,6 +9,7 @@ import (
 )
 
 type Waiter struct {
+	blockingMode  bool
 	waitGroup     sync.WaitGroup
 	closeHandlers []*CloseHandler
 }
@@ -53,11 +54,21 @@ func (w *Waiter) Halt(err error) {
 	} else {
 		log.Info("Program was terminated gracefully.")
 	}
+	if !w.blockingMode {
+		if err != nil {
+			os.Exit(1)
+		} else {
+			os.Exit(0)
+		}
+	}
 }
 
-func (w *Waiter) Wait() {
-	log.Info("Waiting...")
-	w.waitGroup.Wait()
+func (w *Waiter) Wait(blockingMode bool) {
+	w.blockingMode = blockingMode
+	if blockingMode {
+		log.Info("Waiting...")
+		w.waitGroup.Wait()
+	}
 }
 
 func (w *Waiter) onSignal(sig os.Signal) {
@@ -67,6 +78,7 @@ func (w *Waiter) onSignal(sig os.Signal) {
 
 func NewWaiter() *Waiter {
 	w := Waiter{
+		true,
 		sync.WaitGroup{},
 		make([]*CloseHandler, 0),
 	}