package main import ( "context" "flag" "log" "os" "os/signal" "syscall" "time" "git.pyer.club/kingecg/gotidb/pkg/api" "git.pyer.club/kingecg/gotidb/pkg/manager" "git.pyer.club/kingecg/gotidb/pkg/messaging" "git.pyer.club/kingecg/gotidb/pkg/model" "git.pyer.club/kingecg/gotidb/pkg/monitoring" "git.pyer.club/kingecg/gotidb/pkg/storage" ) var ( restAddr = flag.String("rest-addr", ":8080", "REST API服务地址") wsAddr = flag.String("ws-addr", ":8081", "WebSocket服务地址") metricsAddr = flag.String("metrics-addr", ":8082", "指标服务地址") natsURL = flag.String("nats-url", "", "NATS服务器地址") persistenceType = flag.String("persistence", "none", "持久化类型 (none, wal)") persistenceDir = flag.String("persistence-dir", "./data", "持久化目录") syncEvery = flag.Int("sync-every", 100, "每写入多少条数据同步一次") configPath = flag.String("config", "config.yaml", "配置文件路径") // 定义配置QUIC服务的命令行参数 quicAddr = flag.String("quic-addr", ":8083", "QUIC服务地址") genSampleConfig = flag.Bool("gen-sample-config", false, "生成示例配置文件") ) func main() { if *genSampleConfig { err := GenerateSampleConfig("./config.yaml.sample") if err != nil { log.Fatalf("生成示例配置文件失败: %v", err) } log.Println("示例配置文件已生成") return } flag.Parse() // 保存命令行原始值,用于后续判断是否被用户显式设置 originalRestAddr := *restAddr originalWsAddr := *wsAddr originalMetricsAddr := *metricsAddr originalQuicAddr := *quicAddr originalNatsURL := *natsURL originalPersistenceType := *persistenceType originalPersistenceDir := *persistenceDir originalSyncEvery := *syncEvery var config *Config if *configPath != "" { var err error config, err = LoadConfig(*configPath) if err != nil { log.Fatalf("Failed to load config: %v", err) } // 只有当命令行参数为默认值时,才使用配置文件中的值 if *restAddr == originalRestAddr { restAddr = &config.RestAddr } if *wsAddr == originalWsAddr { wsAddr = &config.WsAddr } if *metricsAddr == originalMetricsAddr { metricsAddr = &config.MetricsAddr } if *quicAddr == originalQuicAddr { quicAddr = &config.QuicAddr } if *natsURL == originalNatsURL { natsURL = &config.NATSURL } if *persistenceType == originalPersistenceType { persistenceType = &config.PersistenceType } if *persistenceDir == originalPersistenceDir { persistenceDir = &config.PersistenceDir } if *syncEvery == originalSyncEvery { syncEvery = &config.SyncEvery } } // 创建存储引擎 engine := storage.NewMemoryEngine() // 如果启用了持久化,配置持久化 if *persistenceType != "none" { persistenceConfig := storage.PersistenceConfig{ Type: storage.PersistenceType(*persistenceType), Directory: *persistenceDir, SyncEvery: *syncEvery, } if err := engine.EnablePersistence(persistenceConfig); err != nil { log.Fatalf("Failed to enable persistence: %v", err) } log.Printf("Persistence enabled: type=%s, directory=%s", *persistenceType, *persistenceDir) } // 创建数据管理器 dataManager := manager.NewDataManager(engine) // 创建指标收集器 metrics := monitoring.NewMetrics() // 创建REST API服务 restServer := api.NewRESTServer(dataManager) // 创建WebSocket服务 wsServer := api.NewWebSocketServer(dataManager) // 创建QUIC服务器 var quicServer *api.QUICServer var quicConfig = DefaultQuicConfig() // 使用默认配置 if config != nil && config.QuicConfig != nil { quicConfig = config.QuicConfig // 如果配置文件中有配置,则使用配置文件中的配置 } quicServer, err := api.NewQUICServer(dataManager, quicConfig) if err != nil { log.Printf("Failed to create QUIC server: %v", err) log.Println("Continuing without QUIC server") quicServer = nil } // 创建NATS消息系统 natsConfig := messaging.NATSConfig{ URL: *natsURL, } nats, err := messaging.NewNATSMessaging(natsConfig) if err != nil { log.Printf("Failed to create NATS messaging: %v", err) log.Println("Continuing without NATS messaging") } else { // 订阅消息 err = nats.Subscribe(func(msg messaging.DataMessage) error { // 创建数据点 id := manager.CreateDataPoint( msg.DeviceID, msg.MetricCode, msg.Labels, msg.Value, ) // 写入数据 value := model.DataValue{ Timestamp: msg.Timestamp, Value: msg.Value, } return dataManager.Write(context.Background(), id, value) }) if err != nil { log.Printf("Failed to subscribe to NATS: %v", err) } else { log.Println("NATS messaging enabled") } } // 启动服务 go func() { log.Printf("Starting REST API server on %s", *restAddr) if err := restServer.Start(*restAddr); err != nil { log.Fatalf("Failed to start REST API server: %v", err) } }() go func() { log.Printf("Starting WebSocket server on %s", *wsAddr) if err := wsServer.Start(*wsAddr); err != nil { log.Fatalf("Failed to start WebSocket server: %v", err) } }() go func() { log.Printf("Starting metrics server on %s", *metricsAddr) if err := metrics.StartServer(*metricsAddr); err != nil { log.Fatalf("Failed to start metrics server: %v", err) } }() // 启动QUIC服务 if quicServer != nil { go func() { log.Printf("Starting QUIC server on %s", *quicAddr) if err := quicServer.Start(*quicAddr); err != nil { log.Printf("Failed to start QUIC server: %v", err) } }() } // 等待信号 sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) <-sigCh // 优雅关闭 log.Println("Shutting down...") // 创建关闭上下文 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() // 关闭REST API服务 if err := restServer.Stop(ctx); err != nil { log.Printf("Failed to stop REST API server: %v", err) } // 关闭WebSocket服务 if err := wsServer.Stop(ctx); err != nil { log.Printf("Failed to stop WebSocket server: %v", err) } // 关闭NATS消息系统 if nats != nil { if err := nats.Close(); err != nil { log.Printf("Failed to close NATS messaging: %v", err) } } // 关闭QUIC服务 if quicServer != nil { if err := quicServer.Stop(ctx); err != nil { log.Printf("Failed to stop QUIC server: %v", err) } } // 关闭数据管理器 if err := dataManager.Close(); err != nil { log.Printf("Failed to close data manager: %v", err) } log.Println("Shutdown complete") }