From 4cb24e34ae5826fef056f5b20a4c802aee6096fa Mon Sep 17 00:00:00 2001
From: Benjamin Wang <benjamin.ahrtr@gmail.com>
Date: Fri, 17 Jan 2025 19:16:12 +0000
Subject: [PATCH] Enhance method (*serveCtx) serve to wait for all goroutines
 to complete before it returns

Signed-off-by: Benjamin Wang <benjamin.ahrtr@gmail.com>
---
 server/embed/serve.go | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/server/embed/serve.go b/server/embed/serve.go
index c888e3b1cfb..5615a0c2120 100644
--- a/server/embed/serve.go
+++ b/server/embed/serve.go
@@ -68,6 +68,8 @@ type serveCtx struct {
 	serviceRegister func(*grpc.Server)
 	serversC        chan *servers
 	closeOnce       sync.Once
+
+	wg sync.WaitGroup
 }
 
 type servers struct {
@@ -182,13 +184,17 @@ func (sctx *serveCtx) serve(
 			server = m.Serve
 
 			httpl := m.Match(cmux.HTTP1())
+			sctx.wg.Add(1)
 			go func(srvhttp *http.Server, tlsLis net.Listener) {
+				defer sctx.wg.Done()
 				errHandler(srvhttp.Serve(tlsLis))
 			}(srv, httpl)
 
 			if grpcEnabled {
 				grpcl := m.Match(cmux.HTTP2())
+				sctx.wg.Add(1)
 				go func(gs *grpc.Server, l net.Listener) {
+					defer sctx.wg.Done()
 					errHandler(gs.Serve(l))
 				}(gs, grpcl)
 			}
@@ -237,7 +243,7 @@ func (sctx *serveCtx) serve(
 				TLSConfig: tlscfg,
 				ErrorLog:  logger, // do not log user error
 			}
-			if err := configureHTTPServer(srv, s.Cfg); err != nil {
+			if err = configureHTTPServer(srv, s.Cfg); err != nil {
 				sctx.lg.Error("Configure https server failed", zap.Error(err))
 				return err
 			}
@@ -248,11 +254,13 @@ func (sctx *serveCtx) serve(
 		} else {
 			server = m.Serve
 
-			tlsl, err := transport.NewTLSListener(m.Match(cmux.Any()), tlsinfo)
-			if err != nil {
-				return err
+			tlsl, tlsErr := transport.NewTLSListener(m.Match(cmux.Any()), tlsinfo)
+			if tlsErr != nil {
+				return tlsErr
 			}
+			sctx.wg.Add(1)
 			go func(srvhttp *http.Server, tlsl net.Listener) {
+				defer sctx.wg.Done()
 				errHandler(srvhttp.Serve(tlsl))
 			}(srv, tlsl)
 		}
@@ -265,7 +273,11 @@ func (sctx *serveCtx) serve(
 		)
 	}
 
-	return server()
+	err = server()
+	sctx.close()
+	// ensure all goroutines, which are created by this method, to complete before this method returns.
+	sctx.wg.Wait()
+	return err
 }
 
 func configureHTTPServer(srv *http.Server, cfg config.ServerConfig) error {
@@ -334,7 +346,9 @@ func (sctx *serveCtx) registerGateway(dial func(ctx context.Context) (*grpc.Clie
 			return nil, err
 		}
 	}
+	sctx.wg.Add(1)
 	go func() {
+		defer sctx.wg.Done()
 		<-ctx.Done()
 		if cerr := conn.Close(); cerr != nil {
 			sctx.lg.Warn(