Repository: cloudwego/hertz
Pull request #249 (2022-09-16 09:24:20)

fix: add channel signal judge to allow onShutdownHook to complete or timeout
#### What type of PR is this?
- fix

#### Check the PR title.
- [x] This PR title match the format: \<type\>(optional scope): \<description\>
- [x] The description of this PR title is user-oriented and clear enough for others to understand.


#### (Optional) More detail description for this PR(en: English/zh: Chinese).

- en: add channel signal judge to allow onShutdownHook to complete or timeout
- zh(optional):  添加`channel`信号判断，允许 `onShutdownHook` 执行完成或超时
#### Which issue(s) this PR fixes:


codecov[bot]: # [Codecov](https://codecov.io/gh/cloudwego/hertz/pull/249?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) Report
Base: **59.17**% // Head: **59.17**% // No change to project coverage :thumbsup:
> Coverage data is based on head [(`dcb2665`)](https://codecov.io/gh/cloudwego/hertz/pull/249?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) compared to base [(`97bb37e`)](https://codecov.io/gh/cloudwego/hertz/commit/97bb37e14eec8fd29e7d23e0ce25fe4e01d76954?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).
> Patch coverage: 0.00% of modified lines in pull request are covered.

> :exclamation: Current head dcb2665 differs from pull request most recent head 5fb7ec1. Consider uploading reports for the commit 5fb7ec1 to get more accurate results

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff            @@
##           develop     #249   +/-   ##
========================================
  Coverage    59.17%   59.17%           
========================================
  Files           79       79           
  Lines         8111     8111           
========================================
  Hits          4800     4800           
  Misses        2957     2957           
  Partials       354      354           
```


| [Impacted Files](https://codecov.io/gh/cloudwego/hertz/pull/249?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) | Coverage Δ | |
|---|---|---|
| [pkg/route/engine.go](https://codecov.io/gh/cloudwego/hertz/pull/249/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3JvdXRlL2VuZ2luZS5nbw==) | `64.50% <0.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/cloudwego/hertz/pull/249?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).

welkeyever: [coverage.txt](https://github.com/cloudwego/hertz/pull/249/files#diff-093b633c1617628a597c0ebdbc55ffc8e5c48eb1458ff70598ba5ae55dac4ba7) 这个去掉吧..
li-jin-gou: > [coverage.txt](https://github.com/cloudwego/hertz/pull/249/files#diff-093b633c1617628a597c0ebdbc55ffc8e5c48eb1458ff70598ba5ae55dac4ba7) 这个去掉吧..
done

Patch:
diff --git a/pkg/protocol/header.go b/pkg/protocol/header.go
index 0babea5..bb2e033 100644
--- a/pkg/protocol/header.go
+++ b/pkg/protocol/header.go
@@ -52,6 +52,7 @@ import (
 	"github.com/cloudwego/hertz/internal/bytestr"
 	"github.com/cloudwego/hertz/internal/nocopy"
 	errs "github.com/cloudwego/hertz/pkg/common/errors"
+	"github.com/cloudwego/hertz/pkg/common/hlog"
 	"github.com/cloudwego/hertz/pkg/common/utils"
 	"github.com/cloudwego/hertz/pkg/protocol/consts"
 )
@@ -447,9 +448,19 @@ func (h *RequestHeader) SetHost(host string) {
 
 // SetStatusCode sets response status code.
 func (h *ResponseHeader) SetStatusCode(statusCode int) {
+	checkWriteHeaderCode(statusCode)
 	h.statusCode = statusCode
 }
 
+func checkWriteHeaderCode(code int) {
+	// For now, we only emit a warning for bad codes.
+	// In the future we might block things over 599 or under 100
+	if code < 100 || code > 599 {
+		hlog.Warnf("Invalid StatusCode code %v, status code should not be under 100 or over 599.\n"+
+			"For more info: https://www.rfc-editor.org/rfc/rfc9110.html#name-status-codes", code)
+	}
+}
+
 func (h *ResponseHeader) ResetSkipNormalize() {
 	h.noHTTP11 = false
 	h.connectionClose = false
diff --git a/pkg/route/engine.go b/pkg/route/engine.go
index a3e1baa..52849d4 100644
--- a/pkg/route/engine.go
+++ b/pkg/route/engine.go
@@ -251,7 +251,7 @@ func (engine *Engine) NewContext() *app.RequestContext {
 
 // Shutdown starts the server's graceful exit by next steps:
 //
-// 1. Trigger OnShutdown hooks concurrently, but don't wait them
+// 1. Trigger OnShutdown hooks concurrently and wait them until wait timeout or finish
 // 2. Close the net listener, which means new connection won't be accepted
 // 3. Wait all connections get closed:
 // 	One connection gets closed after reaching out the shorter time of processing
@@ -265,12 +265,21 @@ func (engine *Engine) Shutdown(ctx context.Context) (err error) {
 		return
 	}
 
+	ch := make(chan struct{})
 	// trigger hooks if any
-	for i := range engine.OnShutdown {
-		go func(index int) {
-			engine.OnShutdown[index](ctx)
-		}(i)
-	}
+	go engine.executeOnShutdownHooks(ctx, ch)
+
+	defer func() {
+		// ensure that the hook is executed until wait timeout or finish
+		select {
+		case <-ctx.Done():
+			hlog.Infof("HERTZ: Execute OnShutdownHooks timeout: error=%v", ctx.Err())
+			return
+		case <-ch:
+			hlog.Info("HERTZ: Execute OnShutdownHooks finish")
+			return
+		}
+	}()
 
 	if opt := engine.options; opt != nil && opt.Registry != nil {
 		if err = opt.Registry.Deregister(opt.RegistryInfo); err != nil {
@@ -283,9 +292,23 @@ func (engine *Engine) Shutdown(ctx context.Context) (err error) {
 	if err := engine.transport.Shutdown(ctx); err != ctx.Err() {
 		return err
 	}
+
 	return
 }
 
+func (engine *Engine) executeOnShutdownHooks(ctx context.Context, ch chan struct{}) {
+	wg := sync.WaitGroup{}
+	for i := range engine.OnShutdown {
+		wg.Add(1)
+		go func(index int) {
+			defer wg.Done()
+			engine.OnShutdown[index](ctx)
+		}(i)
+	}
+	wg.Wait()
+	ch <- struct{}{}
+}
+
 func (engine *Engine) Run() (err error) {
 	if err = engine.Init(); err != nil {
 		return err


Test patch:
diff --git a/pkg/app/server/hertz_test.go b/pkg/app/server/hertz_test.go
index 0d27901..de6b359 100644
--- a/pkg/app/server/hertz_test.go
+++ b/pkg/app/server/hertz_test.go
@@ -94,6 +94,7 @@ func TestHertz_GracefulShutdown(t *testing.T) {
 		atomic.StoreUint32(&testint2, 2)
 	})
 	engine.Engine.OnShutdown = append(engine.OnShutdown, func(ctx context.Context) {
+		time.Sleep(2 * time.Second)
 		atomic.StoreUint32(&testint3, 3)
 	})
 
@@ -196,7 +197,6 @@ func TestHertz_Spin(t *testing.T) {
 		t.Fatal(err)
 	}
 	t.Logf("[%v]end SIGHUP\n", time.Now())
-
 	<-ch
 	assert.Nil(t, err)
 	assert.NotNil(t, resp)
diff --git a/pkg/protocol/header_test.go b/pkg/protocol/header_test.go
index 301de8a..9a75b3d 100644
--- a/pkg/protocol/header_test.go
+++ b/pkg/protocol/header_test.go
@@ -42,13 +42,14 @@
 package protocol
 
 import (
+	"bytes"
 	"fmt"
 	"strings"
 	"testing"
 
-	"github.com/cloudwego/hertz/pkg/common/test/assert"
-
 	"github.com/cloudwego/hertz/internal/bytestr"
+	"github.com/cloudwego/hertz/pkg/common/hlog"
+	"github.com/cloudwego/hertz/pkg/common/test/assert"
 	"github.com/cloudwego/hertz/pkg/protocol/consts"
 )
 
@@ -193,6 +194,22 @@ func TestResponseHeaderDelClientCookie(t *testing.T) {
 	ReleaseCookie(c)
 }
 
+func TestCheckWriteHeaderCode(t *testing.T) {
+	buffer := bytes.NewBuffer(make([]byte, 0, 1024))
+	hlog.SetOutput(buffer)
+	checkWriteHeaderCode(99)
+	assert.True(t, strings.Contains(buffer.String(), "[Warn] Invalid StatusCode code"))
+	buffer.Reset()
+	checkWriteHeaderCode(600)
+	assert.True(t, strings.Contains(buffer.String(), "[Warn] Invalid StatusCode code"))
+	buffer.Reset()
+	checkWriteHeaderCode(100)
+	assert.False(t, strings.Contains(buffer.String(), "[Warn] Invalid StatusCode code"))
+	buffer.Reset()
+	checkWriteHeaderCode(599)
+	assert.False(t, strings.Contains(buffer.String(), "[Warn] Invalid StatusCode code"))
+}
+
 func TestResponseHeaderAdd(t *testing.T) {
 	t.Parallel()