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://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")
}