From 486382b9d62b0aac2eda91fb3bb9dafdd528bc5d Mon Sep 17 00:00:00 2001
From: jar3b <hellotan@live.ru>
Date: Fri, 30 Oct 2020 18:20:27 +0300
Subject: [PATCH] fix: rework sync logic

---
 close_handler.go |  7 -------
 go.mod           |  4 +++-
 go.sum           | 10 ++++++++++
 waiter.go        | 30 ++++++++++++++++++++++++------
 4 files changed, 37 insertions(+), 14 deletions(-)
 create mode 100644 go.sum

diff --git a/close_handler.go b/close_handler.go
index 7a9592f..9249d62 100644
--- a/close_handler.go
+++ b/close_handler.go
@@ -1,9 +1,6 @@
 package grawt
 
-import "sync"
-
 type CloseHandler struct {
-	sync.Mutex
 	waiter      *Waiter
 	Quit        chan struct{}
 	active      bool
@@ -14,7 +11,3 @@ type CloseHandler struct {
 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
index a308930..eca0276 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,5 @@
 module github.com/jar3b/grawt
 
-require github.com/sirupsen/logrus v1.4.1
+go 1.15
+
+require github.com/sirupsen/logrus v1.7.0
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..4d74a1e
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,10 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/sirupsen/logrus v1.7.0 h1:ShrD1U9pZB12TX0cVy0DtePoCH97K8EtX+mg7ZARUtM=
+github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
+github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
+github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
+golang.org/x/sys v0.0.0-20191026070338-33540a1f6037 h1:YyJpGZS1sBuBCzLAR1VEpK193GlqGZbnPFnPV/5Rsb4=
+golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
diff --git a/waiter.go b/waiter.go
index 21f14a4..230941b 100644
--- a/waiter.go
+++ b/waiter.go
@@ -9,9 +9,11 @@ import (
 )
 
 type Waiter struct {
+	sync.RWMutex
 	blockingMode  bool
 	waitGroup     sync.WaitGroup
 	closeHandlers []*CloseHandler
+	isHalting     bool
 }
 
 func (w *Waiter) addHandler(f *func(), autoDone bool) *CloseHandler {
@@ -22,25 +24,31 @@ func (w *Waiter) addHandler(f *func(), autoDone bool) *CloseHandler {
 		autoDone:    autoDone,
 		handlerFunc: f,
 	}
+
+	w.Lock()
 	w.waitGroup.Add(1)
 	w.closeHandlers = append(w.closeHandlers, &ch)
+	w.Unlock()
 
 	return &ch
 }
 
 func (w *Waiter) terminateHandler(h *CloseHandler, forceWaitGroupDone bool) {
+	if !h.active {
+		return
+	}
 	if h.handlerFunc != nil && *h.handlerFunc != nil {
 		(*h.handlerFunc)()
 	}
-	h.Lock()
+
 	if h.active {
 		close(h.Quit)
 	}
 	if h.autoDone || forceWaitGroupDone {
 		w.waitGroup.Done()
 	}
+
 	h.active = false
-	h.Unlock()
 }
 
 func (w *Waiter) AddCloseHandler(f func(), waitForChannel bool) *CloseHandler {
@@ -48,11 +56,17 @@ func (w *Waiter) AddCloseHandler(f func(), waitForChannel bool) *CloseHandler {
 }
 
 func (w *Waiter) Halt(err error) {
-	for _, h := range w.closeHandlers {
-		if h.active {
-			w.terminateHandler(h, false)
-		}
+	if w.isHalting {
+		return
 	}
+	w.isHalting = true
+
+	w.RLock()
+	for _, h := range w.closeHandlers {
+		w.terminateHandler(h, false)
+	}
+	w.RUnlock()
+
 	if err != nil {
 		log.Errorf("Program was terminated with error: %v", err)
 	} else {
@@ -65,6 +79,8 @@ func (w *Waiter) Halt(err error) {
 			os.Exit(0)
 		}
 	}
+
+	w.isHalting = false
 }
 
 func (w *Waiter) Wait() {
@@ -80,9 +96,11 @@ func (w *Waiter) onSignal(sig os.Signal) {
 
 func NewWaiter() *Waiter {
 	w := Waiter{
+		sync.RWMutex{},
 		false,
 		sync.WaitGroup{},
 		make([]*CloseHandler, 0),
+		false,
 	}
 	sigs := make(chan os.Signal, 1)
 	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)