The signal.Notify function support stop receiving signals (#155)

no-src · Apr 25, 2023 · 90b0bf4 · 90b0bf4
1 parent c37b122
commit 90b0bf4
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 10 deletions.
diff --git a/cmd/gofs.go b/cmd/gofs.go
@@ -107,14 +107,15 @@ func runWithConfig(c conf.Config, result result.Result) {
 			return
 		}
 
-		ns := signal.Notify(daemon.Shutdown)
+		ns, ss := signal.Notify(daemon.Shutdown)
 		go func() {
 			result.RegisterNotifyHandler(ns)
 		}()
 		result.InitDone()
 		w := wait.NewWaitDone()
 		go daemon.Run(args, c.DaemonPid, c.DaemonDelay.Duration(), c.DaemonMonitorDelay.Duration(), w)
 		err = w.Wait()
+		ss()
 		return
 	}
 
@@ -170,7 +171,7 @@ func runWithConfig(c conf.Config, result result.Result) {
 	// start monitor
 	log.Info("monitor is starting...")
 	defer log.Info("gofs exited")
-	ns := signal.Notify(m.Shutdown)
+	ns, ss := signal.Notify(m.Shutdown)
 	go func() {
 		result.RegisterNotifyHandler(ns)
 	}()
@@ -182,6 +183,7 @@ func runWithConfig(c conf.Config, result result.Result) {
 		return
 	}
 	err = log.ErrorIf(w.Wait(), "monitor running failed")
+	ss()
 }
 
 func parseConfigFile(cp *conf.Config) error {

diff --git a/internal/signal/signal.go b/internal/signal/signal.go
@@ -19,8 +19,11 @@ var (
 // NotifySignal sends a signal with timeout
 type NotifySignal func(s os.Signal, timeout ...time.Duration) error
 
+// StopSignal stop receiving signals
+type StopSignal func()
+
 // Notify receive signal and try to shut down
-func Notify(shutdown func() error) NotifySignal {
+func Notify(shutdown func() error) (NotifySignal, StopSignal) {
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, syscall.SIGHUP, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGABRT, syscall.SIGTERM)
 	go func() {
@@ -42,6 +45,10 @@ func Notify(shutdown func() error) NotifySignal {
 			}
 		}
 	}()
+
+	ss := func() {
+		signal.Stop(c)
+	}
 	return func(s os.Signal, timeout ...time.Duration) error {
 		t := defaultSendSignalTimeout
 		if len(timeout) > 0 {
@@ -55,5 +62,5 @@ func Notify(shutdown func() error) NotifySignal {
 			log.Warn("[timeout] send a signal [%s] by user", s.String())
 			return fmt.Errorf("%w => %s", errSendSignalTimeout, s.String())
 		}
-	}
+	}, ss
 }
diff --git a/internal/signal/signal_test.go b/internal/signal/signal_test.go
@@ -9,7 +9,7 @@ import (
 )
 
 func TestNotify(t *testing.T) {
-	st := Notify(func() error {
+	ns, ss := Notify(func() error {
 		return nil
 	})
 
@@ -25,13 +25,14 @@ func TestNotify(t *testing.T) {
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			st(tc.signal, time.Second)
+			ns(tc.signal, time.Second)
+			ss()
 		})
 	}
 }
 
 func TestNotify_ShutdownError(t *testing.T) {
-	st := Notify(func() error {
+	ns, ss := Notify(func() error {
 		return errors.New("shutdown error mock")
 	})
 
@@ -51,13 +52,14 @@ func TestNotify_ShutdownError(t *testing.T) {
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			st(tc.signal)
+			ns(tc.signal)
+			ss()
 		})
 	}
 }
 
 func TestNotify_IgnoreSignal(t *testing.T) {
-	st := Notify(func() error {
+	ns, ss := Notify(func() error {
 		return nil
 	})
 
@@ -71,7 +73,8 @@ func TestNotify_IgnoreSignal(t *testing.T) {
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			st(tc.signal)
+			ns(tc.signal)
+			ss()
 		})
 	}
 }