websocket/config/server.go

// Package config ...
//
// Description : WS-Server 相关配置
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2021-04-17 2:32 下午
package config

import (
	"git.zhangdeman.cn/zhangdeman/consts"
	"time"

	"git.zhangdeman.cn/zhangdeman/websocket/storage"
	"go.uber.org/zap/zapcore"
	"gopkg.in/olahol/melody.v1"
)

const (
	// RunModeProduct 生产环境
	RunModeProduct = "product"
	// RunModeDebug debug环境
	RunModeDebug = "debug"
)

const (
	// LogSplitIntervalHour 按小时切割日志
	LogSplitIntervalHour = "hour"
	// LogSplitIntervalDay 按天切割日志
	LogSplitIntervalDay = "day"
)

// 定义相关默认值
const (
	// DefaultLogEnable 默认关闭日志
	DefaultLogEnable = false
	// DefaultLogConsole 默认开启控制台输出
	DefaultLogConsole = false
	// DefaultMode 默认为Debug模式
	DefaultMode = RunModeDebug
	// DefaultLogLevel 默认的日志级别
	DefaultLogLevel = zapcore.DebugLevel
	// DefaultLogSplitInterval 默认的日志切割时间
	DefaultLogSplitInterval = consts.LogSplitHour
	// DefaultStoreConnection 默认存储连接
	DefaultStoreConnection = true
	// DefaultEnablePprof 默认关闭pprof
	DefaultEnablePprof = false
	// DefaultMaxMessageSize 默认的最大消息大小
	DefaultMaxMessageSize = 4096
	// DefaultMaxMessageBufferSize 默认的消息缓冲区大小
	DefaultMaxMessageBufferSize = 8192
	// DefaultWriteWait 默认的消息写入等待时间
	DefaultWriteWait = 10
	// DefaultPongWait 默认的等待响应时间
	DefaultPongWait = 10
	// DefaultPingPeriod 多久探活一次
	DefaultPingPeriod = 10
)

// WSServerConfig WS-Server的配置
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 7:02 下午 2021/4/17
type WSServerConfig struct {
	Mode              string              // 运行模式
	LogEnable         bool                // 开启日志
	LogConsole        bool                // 开启控制台日志输出
	LogPath           string              // 日志路径
	LogFile           string              // 日志文件名
	LogLevel          zapcore.Level       // 日志等级
	LogSplitInterval  string              // 日至切割的时间间隔
	StoreConnection   bool                // 存储连接
	ConnectionManager storage.IConnection // 连接管理实例
	EnablePprof       bool                // 开启pprof, 默认关闭
	PprofPort         int                 // pprof监听的端口
	melody.Config                         // 长连接配置
}

// SetWSServerConfig 设置WS-Server的配置
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 7:03 下午 2021/4/17
type SetWSServerConfig func(wsc *WSServerConfig)

// SetWSServerLogEnable 开启日志记录
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 7:25 下午 2021/4/17
func SetWSServerLogEnable(logPath string, logFile string, logLevel zapcore.Level, splitInterval string) SetWSServerConfig {
	return func(wsc *WSServerConfig) {
		if len(logPath) == 0 {
			logPath = "./logs"
		}

		wsc.LogEnable = true
		wsc.LogPath = logPath
		wsc.LogFile = logFile
		wsc.LogLevel = logLevel
		wsc.LogSplitInterval = splitInterval
	}
}

// EnableConsoleLog 开启控制台日志输出
//
// Author : zhangdeman001@ke.com<白茶清欢>
//
// Date : 11:25 下午 2021/4/22
func EnableConsoleLog() SetWSServerConfig {
	return func(wsc *WSServerConfig) {
		wsc.LogConsole = true
	}
}

// EnablePprof 开启PProf,由于多模块共享一个进程,任意一个模块开启,就认为是开启
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 11:24 下午 2021/4/18
func EnablePprof(pprofPort int) SetWSServerConfig {
	return func(wsc *WSServerConfig) {
		wsc.EnablePprof = true
		wsc.PprofPort = pprofPort
	}
}

// DisableStoreConnection 禁用连接存储
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 11:10 下午 2021/4/17
func DisableStoreConnection() SetWSServerConfig {
	return func(wsc *WSServerConfig) {
		wsc.StoreConnection = true
	}
}

// SetConnectionManager 连接管理实例
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:48 下午 2021/4/22
func SetConnectionManager(manager storage.IConnection) SetWSServerConfig {
	return func(wsc *WSServerConfig) {
		if nil == manager {
			manager = storage.NewDefaultConnectionManager()
		}
		wsc.ConnectionManager = manager
	}
}

// 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<白茶清欢>
//
// Date : 7:21 下午 2021/4/17
func NewWSServerConfig(optionList ...SetWSServerConfig) *WSServerConfig {
	c := &WSServerConfig{
		Mode:              DefaultMode,
		LogEnable:         DefaultLogEnable,
		LogConsole:        DefaultLogConsole,
		LogPath:           "",
		LogLevel:          DefaultLogLevel,
		LogSplitInterval:  DefaultLogSplitInterval,
		StoreConnection:   DefaultStoreConnection,
		ConnectionManager: storage.NewDefaultConnectionManager(),
		EnablePprof:       DefaultEnablePprof,
	}
	c.MaxMessageSize = DefaultMaxMessageSize
	c.MessageBufferSize = DefaultMaxMessageBufferSize
	c.PingPeriod = time.Duration(DefaultPingPeriod) * time.Second
	c.WriteWait = time.Duration(DefaultWriteWait) * time.Second
	c.PongWait = time.Duration(DefaultPongWait) * time.Second
	for _, o := range optionList {
		o(c)
	}
	return c
}