From c2ff00883aaffbae79207e64ade4bf2149d3ff43 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Mon, 23 Oct 2023 21:03:05 +0800
Subject: [PATCH] chore: update restful/grpc servers shutdown stages (#3662)

---
 rest/internal/starter.go        | 2 +-
 zrpc/internal/rpcserver.go      | 2 +-
 zrpc/internal/rpcserver_test.go | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/rest/internal/starter.go b/rest/internal/starter.go
index 08cca832..17430334 100644
--- a/rest/internal/starter.go
+++ b/rest/internal/starter.go
@@ -43,7 +43,7 @@ func start(host string, port int, handler http.Handler, run func(svr *http.Serve
 	}
 	healthManager := health.NewHealthManager(fmt.Sprintf("%s-%s:%d", probeNamePrefix, host, port))
 
-	waitForCalled := proc.AddWrapUpListener(func() {
+	waitForCalled := proc.AddShutdownListener(func() {
 		healthManager.MarkNotReady()
 		if e := server.Shutdown(context.Background()); e != nil {
 			logx.Error(e)
diff --git a/zrpc/internal/rpcserver.go b/zrpc/internal/rpcserver.go
index f1b90e42..ed4aa235 100644
--- a/zrpc/internal/rpcserver.go
+++ b/zrpc/internal/rpcserver.go
@@ -76,7 +76,7 @@ func (s *rpcServer) Start(register RegisterFn) error {
 
 	// we need to make sure all others are wrapped up,
 	// so we do graceful stop at shutdown phase instead of wrap up phase
-	waitForCalled := proc.AddWrapUpListener(func() {
+	waitForCalled := proc.AddShutdownListener(func() {
 		if s.health != nil {
 			s.health.Shutdown()
 		}
diff --git a/zrpc/internal/rpcserver_test.go b/zrpc/internal/rpcserver_test.go
index 99b48eb4..7af40110 100644
--- a/zrpc/internal/rpcserver_test.go
+++ b/zrpc/internal/rpcserver_test.go
@@ -47,7 +47,7 @@ func TestRpcServer(t *testing.T) {
 	grpcServer.GracefulStop()
 	lock.Unlock()
 
-	proc.WrapUp()
+	proc.Shutdown()
 	wgDone.Wait()
 }