155 lines
3.6 KiB
Go
155 lines
3.6 KiB
Go
package main
|
||
|
||
import (
|
||
"flag"
|
||
"log"
|
||
"os"
|
||
"runtime/debug"
|
||
"time"
|
||
"xk-websocket/controller"
|
||
"xk-websocket/route"
|
||
"xk-websocket/utils"
|
||
|
||
"gorm.io/driver/mysql"
|
||
"gorm.io/gorm"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
/**
|
||
* ensureLogDir
|
||
* 功能：检查指定的日志目录是否存在，如果不存在则尝试创建该目录。
|
||
* @param logDir string 日志目录路径
|
||
*/
|
||
func ensureLogDir(logDir string) {
|
||
if _, err := os.Stat(logDir); os.IsNotExist(err) {
|
||
log.Printf("📂 日志目录不存在，正在创建: %s", logDir)
|
||
if err := os.MkdirAll(logDir, 0755); err != nil {
|
||
log.Fatalf("❌ 无法创建日志目录: %v", err)
|
||
}
|
||
log.Printf("✅ 日志目录创建成功: %s", logDir)
|
||
}
|
||
}
|
||
|
||
/**
|
||
* setupLogger
|
||
* 功能：配置全局日志系统。
|
||
* 逻辑：
|
||
* 1. 读取环境变量配置日志目录和轮转频率。
|
||
* 2. 创建自定义的 DailyFileWriter 以支持按天/按小时分割日志。
|
||
* 3. 将标准 log 的输出重定向到文件系统。
|
||
*/
|
||
func setupLogger() {
|
||
// 设置日志目录（优先级：环境变量 > 默认值）
|
||
logDir := os.Getenv("LOG_DIR")
|
||
if logDir == "" {
|
||
logDir = "logs" // 默认日志目录
|
||
}
|
||
|
||
// 确保日志目录存在
|
||
ensureLogDir(logDir)
|
||
|
||
// 设置日志轮转频率
|
||
logFreq := os.Getenv("LOG_ROTATE_FREQ")
|
||
if logFreq == "" {
|
||
logFreq = "daily" // 默认每天轮转
|
||
}
|
||
|
||
// 创建文件日志写入器
|
||
logger := log.Default()
|
||
fileWriter := utils.NewDailyFileWriter(logDir, logFreq, logger, "xk-websocket")
|
||
|
||
// 同时输出到控制台和文件
|
||
log.SetOutput(fileWriter)
|
||
}
|
||
|
||
/**
|
||
* initDB
|
||
* 功能：初始化 MySQL 数据库连接。
|
||
* 返回：*gorm.DB 数据库连接实例
|
||
*/
|
||
func initDB() *gorm.DB {
|
||
dsn := utils.GetEnv("DB_DSN", "root:password@tcp(localhost:3306)/xk_chat?charset=utf8mb4&parseTime=True&loc=Local")
|
||
db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
|
||
if err != nil {
|
||
log.Fatalf("❌ 数据库连接失败: %v", err)
|
||
}
|
||
|
||
//// 自动迁移表结构
|
||
//if err := db.AutoMigrate(&models.XkChatMessage{}); err != nil {
|
||
// log.Fatalf("❌ 数据库迁移失败: %v", err)
|
||
//}
|
||
|
||
log.Println("✅ 数据库连接成功")
|
||
return db
|
||
}
|
||
|
||
/**
|
||
* main
|
||
* 功能：程序主入口。
|
||
* 流程：
|
||
* 1. 初始化日志。
|
||
* 2. 解析命令行参数（端口、节点ID）。
|
||
* 3. 初始化数据库和WebSocket控制器。
|
||
* 4. 启动 Redis 订阅（用于集群消息同步）。
|
||
* 5. 注册 HTTP/WebSocket 路由。
|
||
* 6. 启动 Web Server。
|
||
*/
|
||
func main() {
|
||
// 初始化日志系统（必须放在最先）
|
||
setupLogger()
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
log.Printf("❌ 程序崩溃: %v", r)
|
||
debug.PrintStack()
|
||
log.Println("程序将在10秒后退出...")
|
||
time.Sleep(10 * time.Second)
|
||
}
|
||
}()
|
||
|
||
log.Println("🚀 启动 WebSocket 服务...")
|
||
|
||
// 添加命令行参数解析
|
||
var port, nodeID string
|
||
flag.StringVar(&port, "port", "", "运行端口号")
|
||
flag.StringVar(&nodeID, "nodeId", "", "节点标识符")
|
||
flag.Parse()
|
||
|
||
// 初始化数据库
|
||
db := initDB()
|
||
|
||
// 初始化控制器
|
||
wsCtrl := controller.NewWebSocketController(db)
|
||
|
||
// 设置参数优先级
|
||
if nodeID != "" {
|
||
wsCtrl.NodeID = nodeID
|
||
}
|
||
if port != "" {
|
||
wsCtrl.Port = port
|
||
}
|
||
|
||
// 配置系统
|
||
wsCtrl.ConfigureSystem()
|
||
wsCtrl.InitRedisClient()
|
||
wsCtrl.PrintStartupInfo()
|
||
|
||
// 创建Gin路由器
|
||
router := gin.Default()
|
||
|
||
// 设置路由
|
||
route.SetupRoutes(router, wsCtrl)
|
||
|
||
// 启动Redis消息订阅
|
||
go wsCtrl.SubscribeToRedis()
|
||
|
||
// 启动HTTP服务
|
||
addr := "0.0.0.0:" + wsCtrl.Port
|
||
log.Printf("🌐 监听地址: %s", addr)
|
||
log.Printf("🔌 WebSocket连接地址: ws://%s/ws", addr)
|
||
|
||
if err := router.Run(addr); err != nil {
|
||
log.Fatalf("❌ 服务启动失败: %v", err)
|
||
}
|
||
}
|