From c379630c2966fda0b3cba268bac48774de6efca8 Mon Sep 17 00:00:00 2001 From: Bo-Yi Wu Date: Wed, 5 Feb 2020 23:23:16 +0800 Subject: [PATCH] chore(graceful): support custom timeout value (#466) fixed: https://github.com/appleboy/gorush/issues/465 --- README.md | 1 + config/config.go | 3 +++ config/config_test.go | 2 ++ config/testdata/config.yml | 1 + gorush/server_normal.go | 7 +++++++ main.go | 2 +- rpc/server.go | 4 ++++ 7 files changed, 19 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index bfda43f3..e0df2e5c 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,7 @@ See the default [YAML config example](config/config.yml): core: enabled: true # enabale httpd server address: "" # ip address to bind (default: any) + shutdown_timeout: 30 # default is 30 second port: "8088" # ignore this port number if auto_tls is enabled (listen 443). worker_num: 0 # default worker number is runtime.NumCPU() queue_num: 0 # default queue number is 8192 diff --git a/config/config.go b/config/config.go index 13c2b4e0..fbf0b241 100644 --- a/config/config.go +++ b/config/config.go @@ -14,6 +14,7 @@ var defaultConf = []byte(` core: enabled: true # enabale httpd server address: "" # ip address to bind (default: any) + shutdown_timeout: 30 # default is 30 second port: "8088" # ignore this port number if auto_tls is enabled (listen 443). worker_num: 0 # default worker number is runtime.NumCPU() queue_num: 0 # default queue number is 8192 @@ -103,6 +104,7 @@ type ConfYaml struct { type SectionCore struct { Enabled bool `yaml:"enabled"` Address string `yaml:"address"` + ShutdownTimeout int64 `yaml:"shutdown_timeout"` Port string `yaml:"port"` MaxNotification int64 `yaml:"max_notification"` WorkerNum int64 `yaml:"worker_num"` @@ -253,6 +255,7 @@ func LoadConf(confPath string) (ConfYaml, error) { // Core conf.Core.Address = viper.GetString("core.address") conf.Core.Port = viper.GetString("core.port") + conf.Core.ShutdownTimeout = int64(viper.GetInt("core.shutdown_timeout")) conf.Core.Enabled = viper.GetBool("core.enabled") conf.Core.WorkerNum = int64(viper.GetInt("core.worker_num")) conf.Core.QueueNum = int64(viper.GetInt("core.queue_num")) diff --git a/config/config_test.go b/config/config_test.go index 79bbedeb..71ece0e4 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -39,6 +39,7 @@ func (suite *ConfigTestSuite) TestValidateConfDefault() { // Core assert.Equal(suite.T(), "", suite.ConfGorushDefault.Core.Address) assert.Equal(suite.T(), "8088", suite.ConfGorushDefault.Core.Port) + assert.Equal(suite.T(), int64(30), suite.ConfGorushDefault.Core.ShutdownTimeout) assert.Equal(suite.T(), true, suite.ConfGorushDefault.Core.Enabled) assert.Equal(suite.T(), int64(runtime.NumCPU()), suite.ConfGorushDefault.Core.WorkerNum) assert.Equal(suite.T(), int64(8192), suite.ConfGorushDefault.Core.QueueNum) @@ -112,6 +113,7 @@ func (suite *ConfigTestSuite) TestValidateConfDefault() { func (suite *ConfigTestSuite) TestValidateConf() { // Core assert.Equal(suite.T(), "8088", suite.ConfGorush.Core.Port) + assert.Equal(suite.T(), int64(30), suite.ConfGorush.Core.ShutdownTimeout) assert.Equal(suite.T(), true, suite.ConfGorush.Core.Enabled) assert.Equal(suite.T(), int64(runtime.NumCPU()), suite.ConfGorush.Core.WorkerNum) assert.Equal(suite.T(), int64(8192), suite.ConfGorush.Core.QueueNum) diff --git a/config/testdata/config.yml b/config/testdata/config.yml index 6bd72826..56c8b2da 100644 --- a/config/testdata/config.yml +++ b/config/testdata/config.yml @@ -1,6 +1,7 @@ core: enabled: true # enabale httpd server address: "" # ip address to bind (default: any) + shutdown_timeout: 30 # default is 30 second port: "8088" # ignore this port number if auto_tls is enabled (listen 443). worker_num: 0 # default worker number is runtime.NumCPU() queue_num: 0 # default queue number is 8192 diff --git a/gorush/server_normal.go b/gorush/server_normal.go index 66131154..3b7b6747 100644 --- a/gorush/server_normal.go +++ b/gorush/server_normal.go @@ -8,6 +8,7 @@ import ( "encoding/base64" "errors" "net/http" + "time" "golang.org/x/sync/errgroup" ) @@ -73,6 +74,9 @@ func listenAndServe(ctx context.Context, s *http.Server) error { g.Go(func() error { select { case <-ctx.Done(): + timeout := time.Duration(PushConf.Core.ShutdownTimeout) * time.Second + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() return s.Shutdown(ctx) } }) @@ -87,6 +91,9 @@ func listenAndServeTLS(ctx context.Context, s *http.Server) error { g.Go(func() error { select { case <-ctx.Done(): + timeout := time.Duration(PushConf.Core.ShutdownTimeout) * time.Second + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() return s.Shutdown(ctx) } }) diff --git a/main.go b/main.go index 0ac06aab..fec5fd5d 100644 --- a/main.go +++ b/main.go @@ -252,7 +252,7 @@ func main() { wg := &sync.WaitGroup{} wg.Add(int(gorush.PushConf.Core.WorkerNum)) ctx := withContextFunc(context.Background(), func() { - gorush.LogAccess.Info("close the notification queue channel") + gorush.LogAccess.Info("close the notification queue channel, current queue len: ", len(gorush.QueueNotification)) close(gorush.QueueNotification) wg.Wait() close(finished) diff --git a/rpc/server.go b/rpc/server.go index 46e28946..5e6b587c 100644 --- a/rpc/server.go +++ b/rpc/server.go @@ -4,6 +4,7 @@ import ( "context" "net/http" "sync" + "time" "github.com/appleboy/gorush/gorush" "github.com/appleboy/gorush/rpc/proto" @@ -122,6 +123,9 @@ func RunGRPCServer(ctx context.Context) error { g.Go(func() error { select { case <-ctx.Done(): + timeout := time.Duration(gorush.PushConf.Core.ShutdownTimeout) * time.Second + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() return srv.Shutdown(ctx) } })