chore(memory): bound amount of requests in flight on WS connection

Author: pragmaxim

server/public_test.go

@@ -1353,6 +1353,63 @@ func Test_WebsocketRejectsOversizedMessage(t *testing.T) {
 	t.Fatalf("unexpected websocket error after oversized message: %v", err)
 }
 
+func Test_WebsocketClosesWhenPendingRequestLimitExceeded(t *testing.T) {
+	parser, chain := setupChain(t)
+
+	s, dbpath := setupPublicHTTPServer(parser, chain, t, false)
+	defer closeAndDestroyPublicServer(t, s, dbpath)
+	s.ConnectFullPublicInterface()
+
+	releaseRequests := make(chan struct{})
+	defer close(releaseRequests)
+	startedRequests := make(chan struct{}, maxWebsocketPendingRequests)
+	originalPingHandler := requestHandlers["ping"]
+	requestHandlers["ping"] = func(s *WebsocketServer, c *websocketChannel, req *WsReq) (interface{}, error) {
+		startedRequests <- struct{}{}
+		<-releaseRequests
+		return struct{}{}, nil
+	}
+	defer func() {
+		requestHandlers["ping"] = originalPingHandler
+	}()
+
+	ts := httptest.NewServer(s.https.Handler)
+	defer ts.Close()
+
+	ws := connectWebsocket(t, ts)
+	defer ws.Close()
+
+	for i := 0; i < maxWebsocketPendingRequests; i++ {
+		if err := ws.WriteJSON(websocketReq{ID: strconv.Itoa(i), Method: "ping"}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	for i := 0; i < maxWebsocketPendingRequests; i++ {
+		select {
+		case <-startedRequests:
+		case <-time.After(2 * time.Second):
+			t.Fatalf("timed out waiting for pending request %d", i)
+		}
+	}
+
+	if err := ws.WriteJSON(websocketReq{ID: "overflow", Method: "ping"}); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := ws.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil {
+		t.Fatal(err)
+	}
+	_, _, err := ws.ReadMessage()
+	ws.SetReadDeadline(time.Time{})
+	if err == nil {
+		t.Fatal("expected websocket read error after pending request limit was exceeded")
+	}
+	var netErr net.Error
+	if errors.As(err, &netErr) && netErr.Timeout() {
+		t.Fatal("expected connection close after pending request limit was exceeded, got timeout")
+	}
+}
+
 var websocketTestsBitcoinType = []websocketTest{
 	{
 		name: "websocket getInfo",

server/websocket.go

@@ -28,6 +28,7 @@ const outChannelSize = 500
 const defaultTimeout = 60 * time.Second
 const unknownMethodLabel = "unknown"
 const maxWebsocketMessageBytes int64 = 4 * 1024 * 1024
+const maxWebsocketPendingRequests = 48
 const websocketLogPreviewBytes = 256
 
 // allRates is a special "currency" parameter that means all available currencies
@@ -44,6 +45,7 @@ type websocketChannel struct {
 	id                           uint64
 	conn                         *websocket.Conn
 	out                          chan *WsRes
+	pendingRequests              chan struct{}
 	ip                           string
 	requestHeader                http.Header
 	alive                        bool
@@ -221,12 +223,13 @@ func (s *WebsocketServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 	conn.SetReadLimit(maxWebsocketMessageBytes)
 	c := &websocketChannel{
-		id:            atomic.AddUint64(&connectionCounter, 1),
-		conn:          conn,
-		out:           make(chan *WsRes, outChannelSize),
-		ip:            getIP(r),
-		requestHeader: r.Header,
-		alive:         true,
+		id:              atomic.AddUint64(&connectionCounter, 1),
+		conn:            conn,
+		out:             make(chan *WsRes, outChannelSize),
+		pendingRequests: make(chan struct{}, maxWebsocketPendingRequests),
+		ip:              getIP(r),
+		requestHeader:   r.Header,
+		alive:           true,
 	}
 	if s.is.WsGetAccountInfoLimit > 0 {
 		c.getAddressInfoDescriptors = make(map[string]struct{})
@@ -290,6 +293,19 @@ func (c *websocketChannel) DataOut(data *WsRes) {
 	}
 }
 
+func (c *websocketChannel) acquireRequestSlot() bool {
+	select {
+	case c.pendingRequests <- struct{}{}:
+		return true
+	default:
+		return false
+	}
+}
+
+func (c *websocketChannel) releaseRequestSlot() {
+	<-c.pendingRequests
+}
+
 func (s *WebsocketServer) inputLoop(c *websocketChannel) {
 	defer func() {
 		if r := recover(); r != nil {
@@ -313,7 +329,15 @@ func (s *WebsocketServer) inputLoop(c *websocketChannel) {
 				s.closeChannel(c, "protocol_error")
 				return
 			}
-			go s.onRequest(c, &req)
+			if !c.acquireRequestSlot() {
+				glog.Warning("Client ", c.id, " exceeded pending websocket request limit, ", c.ip)
+				s.closeChannel(c, "pending_requests_limit")
+				return
+			}
+			go func(req WsReq) {
+				defer c.releaseRequestSlot()
+				s.onRequest(c, &req)
+			}(req)
 		case websocket.BinaryMessage:
 			glog.Error("Binary message received from ", c.id, ", ", c.ip)
 			s.closeChannel(c, "protocol_error")