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://localhost:4222", "NATS服务器地址") persistenceType = flag.String("persistence", "none", "持久化类型 (none, wal)") persistenceDir = flag.String("persistence-dir", "./data", "持久化目录") syncEvery = flag.Int("sync-every", 100, "每写入多少条数据同步一次") ) func main() { flag.Parse() // 创建存储引擎 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) // 创建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) } }() // 等待信号 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) } } // 关闭数据管理器 if err := dataManager.Close(); err != nil { log.Printf("Failed to close data manager: %v", err) } log.Println("Shutdown complete") }