fix: add channel signal judge to allow onShutdownHook to complete or timeout

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)

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