enhance: gproc.signal enhance #3325

os/gproc/gproc_signal.go

@@ -8,13 +8,12 @@ package gproc
 
 import (
 	"context"
+	"github.com/gogf/gf/v2/internal/intlog"
+	"github.com/gogf/gf/v2/util/gutil"
 	"os"
 	"os/signal"
 	"sync"
 	"syscall"
-
-	"github.com/gogf/gf/v2/internal/intlog"
-	"github.com/gogf/gf/v2/util/gutil"
 )
 
 // SigHandler defines a function type for signal handling.
@@ -23,6 +22,8 @@ type SigHandler func(sig os.Signal)
 var (
 	// Use internal variable to guarantee concurrent safety
 	// when multiple Listen happen.
+	listenOnce        = sync.Once{}
+	waitChan          = make(chan struct{})
 	signalChan        = make(chan os.Signal, 1)
 	signalHandlerMu   sync.Mutex
 	signalHandlerMap  = make(map[os.Signal][]SigHandler)
@@ -48,6 +49,7 @@ func AddSigHandler(handler SigHandler, signals ...os.Signal) {
 	for _, sig := range signals {
 		signalHandlerMap[sig] = append(signalHandlerMap[sig], handler)
 	}
+	notifySignals()
 }
 
 // AddSigHandlerShutdown adds custom signal handler for shutdown signals:
@@ -64,17 +66,26 @@ func AddSigHandlerShutdown(handler ...SigHandler) {
 			signalHandlerMap[sig] = append(signalHandlerMap[sig], h)
 		}
 	}
+	notifySignals()
 }
 
 // Listen blocks and does signal listening and handling.
 func Listen() {
+	listenOnce.Do(func() {
+		go listen()
+	})
+
+	<-waitChan
+}
+
+func listen() {
+	defer close(waitChan)
+
 	var (
-		signals = getHandlerSignals()
-		ctx     = context.Background()
-		wg      = sync.WaitGroup{}
-		sig     os.Signal
+		ctx = context.Background()
+		wg  = sync.WaitGroup{}
+		sig os.Signal
 	)
-	signal.Notify(signalChan, signals...)
 	for {
 		sig = <-signalChan
 		intlog.Printf(ctx, `signal received: %s`, sig.String())
@@ -108,14 +119,12 @@ func Listen() {
 	}
 }
 
-func getHandlerSignals() []os.Signal {
-	signalHandlerMu.Lock()
-	defer signalHandlerMu.Unlock()
+func notifySignals() {
 	var signals = make([]os.Signal, 0)
 	for s := range signalHandlerMap {
 		signals = append(signals, s)
 	}
-	return signals
+	signal.Notify(signalChan, signals...)
 }
 
 func getHandlersBySignal(sig os.Signal) []SigHandler {

os/gproc/gproc_z_signal_test.go

@@ -0,0 +1,103 @@
+// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
+//
+// This Source Code Form is subject to the terms of the MIT License.
+// If a copy of the MIT was not distributed with this file,
+// You can obtain one at https://github.com/gogf/gf.
+
+package gproc
+
+import (
+	"github.com/gogf/gf/v2/test/gtest"
+	"os"
+	"syscall"
+	"testing"
+	"time"
+)
+
+func Test_Signal(t *testing.T) {
+	go Listen()
+
+	// non shutdown signal
+	gtest.C(t, func(t *gtest.T) {
+		sigRec := make(chan os.Signal, 1)
+		AddSigHandler(func(sig os.Signal) {
+			sigRec <- sig
+		}, syscall.SIGUSR1, syscall.SIGUSR2)
+
+		sendSignal(syscall.SIGUSR1)
+		select {
+		case s := <-sigRec:
+			t.AssertEQ(s, syscall.SIGUSR1)
+			t.AssertEQ(false, isWaitChClosed())
+		case <-time.After(time.Second):
+			t.Error("signal SIGUSR1 handler timeout")
+		}
+
+		sendSignal(syscall.SIGUSR2)
+		select {
+		case s := <-sigRec:
+			t.AssertEQ(s, syscall.SIGUSR2)
+			t.AssertEQ(false, isWaitChClosed())
+		case <-time.After(time.Second):
+			t.Error("signal SIGUSR2 handler timeout")
+		}
+
+		sendSignal(syscall.SIGHUP)
+		select {
+		case <-sigRec:
+			t.Error("signal SIGHUP should not be listen")
+		case <-time.After(time.Millisecond * 100):
+		}
+
+		// multiple listen
+		go Listen()
+		go Listen()
+		sendSignal(syscall.SIGUSR1)
+		cnt := 0
+		timeout := time.After(time.Second)
+		for {
+			select {
+			case <-sigRec:
+				cnt++
+			case <-timeout:
+				if cnt == 0 {
+					t.Error("signal SIGUSR2 handler timeout")
+				}
+				if cnt != 1 {
+					t.Error("multi Listen() repetitive execution")
+				}
+				return
+			}
+		}
+	})
+
+	// test shutdown signal
+	gtest.C(t, func(t *gtest.T) {
+		sigRec := make(chan os.Signal, 1)
+		AddSigHandlerShutdown(func(sig os.Signal) {
+			sigRec <- sig
+		})
+
+		sendSignal(syscall.SIGTERM)
+		select {
+		case s := <-sigRec:
+			t.AssertEQ(s, syscall.SIGTERM)
+			t.AssertEQ(true, isWaitChClosed())
+		case <-time.After(time.Second):
+			t.Error("signal SIGUSR2 handler timeout")
+		}
+	})
+}
+
+func sendSignal(sig os.Signal) {
+	signalChan <- sig
+}
+
+func isWaitChClosed() bool {
+	select {
+	case _, ok := <-waitChan:
+		return !ok
+	default:
+		return false
+	}
+}