checkpoint fixing data race to process h2 and grpc requests

server/rpc_server.go

@@ -63,20 +63,33 @@ func (r rpcRouter) ServeRequest(ctx context.Context, req Request, rsp Response)
 
 // ServeConn serves a single connection
 func (s *rpcServer) ServeConn(sock transport.Socket) {
+	var wg sync.WaitGroup
+	var mtx sync.RWMutex
+	// streams are multiplexed on Micro-Stream or Micro-Id header
+	sockets := make(map[string]*socket.Socket)
+
 	defer func() {
-		// close socket
+		// wait till done
+		wg.Wait()
+
+		// close underlying socket
 		sock.Close()
 
+		// close the sockets
+		mtx.Lock()
+		for id, psock := range sockets {
+			psock.Close()
+			delete(sockets, id)
+		}
+		mtx.Unlock()
+
+		// recover any panics
 		if r := recover(); r != nil {
 			log.Log("panic recovered: ", r)
 			log.Log(string(debug.Stack()))
 		}
 	}()
 
-	// multiplex the streams on a single socket by Micro-Stream
-	var mtx sync.RWMutex
-	sockets := make(map[string]*socket.Socket)
-
 	for {
 		var msg transport.Message
 		if err := sock.Recv(&msg); err != nil {
@@ -94,6 +107,9 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 			id = msg.Header["Micro-Id"]
 		}
 
+		// we're starting processing
+		wg.Add(1)
+
 		// add to wait group if "wait" is opt-in
 		if s.wg != nil {
 			s.wg.Add(1)
@@ -119,6 +135,8 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 				s.wg.Done()
 			}
 
+			wg.Done()
+
 			// continue to the next message
 			continue
 		}
@@ -136,9 +154,14 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 		sockets[id] = psock
 		mtx.Unlock()
 
+		// wait for processing to exit
+		wg.Add(1)
+
 		// process the outbound messages from the socket
 		go func(id string, psock *socket.Socket) {
-			defer psock.Close()
+			defer func() {
+				wg.Done()
+			}()
 
 			for {
 				// get the message from our internal handler/stream
@@ -155,6 +178,7 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 				if err := sock.Send(m); err != nil {
 					return
 				}
+
 			}
 		}(id, psock)
 
@@ -205,9 +229,13 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 					},
 					Body: []byte(err.Error()),
 				})
+
 				if s.wg != nil {
 					s.wg.Done()
 				}
+
+				wg.Done()
+
 				return
 			}
 		}
@@ -261,6 +289,8 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 
 		// serve the request in a go routine as this may be a stream
 		go func(id string, psock *socket.Socket) {
+			defer psock.Close()
+
 			// serve the actual request using the request router
 			if err := r.ServeRequest(ctx, request, response); err != nil {
 				// write an error response
@@ -285,8 +315,9 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 				s.wg.Done()
 			}
 
+			// done with this socket
+			wg.Done()
 		}(id, psock)
-
 	}
 }