From 4ab6503a0e053f76940bafb29bd303dd20a87a9e Mon Sep 17 00:00:00 2001
From: smallnest <smallnest@gmail.com>
Date: Thu, 22 Aug 2024 19:47:56 +0800
Subject: [PATCH] fix data race in case of unregister services

---
 server/server.go  | 10 ++++------
 server/service.go |  7 ++++++-
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/server/server.go b/server/server.go
index 26a176ad..644f9ce8 100644
--- a/server/server.go
+++ b/server/server.go
@@ -126,7 +126,8 @@ type Server struct {
 	ServerErrorFunc func(res *protocol.Message, err error) string
 
 	// The server is started.
-	Started chan struct{}
+	Started           chan struct{}
+	unregisterAllOnce sync.Once
 }
 
 // NewServer returns a server.
@@ -955,11 +956,8 @@ func (s *Server) Shutdown(ctx context.Context) error {
 		s.mu.Lock()
 
 		// 主动注销注册的服务
-		if s.Plugins != nil {
-			for name := range s.serviceMap {
-				s.Plugins.DoUnregister(name)
-			}
-		}
+		s.UnregisterAll()
+
 		if s.ln != nil {
 			s.ln.Close()
 		}
diff --git a/server/service.go b/server/service.go
index 2a62d81e..75318b3f 100644
--- a/server/service.go
+++ b/server/service.go
@@ -325,7 +325,12 @@ func suitableMethods(typ reflect.Type, reportErr bool) map[string]*methodType {
 // UnregisterAll unregisters all services.
 // You can call this method when you want to shutdown/upgrade this node.
 func (s *Server) UnregisterAll() error {
-	return s.unregisterAll()
+	var err error
+	s.unregisterAllOnce.Do(func() {
+		err = s.unregisterAll()
+	})
+
+	return err
 }
 
 func (s *Server) unregisterAll() error {