From 70f58b264a10558edbf1a7cb5503e8ee81f8e12d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8C=85=E5=AD=90?= <baozhecheng@foxmail.com>
Date: Thu, 28 Oct 2021 11:09:17 +0800
Subject: [PATCH] fix(app): use new context when app stop (#1589)

---
 app.go     | 5 ++++-
 options.go | 6 ++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/app.go b/app.go
index aa7ba8248..bfbb85e04 100644
--- a/app.go
+++ b/app.go
@@ -42,6 +42,7 @@ func New(opts ...Option) *App {
 		logger:           log.NewHelper(log.DefaultLogger),
 		sigs:             []os.Signal{syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT},
 		registrarTimeout: 10 * time.Second,
+		stopTimeout:      10 * time.Second,
 	}
 	if id, err := uuid.NewUUID(); err == nil {
 		o.id = id.String()
@@ -90,7 +91,9 @@ func (a *App) Run() error {
 		srv := srv
 		eg.Go(func() error {
 			<-ctx.Done() // wait for stop signal
-			return srv.Stop(ctx)
+			sctx, cancel := context.WithTimeout(context.Background(), a.opts.stopTimeout)
+			defer cancel()
+			return srv.Stop(sctx)
 		})
 		wg.Add(1)
 		eg.Go(func() error {
diff --git a/options.go b/options.go
index 29349f606..d18d75da1 100644
--- a/options.go
+++ b/options.go
@@ -28,6 +28,7 @@ type options struct {
 	logger           *log.Helper
 	registrar        registry.Registrar
 	registrarTimeout time.Duration
+	stopTimeout      time.Duration
 	servers          []transport.Server
 }
 
@@ -87,3 +88,8 @@ func Registrar(r registry.Registrar) Option {
 func RegistrarTimeout(t time.Duration) Option {
 	return func(o *options) { o.registrarTimeout = t }
 }
+
+// StopTimeout with app stop timeout.
+func StopTimeout(t time.Duration) Option {
+	return func(o *options) { o.stopTimeout = t }
+}