257 lines
6.8 KiB
Go
257 lines
6.8 KiB
Go
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
|
||
}
|
||
}
|
||
|
||
// 创建存储引擎配置
|
||
engineConfig := storage.EngineConfig{
|
||
PersistenceType: storage.PersistenceType(*persistenceType),
|
||
PersistenceDir: *persistenceDir,
|
||
SyncEvery: *syncEvery,
|
||
}
|
||
|
||
// 如果是BoltDB，添加BoltDB特有配置
|
||
if *persistenceType == "boltdb" {
|
||
if config != nil {
|
||
engineConfig.BoltDBFilename = config.BoltDBFilename
|
||
engineConfig.BoltDBOptions = &storage.BoltDBConfig{
|
||
BucketSize: config.BoltDBBucketSize,
|
||
}
|
||
} else {
|
||
// 使用默认值
|
||
engineConfig.BoltDBFilename = "gotidb.db"
|
||
engineConfig.BoltDBOptions = &storage.BoltDBConfig{
|
||
BucketSize: 30,
|
||
}
|
||
}
|
||
}
|
||
|
||
// 使用工厂方法创建存储引擎
|
||
engine, err := storage.NewStorageEngine(engineConfig)
|
||
if err != nil {
|
||
log.Fatalf("Failed to create storage engine: %v", err)
|
||
}
|
||
|
||
log.Printf("Storage engine created: 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")
|
||
}
|