MM-21012: Revamp websocket implementation (#16620)

* MM-21012: Revamp websocket implementation We replace the old gorilla/websocket implementation with the gobwas/ws library. The gorilla library was in maintenance mode and had a high level API due to which we cannot use that for situations where a large number of concurrent connections needs to be supported. The ws library is a very low-level library that allows us to work with raw net.Conns. We make several improvements: - We completely remove the reader goroutines, and instead replace them with a manual epoll implementation which sends off messages to be read when it receives any data on the connection. This lets us scale to a much larger number of connections. - The reader buffer is eliminated, because we directly read from the connection now. https://mattermost.atlassian.net/browse/MM-21012 ```release-notes Improved the websocket implementation by using epoll manually to read from a websocket. As a result, the number of goroutines is expected to go down by half. ``` * fix tests * fix shadowing errors * final changes * windows support! * Remove pointer to waitgroup * Fix edge case * Trigger CI * Trigger CI Co-authored-by: Mattermod <mattermod@users.noreply.github.com>
2021-02-13 23:42:11 +05:30
--- a/app/app_iface.go
+++ b/app/app_iface.go
@@ -12,12 +12,12 @@ import (
 	"crypto/ecdsa"
 	"io"
 	"mime/multipart"
+	"net"
 	"net/http"
 	"net/url"
 	"time"

 	"github.com/dyatlov/go-opengraph/opengraph"
-	"github.com/gorilla/websocket"
 	"github.com/mattermost/go-i18n/i18n"
 	goi18n "github.com/mattermost/go-i18n/i18n"
 	"github.com/mattermost/mattermost-server/v5/audit"
@@ -231,7 +231,7 @@ type AppIface interface {
 	// function is only exposed to sysadmins and the possibility of this edge case is relatively small.
 	MoveChannel(team *model.Team, channel *model.Channel, user *model.User) *model.AppError
 	// NewWebConn returns a new WebConn instance.
-	NewWebConn(ws *websocket.Conn, session model.Session, t goi18n.TranslateFunc, locale string) *WebConn
+	NewWebConn(ws net.Conn, session model.Session, t goi18n.TranslateFunc, locale string) *WebConn
 	// NewWebHub creates a new Hub.
 	NewWebHub() *Hub
 	// NotifySessionsExpired is called periodically from the job server to notify any mobile sessions that have expired.
--- a/app/opentracing/opentracing_layer.go
+++ b/app/opentracing/opentracing_layer.go
@@ -12,12 +12,12 @@ import (
 	"crypto/ecdsa"
 	"io"
 	"mime/multipart"
+	"net"
 	"net/http"
 	"net/url"
 	"time"

 	"github.com/dyatlov/go-opengraph/opengraph"
-	"github.com/gorilla/websocket"
 	"github.com/mattermost/go-i18n/i18n"
 	goi18n "github.com/mattermost/go-i18n/i18n"
 	"github.com/mattermost/mattermost-server/v5/app"
@@ -10858,7 +10858,7 @@ func (a *OpenTracingAppLayer) NewPluginAPI(manifest *model.Manifest) plugin.API
 	return resultVar0
 }

-func (a *OpenTracingAppLayer) NewWebConn(ws *websocket.Conn, session model.Session, t goi18n.TranslateFunc, locale string) *app.WebConn {
+func (a *OpenTracingAppLayer) NewWebConn(ws net.Conn, session model.Session, t goi18n.TranslateFunc, locale string) *app.WebConn {
 	origCtx := a.ctx
 	span, newCtx := tracing.StartSpanWithParentByContext(a.ctx, "app.NewWebConn")

--- a/app/server.go
+++ b/app/server.go
@@ -30,6 +30,7 @@ import (
 	"github.com/getsentry/sentry-go"
 	sentryhttp "github.com/getsentry/sentry-go/http"
 	"github.com/gorilla/mux"
+	"github.com/mailru/easygo/netpoll"
 	"github.com/pkg/errors"
 	"github.com/rs/cors"
 	"golang.org/x/crypto/acme/autocert"
@@ -102,8 +103,11 @@ type Server struct {

 	EmailService *EmailService

-	hubs     []*Hub
-	hashSeed maphash.Seed
+	hubs          []*Hub
+	hashSeed      maphash.Seed
+	poller        netpoll.Poller
+	webConnSema   chan struct{}
+	webConnSemaWg sync.WaitGroup

 	PushNotificationsHub   PushNotificationsHub
 	pushNotificationClient *http.Client // TODO: move this to it's own package
@@ -226,6 +230,16 @@ func NewServer(options ...Option) (*Server, error) {
 		mlog.Error("Could not initiate logging", mlog.Err(err))
 	}

+	// epoll/kqueue is not available on Windows.
+	if runtime.GOOS != "windows" {
+		poller, err := netpoll.New(nil)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to create a netpoll instance")
+		}
+		s.poller = poller
+		s.webConnSema = make(chan struct{}, runtime.NumCPU()*8) // numCPU * 8 is a good amount of concurrency.
+	}
+
 	// This is called after initLogging() to avoid a race condition.
 	mlog.Info("Server is initializing...", mlog.String("go_version", runtime.Version()))

@@ -259,9 +273,9 @@ func NewServer(options ...Option) (*Server, error) {
 	}

 	if *s.Config().ServiceSettings.EnableOpenTracing {
-		tracer, err := tracing.New()
-		if err != nil {
-			return nil, err
+		tracer, err2 := tracing.New()
+		if err2 != nil {
+			return nil, err2
 		}
 		s.tracer = tracer
 	}
@@ -1201,7 +1215,7 @@ func (a *App) OriginChecker() func(*http.Request) bool {

 		return utils.OriginChecker(allowed)
 	}
-	return nil
+	return utils.SameOriginChecker()
 }

 func (s *Server) checkPushNotificationServerUrl() {
@@ -1621,6 +1635,10 @@ func (s *Server) SetLog(l *mlog.Logger) {
 	s.Log = l
 }

+func (s *Server) Poller() netpoll.Poller {
+	return s.poller
+}
+
 func (a *App) GenerateSupportPacket() []model.FileData {
 	// If any errors we come across within this function, we will log it in a warning.txt file so that we know why certain files did not get produced if any
 	var warnings []string
--- a/app/web_conn.go
+++ b/app/web_conn.go
@@ -7,13 +7,19 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"net"
 	"net/http"
+	"runtime"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"

-	"github.com/gorilla/websocket"
+	"github.com/gobwas/ws"
+	"github.com/gobwas/ws/wsutil"
+	"github.com/mailru/easygo/netpoll"
 	goi18n "github.com/mattermost/go-i18n/i18n"
+	"github.com/pkg/errors"

 	"github.com/mattermost/mattermost-server/v5/mlog"
 	"github.com/mattermost/mattermost-server/v5/model"
@@ -31,29 +37,31 @@ const (
 )

 // WebConn represents a single websocket connection to a user.
-// It contains all the necesarry state to manage sending/receiving data to/from
+// It contains all the necessary state to manage sending/receiving data to/from
 // a websocket.
 type WebConn struct {
 	sessionExpiresAt int64 // This should stay at the top for 64-bit alignment of 64-bit words accessed atomically
 	App              *App
-	WebSocket        *websocket.Conn
+	WebSocket        net.Conn
 	T                goi18n.TranslateFunc
 	Locale           string
 	Sequence         int64
 	UserId           string

+	readMut                   sync.Mutex
 	allChannelMembers         map[string]string
 	lastAllChannelMembersTime int64
 	lastUserActivityAt        int64
 	send                      chan model.WebSocketMessage
 	sessionToken              atomic.Value
 	session                   atomic.Value
+	isWindows                 bool
 	endWritePump              chan struct{}
 	pumpFinished              chan struct{}
 }

 // NewWebConn returns a new WebConn instance.
-func (a *App) NewWebConn(ws *websocket.Conn, session model.Session, t goi18n.TranslateFunc, locale string) *WebConn {
+func (a *App) NewWebConn(ws net.Conn, session model.Session, t goi18n.TranslateFunc, locale string) *WebConn {
 	if session.UserId != "" {
 		a.Srv().Go(func() {
 			a.SetStatusOnline(session.UserId, false)
@@ -69,6 +77,7 @@ func (a *App) NewWebConn(ws *websocket.Conn, session model.Session, t goi18n.Tra
 		UserId:             session.UserId,
 		T:                  t,
 		Locale:             locale,
+		isWindows:          runtime.GOOS == "windows",
 		endWritePump:       make(chan struct{}),
 		pumpFinished:       make(chan struct{}),
 	}
@@ -77,12 +86,22 @@ func (a *App) NewWebConn(ws *websocket.Conn, session model.Session, t goi18n.Tra
 	wc.SetSessionToken(session.Token)
 	wc.SetSessionExpiresAt(session.ExpiresAt)

+	// epoll/kqueue is not available on Windows.
+	if !wc.isWindows {
+		wc.startPoller()
+	}
 	return wc
 }

 // Close closes the WebConn.
 func (wc *WebConn) Close() {
 	wc.WebSocket.Close()
+	if !wc.isWindows {
+		// This triggers the pump exit.
+		// If the pump has already exited, this just becomes a noop.
+		close(wc.endWritePump)
+	}
+	// We wait for the pump to fully exit.
 	<-wc.pumpFinished
 }

@@ -121,8 +140,20 @@ func (wc *WebConn) SetSession(v *model.Session) {
 }

 // Pump starts the WebConn instance. After this, the websocket
-// is ready to send/receive messages.
+// is ready to send messages.
+// This is only used by *nix platforms.
 func (wc *WebConn) Pump() {
+	// writePump is blocking in nature.
+	wc.writePump()
+	// Once it exits, we close everything.
+	wc.App.HubUnregister(wc)
+	close(wc.pumpFinished)
+}
+
+// BlockingPump is the Windows alternative of Pump.
+// It creates two goroutines - one for reading, another
+// for writing.
+func (wc *WebConn) BlockingPump() {
 	var wg sync.WaitGroup
 	wg.Add(1)
 	go func() {
@@ -138,29 +169,108 @@ func (wc *WebConn) Pump() {
 	defer ReturnSessionToPool(wc.GetSession())
 }

-func (wc *WebConn) readPump() {
-	defer func() {
-		wc.WebSocket.Close()
-	}()
-	wc.WebSocket.SetReadLimit(model.SOCKET_MAX_MESSAGE_SIZE_KB)
-	wc.WebSocket.SetReadDeadline(time.Now().Add(pongWaitTime))
-	wc.WebSocket.SetPongHandler(func(string) error {
+// startPoller adds the file descriptor of the connection
+// to the global epoll instance and registers a callback.
+func (wc *WebConn) startPoller() {
+	desc := netpoll.Must(netpoll.HandleRead(wc.WebSocket))
+	wc.App.Srv().Poller().Start(desc, func(wsEv netpoll.Event) {
+		if wsEv&(netpoll.EventReadHup|netpoll.EventHup) != 0 {
+			wc.App.Srv().Poller().Stop(desc)
+			wc.Close()
+			return
+		}
+
+		// Block until we have a token.
+		wc.App.Srv().GetWebConnToken()
+		// Read from conn.
+		go func() {
+			defer wc.App.Srv().ReleaseWebConnToken()
+			err := wc.ReadMsg()
+			if err != nil {
+				mlog.Debug("Error while reading message from websocket", mlog.Err(err))
+				wc.App.Srv().Poller().Stop(desc)
+				// net.ErrClosed is not available until Go 1.16.
+				// https://github.com/golang/go/issues/4373
+				//
+				// Sometimes, the netpoller generates a data event and a HUP event
+				// close to each other. In that case, we don't want to double-close
+				// the connection.
+				if !strings.Contains(err.Error(), "use of closed network connection") {
+					wc.Close()
+				}
+			}
+		}()
+	})
+}
+
+// GetWebConnToken creates backpressure by using
+// a counting semaphore to limit the number of concurrent goroutines.
+func (s *Server) GetWebConnToken() {
+	s.webConnSemaWg.Add(1)
+	s.webConnSema <- struct{}{}
+}
+
+// ReleaseWebConnToken releases a token
+// got from the semaphore
+func (s *Server) ReleaseWebConnToken() {
+	<-s.webConnSema
+	s.webConnSemaWg.Done()
+}
+
+// ReadMsg will read a single message from the websocket connection.
+func (wc *WebConn) ReadMsg() error {
+	r := wsutil.NewReader(wc.WebSocket, ws.StateServerSide)
+	r.MaxFrameSize = model.SOCKET_MAX_MESSAGE_SIZE_KB
+	decoder := json.NewDecoder(r)
+
+	// The reader's methods are not goroutine safe.
+	// We restrict only one reader goroutine per-connection.
+	wc.readMut.Lock()
+	defer wc.readMut.Unlock()
+
+	var req model.WebSocketRequest
+	hdr, err := r.NextFrame()
+	if err != nil {
+		return errors.Wrap(err, "error while getting the next websocket frame")
+	}
+	switch hdr.OpCode {
+	case ws.OpClose:
+		// Return if closed.
+		// We need to return an error for Windows to let the reader exit.
+		if wc.isWindows {
+			return errors.New("connection closed")
+		}
+		return nil
+	case ws.OpPong:
 		wc.WebSocket.SetReadDeadline(time.Now().Add(pongWaitTime))
+		// Handle pongs
 		if wc.IsAuthenticated() {
 			wc.App.Srv().Go(func() {
 				wc.App.SetStatusAwayIfNeeded(wc.UserId, false)
 			})
 		}
-		return nil
-	})
+	default:
+		// Default case of data message.
+		if err := decoder.Decode(&req); err != nil {
+			// We discard any remaining data left in the socket.
+			r.Discard()
+			return errors.Wrap(err, "error during decoding websocket message")
+		}
+
+		wc.App.Srv().WebSocketRouter.ServeWebSocket(wc, &req)
+	}
+	return nil
+}
+
+func (wc *WebConn) readPump() {
+	defer wc.WebSocket.Close()
+	wc.WebSocket.SetReadDeadline(time.Now().Add(pongWaitTime))

 	for {
-		var req model.WebSocketRequest
-		if err := wc.WebSocket.ReadJSON(&req); err != nil {
+		if err := wc.ReadMsg(); err != nil {
 			wc.logSocketErr("websocket.read", err)
 			return
 		}
-		wc.App.Srv().WebSocketRouter.ServeWebSocket(wc, &req)
 	}
 }

@@ -184,7 +294,7 @@ func (wc *WebConn) writePump() {
 		case msg, ok := <-wc.send:
 			if !ok {
 				wc.WebSocket.SetWriteDeadline(time.Now().Add(writeWaitTime))
-				wc.WebSocket.WriteMessage(websocket.CloseMessage, []byte{})
+				wsutil.WriteServerMessage(wc.WebSocket, ws.OpClose, []byte{})
 				return
 			}

@@ -239,7 +349,7 @@ func (wc *WebConn) writePump() {
 			}

 			wc.WebSocket.SetWriteDeadline(time.Now().Add(writeWaitTime))
-			if err := wc.WebSocket.WriteMessage(websocket.TextMessage, buf.Bytes()); err != nil {
+			if err := wsutil.WriteServerMessage(wc.WebSocket, ws.OpText, buf.Bytes()); err != nil {
 				wc.logSocketErr("websocket.send", err)
 				return
 			}
@@ -249,7 +359,7 @@ func (wc *WebConn) writePump() {
 			}
 		case <-ticker.C:
 			wc.WebSocket.SetWriteDeadline(time.Now().Add(writeWaitTime))
-			if err := wc.WebSocket.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
+			if err := wsutil.WriteServerMessage(wc.WebSocket, ws.OpPing, []byte{}); err != nil {
 				wc.logSocketErr("websocket.ticker", err)
 				return
 			}
@@ -436,10 +546,5 @@ func (wc *WebConn) isMemberOfTeam(teamID string) bool {
 }

 func (wc *WebConn) logSocketErr(source string, err error) {
-	// browsers will appear as CloseNoStatusReceived
-	if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseNoStatusReceived) {
-		mlog.Debug(source+": client side closed socket", mlog.String("user_id", wc.UserId))
-	} else {
-		mlog.Debug(source+": closing websocket", mlog.String("user_id", wc.UserId), mlog.Err(err))
-	}
+	mlog.Debug(source+": error during writing to websocket", mlog.String("user_id", wc.UserId), mlog.Err(err))
 }
--- a/app/web_hub.go
+++ b/app/web_hub.go
@@ -114,6 +114,9 @@ func (a *App) InvalidateWebConnSessionCacheForUser(userID string) {
 func (s *Server) HubStop() {
 	mlog.Info("stopping websocket hub connections")

+	// Wait until all messages have finished reading.
+	s.webConnSemaWg.Wait()
+	// Now stop the hub.
 	for _, hub := range s.hubs {
 		hub.Stop()
 	}
--- a/app/web_hub_test.go
+++ b/app/web_hub_test.go
@@ -4,13 +4,17 @@
 package app

 import (
+	"context"
+	"io"
+	"io/ioutil"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 	"time"

-	"github.com/gorilla/websocket"
+	"github.com/gobwas/ws"
+	"github.com/gobwas/ws/wsutil"
 	goi18n "github.com/mattermost/go-i18n/i18n"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -22,16 +26,35 @@ import (

 func dummyWebsocketHandler(t *testing.T) http.HandlerFunc {
 	return func(w http.ResponseWriter, req *http.Request) {
-		upgrader := &websocket.Upgrader{
-			ReadBufferSize:  1024,
-			WriteBufferSize: 1024,
+		upgrader := ws.HTTPUpgrader{
+			Timeout: 5 * time.Second,
+		}
+
+		var hdr ws.Header
+		conn, _, _, err := upgrader.Upgrade(req, w)
+		rd := wsutil.Reader{
+			Source:          conn,
+			State:           ws.StateServerSide,
+			CheckUTF8:       true,
+			SkipHeaderCheck: true,
 		}
-		conn, err := upgrader.Upgrade(w, req, nil)
 		for err == nil {
-			_, _, err = conn.ReadMessage()
+			hdr, err = rd.NextFrame()
+			if err != nil {
+				continue
+			}
+			if hdr.OpCode.IsControl() {
+				continue
+			}
+			if hdr.OpCode&(ws.OpText|ws.OpBinary) == 0 {
+				err = rd.Discard()
+				continue
+			}
+
+			_, err = ioutil.ReadAll(&rd)
 		}
-		if _, ok := err.(*websocket.CloseError); !ok {
-			require.NoError(t, err)
+		if err != io.EOF {
+			require.Fail(t, "unexpected error:", err)
 		}
 	}
 }
@@ -42,8 +65,7 @@ func registerDummyWebConn(t *testing.T, a *App, addr net.Addr, userID string) *W
 	})
 	require.Nil(t, appErr)

-	d := websocket.Dialer{}
-	c, _, err := d.Dial("ws://"+addr.String()+"/ws", nil)
+	c, _, _, err := ws.Dial(context.Background(), "ws://"+addr.String()+"/ws")
 	require.NoError(t, err)

 	wc := a.NewWebConn(c, *session, goi18n.IdentityTfunc(), "en")
@@ -78,8 +100,7 @@ func TestHubStopRaceCondition(t *testing.T) {
 	s := httptest.NewServer(dummyWebsocketHandler(t))

 	th.App.HubStart()
-	wc1 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
-	defer wc1.Close()
+	registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)

 	hub := th.App.Srv().hubs[0]
 	th.App.HubStop()