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}>

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 will close the connection.
+  # 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 will close the connection.
+  # default 60s, 0 means no limit
+  ReadDeadline: 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: 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: 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

@@ -8,28 +8,13 @@ import (
 
 	"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 +142,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 +181,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,7 +198,10 @@ func (c *Client) writePump() {
 	for {
 		select {
 		case message, ok := <-c.send:
-			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 !ok {
 				// The hub closed the channel.
 				logx.Errorf("[WebSocket writePump]: c.send length: %v", len(c.send))
@@ -220,7 +221,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