From c696a859be5c28a14196f2a62d3c0a9a01b98d5e Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@unistack.org>
Date: Sun, 10 Nov 2019 00:52:41 +0300
Subject: [PATCH] fix data race for server Wait option (#931)

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
---
 server/rpc_server.go | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/server/rpc_server.go b/server/rpc_server.go
index 9e4828c6..850fd97b 100644
--- a/server/rpc_server.go
+++ b/server/rpc_server.go
@@ -112,8 +112,12 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 		wg.Add(1)
 
 		// add to wait group if "wait" is opt-in
-		if s.wg != nil {
-			s.wg.Add(1)
+		s.Lock()
+		swg := s.wg
+		s.Unlock()
+
+		if swg != nil {
+			swg.Add(1)
 		}
 
 		// check we have an existing socket
@@ -132,8 +136,8 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 			}
 
 			// done(1)
-			if s.wg != nil {
-				s.wg.Done()
+			if swg != nil {
+				swg.Done()
 			}
 
 			wg.Done()
@@ -203,8 +207,8 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 					Body: []byte(err.Error()),
 				})
 
-				if s.wg != nil {
-					s.wg.Done()
+				if swg != nil {
+					swg.Done()
 				}
 
 				wg.Done()
@@ -324,8 +328,8 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 			mtx.Unlock()
 
 			// signal we're done
-			if s.wg != nil {
-				s.wg.Done()
+			if swg != nil {
+				swg.Done()
 			}
 
 			// done with this socket
@@ -744,9 +748,13 @@ func (s *rpcServer) Start() error {
 			log.Logf("Server %s-%s deregister error: %s", config.Name, config.Id, err)
 		}
 
+		s.Lock()
+		swg := s.wg
+		s.Unlock()
+
 		// wait for requests to finish
-		if s.wg != nil {
-			s.wg.Wait()
+		if swg != nil {
+			swg.Wait()
 		}
 
 		// close transport listener