From 9052135297f88ab2869f6a5422acd7ded800bc61 Mon Sep 17 00:00:00 2001
From: whiteCcinn <471113744@qq.com>
Date: Tue, 25 May 2021 18:19:26 +0800
Subject: [PATCH] Optimize semaphore logic

---
 daemon.go                | 35 +++++++++++++++++----------
 example/daemon_signal.go | 27 ++++++++++++---------
 signal.go                | 52 +++++++++++++++-------------------------
 3 files changed, 57 insertions(+), 57 deletions(-)

diff --git a/daemon.go b/daemon.go
index 16b7562..356c31f 100644
--- a/daemon.go
+++ b/daemon.go
@@ -10,7 +10,9 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"os/signal"
 	"strconv"
+	"strings"
 	"sync"
 	"syscall"
 	"time"
@@ -237,21 +239,27 @@ func (dctx *Context) Run(ctx context.Context) error {
 
 		// child process
 		if cmd == nil {
-			exitFunc := func(sig os.Signal) (err error) {
-				// this is fd(3)
-				pipe := os.NewFile(uintptr(3), "pipe")
-				message := PipeMessage{
-					Type:     ProcessToSupervisor,
-					Behavior: WantSafetyClose,
+			go func() {
+				ch := make(chan os.Signal, 2)
+				exitFunc := func(sig os.Signal) {
+					// this is fd(3)
+					pipe := os.NewFile(uintptr(3), "pipe")
+					message := PipeMessage{
+						Type:     ProcessToSupervisor,
+						Behavior: WantSafetyClose,
+					}
+					err = json.NewEncoder(pipe).Encode(message)
+					if err != nil && !strings.Contains(err.Error(), "broken pipe") {
+						panic(err)
+					}
+					return
 				}
-				err = json.NewEncoder(pipe).Encode(message)
-				if err != nil {
-					panic(err)
+				signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+
+				for sig := range ch {
+					exitFunc(sig)
 				}
-				return
-			}
-			SetSigHandler(exitFunc, syscall.SIGINT, syscall.SIGTERM)
-			go ServeSignals()
+			}()
 			break
 		}
 
@@ -296,6 +304,7 @@ func (dctx *Context) Run(ctx context.Context) error {
 			dctx.log("[supervisor(%d)] [named-pipe-ipc] [listen]\n", dctx.Pid)
 			SetSigHandler(func(sig os.Signal) (err error) {
 				dctx.clean()
+				os.Exit(0)
 				return
 			}, syscall.SIGINT, syscall.SIGTERM)
 			go ServeSignals()
diff --git a/example/daemon_signal.go b/example/daemon_signal.go
index 171e3eb..26d9857 100644
--- a/example/daemon_signal.go
+++ b/example/daemon_signal.go
@@ -5,8 +5,8 @@ import (
 	"github.com/whiteCcinn/daemon"
 	"log"
 	"os"
+	"sync"
 	"syscall"
-	"time"
 )
 
 func main() {
@@ -34,24 +34,29 @@ func main() {
 
 	// belong func main()
 	dctx.WithRecovery(func() {
-		daemon.SetSigHandler(func(sig os.Signal) (err error) {
-			log.Println("sigint")
-			return
-		}, syscall.SIGINT)
+		wg := new(sync.WaitGroup)
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			daemon.SetSigHandler(func(sig os.Signal) (err error) {
+				log.Println("sigint")
+				return daemon.ErrStop
+			}, syscall.SIGINT)
 
-		daemon.SetSigHandler(func(sig os.Signal) (err error) {
-			log.Println("sigterm")
-			return nil
-		}, syscall.SIGTERM)
+			daemon.SetSigHandler(func(sig os.Signal) (err error) {
+				log.Println("sigterm")
+				return daemon.ErrStop
+			}, syscall.SIGTERM)
 
-		go func() {
 			err := daemon.ServeSignals()
 			if err != nil {
 				log.Println(err)
 			}
 		}()
+
 		log.Println(os.Getpid(), "start...")
-		time.Sleep(time.Second * 10)
+		//time.Sleep(time.Second * 10)
 		log.Println(os.Getpid(), "end...")
+		wg.Wait()
 	}, nil)
 }
diff --git a/signal.go b/signal.go
index 04e3651..cbfe6c3 100644
--- a/signal.go
+++ b/signal.go
@@ -4,13 +4,9 @@ import (
 	"errors"
 	"os"
 	"os/signal"
-	"sync"
 	"syscall"
 )
 
-var signalOnce = new(sync.Once)
-var ch = make(chan os.Signal, 8)
-
 // ErrStop should be returned signal handler function
 // for termination of handling signals.
 var ErrStop = errors.New("stop serve signals")
@@ -22,52 +18,42 @@ type SignalHandlerFunc func(sig os.Signal) (err error)
 // SIGTERM has the default handler, he returns ErrStop.
 func SetSigHandler(handler SignalHandlerFunc, signals ...os.Signal) {
 	for _, sig := range signals {
-		handlers[sig] = append(handlers[sig], handler)
+		handlers[sig] = handler
 	}
 }
 
 // ServeSignals calls handlers for system signals.
 func ServeSignals() (err error) {
-	done := false
-	signalOnce.Do(func() {
-		signals := make([]os.Signal, 0, len(handlers))
-		for sig := range handlers {
-			signals = append(signals, sig)
-		}
+	signals := make([]os.Signal, 0, len(handlers))
+	for sig := range handlers {
+		signals = append(signals, sig)
+	}
 
-		signal.Notify(ch, signals...)
-	loop:
-		for sig := range ch {
-			for _, f := range handlers[sig] {
-				err = f(sig)
-				if err == ErrStop {
-					break loop
-				}
-				if err != nil {
-					break
-				}
-			}
+	ch := make(chan os.Signal, 8)
+	signal.Notify(ch, signals...)
+
+	for sig := range ch {
+		err = handlers[sig](sig)
+		if err != nil {
+			break
 		}
+	}
 
-		signal.Stop(ch)
-		done = true
-	})
+	signal.Stop(ch)
 
-	if done == false {
-		signal.Stop(ch)
-		signalOnce = new(sync.Once)
-		return ServeSignals()
+	if err == ErrStop {
+		err = nil
 	}
 
 	return
 }
 
-var handlers = make(map[os.Signal][]SignalHandlerFunc)
+var handlers = make(map[os.Signal]SignalHandlerFunc)
 
 func init() {
-	handlers[syscall.SIGTERM] = []SignalHandlerFunc{sigtermDefaultHandler}
+	handlers[syscall.SIGTERM] = sigtermDefaultHandler
 }
 
 func sigtermDefaultHandler(sig os.Signal) error {
 	return ErrStop
-}
+}
\ No newline at end of file