添加选项

websocket/websocket.go

@@ -3,6 +3,7 @@ package websocket
 import (
 	"github.com/olahol/melody"
 	"net/http"
+	"time"
 )
 
 const (
@@ -19,9 +20,37 @@ type HandleMessageFunc func(groupID string, message []byte, context any)
 
 var managerInstance *Manager
 
-func Init() {
+func Init(opts ...InitOption) {
 	if managerInstance == nil {
 		melodyInstance := melody.New()
+
+		options := new(InitOptions)
+		for _, opt := range opts {
+			opt(options)
+		}
+
+		if options.writeWaitSec != 0 {
+			melodyInstance.Config.WriteWait = time.Duration(options.writeWaitSec) * time.Second
+		}
+
+		if options.pongWaitSec != 0 {
+			melodyInstance.Config.PongWait = time.Duration(options.pongWaitSec) * time.Second
+		}
+
+		if options.pingPeriodSec != 0 {
+			melodyInstance.Config.PingPeriod = time.Duration(options.pingPeriodSec) * time.Second
+		}
+
+		if options.maxMessageSize != 0 {
+			melodyInstance.Config.MaxMessageSize = options.maxMessageSize
+		}
+
+		if options.messageBufferSize != 0 {
+			melodyInstance.Config.MessageBufferSize = options.messageBufferSize
+		}
+
+		melodyInstance.Config.ConcurrentMessageHandling = options.concurrentMessageHandling
+
 		melodyInstance.Upgrader.CheckOrigin = func(r *http.Request) bool { return true }
 		managerInstance = &Manager{melodyInstance: melodyInstance}
 	}
@@ -126,6 +155,53 @@ func (m *Manager) BroadCast(groupID string, msg []byte) error {
 	})
 }
 
+type InitOption func(*InitOptions)
+
+type InitOptions struct {
+	writeWaitSec              int64
+	pongWaitSec               int64
+	pingPeriodSec             int64
+	maxMessageSize            int64
+	messageBufferSize         int
+	concurrentMessageHandling bool
+}
+
+func InitWithWriteWaitSec(writeWaitSec int64) InitOption {
+	return func(options *InitOptions) {
+		options.writeWaitSec = writeWaitSec
+	}
+}
+
+func InitWithPongWaitSec(pongWaitSec int64) InitOption {
+	return func(options *InitOptions) {
+		options.pongWaitSec = pongWaitSec
+	}
+}
+
+func InitWithPingPeriodSec(pingPeriodSec int64) InitOption {
+	return func(options *InitOptions) {
+		options.pingPeriodSec = pingPeriodSec
+	}
+}
+
+func InitWithMaxMessageSize(maxMessageSize int64) InitOption {
+	return func(options *InitOptions) {
+		options.maxMessageSize = maxMessageSize
+	}
+}
+
+func InitWithMaxMessageBufferSize(messageBufferSize int) InitOption {
+	return func(options *InitOptions) {
+		options.messageBufferSize = messageBufferSize
+	}
+}
+
+func InitWithConcurrentMessageHandling(concurrentMessageHandling bool) InitOption {
+	return func(options *InitOptions) {
+		options.concurrentMessageHandling = concurrentMessageHandling
+	}
+}
+
 type ConnectionOption func(sessionMap map[string]any)
 
 func WithConnectionContext(context any) ConnectionOption {