feat: customization of WebSocket-related parameters

Author: huaxiabuluo

app/pages/Schema/SchemaConfig/List/SpaceStats/index.tsx

@@ -127,7 +127,6 @@ const SpaceStats = () => {
     }
   };
   const showTime = updateTime && jobId == null && !loading;
-  console.log('=====updateTime', updateTime);
   return (
     <div className={styles.nebulaStats}>
       <div className={styles.row}>

app/utils/websocket.ts

@@ -225,17 +225,20 @@ export class NgqlRunner {
   onError = (e: Event) => {
     console.error('=====ngqlSocket error', e);
     message.error('WebSocket error, try to reconnect...');
-    this.onDisconnect();
+    this.onDisconnect(e);
   };
 
-  onDisconnect = () => {
-    console.log('=====onDisconnect');
+  onDisconnect = (e?: CloseEvent | Event) => {
+    console.log('=====onDisconnect', e);
     this.socket?.removeEventListener('close', this.onDisconnect);
     this.socket?.removeEventListener('error', this.onError);
 
     this.clearMessageReceiver();
     this.closeSocket();
 
+    const { code, reason } = (e as CloseEvent) || {};
+    reason && message.error(`WebSocket closed unexpectedly, code: ${code}, reason: \`${reason}\`, try to reconnect...`);
+
     // try reconnect
     this.socketUrl && setTimeout(this.reConnect, 3000);
   };

server/api/studio/etc/studio-api.yaml

@@ -12,6 +12,23 @@ Log:
   KeepDays: 7
 Debug:
   Enable: false
+WebSocket:
+  # The maximum wait time (secend) for the pong message from peer.
+  # If a peer does not respond to the ping message within this time, websocket connection will be closed.
+  # default 60s, 0 means no limit
+  WriteDeadline: 60
+  # The maximum wait time (secend) for the ping message from peer.
+  # If a peer does not send a ping message within this time, websocket connection will be closed.
+  # default 60s, 0 means no limit
+  ReadDeadline: 60
+  # The maximum message size allowed from peer.
+  # If a peer sends a message larger than this, an error will be thrown and websocket connection will be closed.
+  # default: 64MB (32 * 1024 * 1024), 0 means no limit or system limit
+  WriteLimit: 33554432
+  # The maximum message size allowed from peer.
+  # If a peer sends a message larger than this, websocket connection will be closed.
+  # default: 8MB (8 * 1024 * 1024), 0 means no limit or system limit
+  ReadLimit: 8388608
 Auth:
   TokenName: "studio_token"
   AccessSecret: "login_secret"

server/api/studio/internal/config/config.go

@@ -32,6 +32,25 @@ type Config struct {
 		TasksDir  string
 	} `json:",optional"`
 
+	WebSocket struct {
+		// The maximum wait time (secend) for the pong message from peer.
+		// If a peer does not respond to the ping message within this time, websocket will close the connection.
+		// default 60s, 0 means no limit
+		WriteDeadline int64 `json:",default=60"`
+		// The maximum wait time (secend) for the ping message from peer.
+		// If a peer does not send a ping message within this time, websocket will close the connection.
+		// default 60s, 0 means no limit
+		ReadDeadline int64 `json:",default=60"`
+		// The maximum message size allowed from peer.
+		// If a peer sends a message larger than this, a `websocket: write limit exceeded` error will be returned.
+		// default: 64MB (32 * 1024 * 1024), 0 means no limit or system limit
+		WriteLimit int64 `json:",default=33554432"`
+		// The maximum message size allowed from peer.
+		// If a peer sends a message larger than this, websocket will close the connection.
+		// default: 8MB (8 * 1024 * 1024), 0 means no limit or system limit
+		ReadLimit int64 `json:",default=8388608"`
+	} `json:",optional"`
+
 	DB struct {
 		LogLevel                  int    `json:",default=4"`
 		IgnoreRecordNotFoundError bool   `json:",default=true"`

server/api/studio/pkg/ws/utils/client.go

@@ -3,33 +3,19 @@ package utils
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
 	uuid "github.com/satori/go.uuid"
+	"github.com/vesoft-inc/nebula-studio/server/api/studio/internal/config"
 	"github.com/zeromicro/go-zero/core/logx"
 )
 
 const (
-	// Time allowed to write a message to the peer.
-	writeWait = 10 * time.Second
-
-	// Time allowed to read the next pong message from the peer.
-	pongWait = 60 * time.Second
-
-	// Send pings to peer with this period. Must be less than pongWait.
-	pingPeriod = (pongWait * 9) / 10
-
-	// Maximum message size allowed from peer.
-	// Max ngql length can be executed is 4MB
-	maxMessageSize = 8 * 1024 * 1024
-
-	// send buffer size
-	bufSize = 512
-
-	heartbeatRequest = "1"
-
+	bufSize           = 512 // send buffer size
+	heartbeatRequest  = "1"
 	heartbeatResponse = "2"
 )
 
@@ -157,10 +143,15 @@ func (c *Client) Destroy() {
 // reads from this goroutine.
 func (c *Client) readPump() {
 	defer c.Destroy()
-
-	c.Conn.SetReadLimit(maxMessageSize)
-	c.Conn.SetReadDeadline(time.Now().Add(pongWait))
-	c.Conn.SetPongHandler(func(string) error { c.Conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
+	wsConfig := config.GetConfig().WebSocket
+	if wsConfig.ReadLimit > 0 {
+		c.Conn.SetReadLimit(wsConfig.ReadLimit)
+	}
+	if wsConfig.ReadDeadline > 0 {
+		pongWait := time.Duration(wsConfig.ReadDeadline) * time.Second
+		c.Conn.SetReadDeadline(time.Now().Add(pongWait))
+		c.Conn.SetPongHandler(func(string) error { c.Conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
+	}
 	for {
 		_, message, err := c.Conn.ReadMessage()
 		if err != nil {
@@ -191,7 +182,15 @@ func (c *Client) readPump() {
 // application ensures that there is at most one writer to a connection by
 // executing all writes from this goroutine.
 func (c *Client) writePump() {
-	ticker := time.NewTicker(pingPeriod)
+	wsConfig := config.GetConfig().WebSocket
+
+	writeWait := 0 * time.Second
+	ticker := time.NewTicker(30 * time.Second)
+
+	if wsConfig.WriteDeadline > 0 {
+		writeWait = time.Duration(wsConfig.WriteDeadline) * time.Second
+	}
+
 	defer func() {
 		ticker.Stop()
 		c.Destroy()
@@ -200,14 +199,25 @@ func (c *Client) writePump() {
 	for {
 		select {
 		case message, ok := <-c.send:
-			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
 			if !ok {
 				// The hub closed the channel.
 				logx.Errorf("[WebSocket writePump]: c.send length: %v", len(c.send))
 				c.Conn.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(time.Second))
 				return
 			}
 
+			if writeWait > 0 {
+				// no timeout limit if writeWait is 0
+				c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
+			}
+
+			msgLen := len(message)
+			if wsConfig.WriteLimit > 0 && int64(msgLen) > wsConfig.WriteLimit {
+				errMsg := websocket.FormatCloseMessage(websocket.CloseMessageTooBig, fmt.Sprintf("message length (%d) exceeds config limit (%d)", msgLen, wsConfig.WriteLimit))
+				c.Conn.WriteControl(websocket.CloseMessage, errMsg, time.Now().Add(time.Second))
+				return
+			}
+
 			w, err := c.Conn.NextWriter(websocket.TextMessage)
 			if err != nil {
 				logx.Errorf("[WebSocket WriteMessage]: %v", err)
@@ -220,7 +230,10 @@ func (c *Client) writePump() {
 				return
 			}
 		case <-ticker.C:
-			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if writeWait > 0 {
+				// no timeout limit if writeWait is 0
+				c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
+			}
 			if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
 				logx.Errorf("[WebSocket ticker error]: %v", err)
 				return