chore(graceful): support custom timeout value (#466)

fixed: #465

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

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"))

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)

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

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)
 		}
 	})

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)

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)
 		}
 	})