gotidb/cmd/server/main.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
		}
	}

	// 创建存储引擎
	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")
}