diff --git a/core/proc/shutdown.go b/core/proc/shutdown.go index 763742a298df..5f35ee1b1478 100644 --- a/core/proc/shutdown.go +++ b/core/proc/shutdown.go @@ -13,16 +13,17 @@ import ( "github.com/zeromicro/go-zero/core/threading" ) -const ( - wrapUpTime = time.Second - // why we use 5500 milliseconds is because most of our queue are blocking mode with 5 seconds - waitTime = 5500 * time.Millisecond -) +type ProcConf struct { + WrapUpTime time.Duration `json:",default=1s"` + WaitTime time.Duration `json:",default=5.5s"` +} var ( - wrapUpListeners = new(listenerManager) - shutdownListeners = new(listenerManager) - delayTimeBeforeForceQuit = waitTime + 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 ) // AddShutdownListener adds fn as a shutdown listener. @@ -42,6 +43,11 @@ func SetTimeToForceQuit(duration time.Duration) { delayTimeBeforeForceQuit = duration } +func Setup(conf ProcConf) { + wrapUpTime = conf.WrapUpTime + delayTimeBeforeForceQuit = conf.WaitTime +} + // Shutdown calls the registered shutdown listeners, only for test purpose. func Shutdown() { shutdownListeners.notifyListeners() diff --git a/core/proc/shutdown_test.go b/core/proc/shutdown_test.go index 79c7deffa8af..30a805488940 100644 --- a/core/proc/shutdown_test.go +++ b/core/proc/shutdown_test.go @@ -58,3 +58,12 @@ func TestNotifyMoreThanOnce(t *testing.T) { t.Fatal("timeout, check error logs") } } + +func TestSetup(t *testing.T) { + Setup(ProcConf{ + WrapUpTime: time.Second * 2, + WaitTime: time.Second * 30, + }) + assert.Equal(t, time.Second*2, wrapUpTime) + assert.Equal(t, time.Second*30, delayTimeBeforeForceQuit) +} diff --git a/core/service/serviceconf.go b/core/service/serviceconf.go index a4999ab8b075..fc66e91e281b 100644 --- a/core/service/serviceconf.go +++ b/core/service/serviceconf.go @@ -37,6 +37,7 @@ type ( Prometheus prometheus.Config `json:",optional"` Telemetry trace.Config `json:",optional"` DevServer DevServerConfig `json:",optional"` + Proc proc.ProcConf `json:",optional"` } ) @@ -61,6 +62,7 @@ func (sc ServiceConf) SetUp() error { sc.Telemetry.Name = sc.Name } trace.StartAgent(sc.Telemetry) + proc.Setup(sc.Proc) proc.AddShutdownListener(func() { trace.StopAgent() })