From 8a19bed17bc908a05a9f0bf1cb5aeb75f6763060 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E5=BE=B7=E6=BB=A1?= <go_developer@163.com>
Date: Thu, 22 Apr 2021 22:03:17 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=AF=B9=E9=95=BF=E8=BF=9E?=
 =?UTF-8?q?=E6=8E=A5=E6=B6=88=E6=81=AF=E4=BD=93=E4=BB=A5=E5=8F=8A=E6=8E=A2?=
 =?UTF-8?q?=E6=B4=BB=E9=85=8D=E7=BD=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 config/server.go | 87 ++++++++++++++++++++++++++++++++++++++++++++++++
 construct.go     | 14 ++++++++
 2 files changed, 101 insertions(+)

diff --git a/config/server.go b/config/server.go
index 55816a9..f0192dc 100644
--- a/config/server.go
+++ b/config/server.go
@@ -8,9 +8,12 @@
 package config
 
 import (
+	"time"
+
 	"github.com/go-developer/gopkg/logger"
 	"github.com/go-developer/websocket/storage"
 	"go.uber.org/zap/zapcore"
+	"gopkg.in/olahol/melody.v1"
 )
 
 const (
@@ -43,6 +46,16 @@ const (
 	DefaultStoreConnection = true
 	// DefaultEnablePprof 默认关闭pprof
 	DefaultEnablePprof = false
+	// DefaultMaxMessageSize 默认的最大消息大小
+	DefaultMaxMessageSize = 4096
+	// DefaultMaxMessageBufferSize 默认的消息缓冲区大小
+	DefaultMaxMessageBufferSize = 8192
+	// DefaultWriteWait 默认的消息写入等待时间
+	DefaultWriteWait = 10
+	// DefaultPongWait 默认的等待响应时间
+	DefaultPongWait = 10
+	// DefaultPingPeriod 多久探活一次
+	DefaultPingPeriod = 10
 )
 
 // WSServerConfig WS-Server的配置
@@ -62,6 +75,7 @@ type WSServerConfig struct {
 	ConnectionManager storage.IConnection     // 连接管理实例
 	EnablePprof       bool                    // 开启pprof, 默认关闭
 	PprofPort         int                     // pprof监听的端口
+	melody.Config                             // 长连接配置
 }
 
 // SetWSServerConfig 设置WS-Server的配置
@@ -121,6 +135,76 @@ func DisableStoreConnection() SetWSServerConfig {
 	}
 }
 
+// SetMaxMessageSize 限制消息大小
+//
+// Author : go_developer@163.com<张德满>
+//
+// Date : 9:29 下午 2021/4/22
+func SetMaxMessageSize(maxMessageSize int64) SetWSServerConfig {
+	return func(wsc *WSServerConfig) {
+		if maxMessageSize <= 0 {
+			maxMessageSize = DefaultMaxMessageSize
+		}
+		wsc.MaxMessageSize = maxMessageSize
+	}
+}
+
+// SetMaxMessageBufferSize 设置消息缓冲区
+//
+// Author : go_developer@163.com<张德满>
+//
+// Date : 9:39 下午 2021/4/22
+func SetMaxMessageBufferSize(maxMessageBufferSize int) SetWSServerConfig {
+	return func(wsc *WSServerConfig) {
+		if maxMessageBufferSize <= 0 {
+			maxMessageBufferSize = DefaultMaxMessageSize
+		}
+		wsc.MessageBufferSize = maxMessageBufferSize
+	}
+}
+
+// SetWriteWait 设置消息写入的等待时间, 单位 : 秒
+//
+// Author : go_developer@163.com<张德满>
+//
+// Date : 9:44 下午 2021/4/22
+func SetWriteWait(writeWaitTime int) SetWSServerConfig {
+	return func(wsc *WSServerConfig) {
+		if writeWaitTime <= 0 {
+			writeWaitTime = DefaultWriteWait
+		}
+		wsc.WriteWait = time.Duration(writeWaitTime) * time.Second
+	}
+}
+
+// SetPongWait 设置探活时,等待响应时间,单位 : 秒
+//
+// Author : go_developer@163.com<张德满>
+//
+// Date : 9:55 下午 2021/4/22
+func SetPongWait(pongWait int) SetWSServerConfig {
+	return func(wsc *WSServerConfig) {
+		if pongWait <= 0 {
+			pongWait = DefaultPongWait
+		}
+		wsc.PongWait = time.Duration(pongWait) * time.Second
+	}
+}
+
+// SetPingPeriod 设置探活时间间隔
+//
+// Author : go_developer@163.com<张德满>
+//
+// Date : 10:00 下午 2021/4/22
+func SetPingPeriod(pingPeriod int) SetWSServerConfig {
+	return func(wsc *WSServerConfig) {
+		if pingPeriod <= 0 {
+			pingPeriod = DefaultPingPeriod
+		}
+		wsc.PingPeriod = time.Duration(pingPeriod) * time.Second
+	}
+}
+
 // NewWSServerConfig 生成新的WS-Server配置
 //
 // Author : go_developer@163.com<张德满>
@@ -136,7 +220,10 @@ func NewWSServerConfig(optionList ...SetWSServerConfig) *WSServerConfig {
 		LogSplitInterval:  DefaultLogSplitInterval,
 		StoreConnection:   DefaultStoreConnection,
 		ConnectionManager: storage.NewDefaultConnectionManager(),
+		EnablePprof:       DefaultEnablePprof,
 	}
+	c.MaxMessageSize = DefaultMaxMessageSize
+	c.MessageBufferSize = DefaultMaxMessageBufferSize
 	for _, o := range optionList {
 		o(c)
 	}
diff --git a/construct.go b/construct.go
index c4a9676..5b28ebf 100644
--- a/construct.go
+++ b/construct.go
@@ -134,7 +134,21 @@ func initServer(wsInstance abstract.IWebsocket) {
 				loggerInstanceTable[wsInstance.GetModuleFlag()] = loggerInstance
 			}
 		}
+		// 对长连接进行配置
+		s.wsServer.Config.MaxMessageSize = s.conf.MaxMessageSize
+		s.wsServer.Config.MessageBufferSize = s.conf.MessageBufferSize
+		s.wsServer.Config.WriteWait = s.conf.WriteWait
+		s.wsServer.Config.PongWait = s.conf.PongWait
+		s.wsServer.Config.PingPeriod = s.conf.PingPeriod
+
 		wsServerTable[wsInstance.GetServerPort()][wsInstance.GetModuleFlag()] = s
+		/*melody.Config{
+			WriteWait:         0,
+			PongWait:          0,
+			PingPeriod:        0,
+			MaxMessageSize:    0,
+			MessageBufferSize: 0,
+		}*/
 	}
 
 	// 初始化指令存储表