server: guard sessions access in mutex

server/context.go

@@ -64,6 +64,12 @@ func (s *SessionManager) NewSession(conn *mysql.Conn) {
 	s.mu.Unlock()
 }
 
+func (s *SessionManager) session(conn *mysql.Conn) sql.Session {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.sessions[conn.ConnectionID]
+}
+
 // NewContext creates a new context for the session at the given conn.
 func (s *SessionManager) NewContext(conn *mysql.Conn) *sql.Context {
 	return s.NewContextWithQuery(conn, "")

server/handler.go

@@ -150,12 +150,11 @@ func (h *Handler) ComQuery(
 // ComQuery callback if the result does not contain any fields,
 // or after the last ComQuery call completes.
 func (h *Handler) WarningCount(c *mysql.Conn) uint16 {
-	sess, ok := h.sm.sessions[c.ConnectionID]
-	if !ok {
-		return 0
+	if sess := h.sm.session(c); sess != nil {
+		return sess.WarningCount()
 	}
 
-	return sess.WarningCount()
+	return 0
 }
 
 func (h *Handler) handleKill(conn *mysql.Conn, query string) (bool, error) {