feat: customization of WebSocket-related parameters

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