chore: refactor shutdown config, to prevent setting zero values (#4533)

Signed-off-by: kevin <wanjunfeng@gmail.com>

core/proc/shutdown.go

@@ -13,19 +13,26 @@ import (
 	"github.com/zeromicro/go-zero/core/threading"
 )
 
-type ProcConf struct {
-	WrapUpTime time.Duration `json:",default=1s"`
-	WaitTime   time.Duration `json:",default=5.5s"`
-}
+const (
+	defaultWrapUpTime = time.Second
+	// why we use 5500 milliseconds is because most of our queue are blocking mode with 5 seconds
+	defaultWaitTime = 5500 * time.Millisecond
+)
 
 var (
 	wrapUpListeners   = new(listenerManager)
 	shutdownListeners = new(listenerManager)
-	wrapUpTime        = time.Second
-	// why we use 5500 milliseconds is because most of our queue are blocking mode with 5 seconds
-	delayTimeBeforeForceQuit = 5500 * time.Millisecond
+	wrapUpTime        = defaultWrapUpTime
+	waitTime          = defaultWaitTime
+	shutdownLock      sync.Mutex
 )
 
+// ShutdownConf defines the shutdown configuration for the process.
+type ShutdownConf struct {
+	WrapUpTime time.Duration `json:",default=1s"`
+	WaitTime   time.Duration `json:",default=5.5s"`
+}
+
 // AddShutdownListener adds fn as a shutdown listener.
 // The returned func can be used to wait for fn getting called.
 func AddShutdownListener(fn func()) (waitForCalled func()) {
@@ -40,12 +47,21 @@ func AddWrapUpListener(fn func()) (waitForCalled func()) {
 
 // SetTimeToForceQuit sets the waiting time before force quitting.
 func SetTimeToForceQuit(duration time.Duration) {
-	delayTimeBeforeForceQuit = duration
+	shutdownLock.Lock()
+	defer shutdownLock.Unlock()
+	waitTime = duration
 }
 
-func Setup(conf ProcConf) {
-	wrapUpTime = conf.WrapUpTime
-	delayTimeBeforeForceQuit = conf.WaitTime
+func Setup(conf ShutdownConf) {
+	shutdownLock.Lock()
+	defer shutdownLock.Unlock()
+
+	if conf.WrapUpTime > 0 {
+		wrapUpTime = conf.WrapUpTime
+	}
+	if conf.WaitTime > 0 {
+		waitTime = conf.WaitTime
+	}
 }
 
 // Shutdown calls the registered shutdown listeners, only for test purpose.
@@ -67,8 +83,12 @@ func gracefulStop(signals chan os.Signal, sig syscall.Signal) {
 	time.Sleep(wrapUpTime)
 	go shutdownListeners.notifyListeners()
 
-	time.Sleep(delayTimeBeforeForceQuit - wrapUpTime)
-	logx.Infof("Still alive after %v, going to force kill the process...", delayTimeBeforeForceQuit)
+	shutdownLock.Lock()
+	remainingTime := waitTime - wrapUpTime
+	shutdownLock.Unlock()
+
+	time.Sleep(remainingTime)
+	logx.Infof("Still alive after %v, going to force kill the process...", waitTime)
 	_ = syscall.Kill(syscall.Getpid(), sig)
 }