LokiStack/backend/main.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/gin-gonic/gin"

	"lokistack/config"
	"lokistack/repository"
	"lokistack/routers"
	"lokistack/seed"
	"lokistack/services"
)

// version 由构建注入（-ldflags "-X main.version=..."），未注入时为 dev
var version = "2.0.0"

// fileExists 判断文件是否存在（用于"是否首次部署"这类轻量判断）。
func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

// runCLI 提供不依赖 Web 的运维命令（备份 CronJob、恢复、自检都走它）。
//
//	lokistack backup --output /app/data/backups/xxx.tar.gz
//	lokistack restore <备份包路径>          # 校验 + 立即应用（会先备份当前 DB）
//	lokistack self-check                    # 打印 JSON 自检结果
func runCLI(args []string) {
	switch args[0] {
	case "backup":
		out := ""
		for i := 1; i < len(args); i++ {
			if args[i] == "--output" && i+1 < len(args) {
				out = args[i+1]
			}
		}
		if out == "" {
			out = "/app/data/backups/lokistack-" + time.Now().Format("20060102-150405") + ".tar.gz"
		}
		if err := os.MkdirAll(filepath.Dir(out), 0o700); err != nil {
			fmt.Fprintf(os.Stderr, "创建备份目录失败: %v\n", err)
			os.Exit(1)
		}
		f, err := os.Create(out)
		if err != nil {
			fmt.Fprintf(os.Stderr, "创建备份文件失败: %v\n", err)
			os.Exit(1)
		}
		defer f.Close()
		m, err := services.CreateBackup(f)
		if err != nil {
			fmt.Fprintf(os.Stderr, "备份失败: %v\n", err)
			os.Exit(1)
		}
		fmt.Printf("备份完成: %s\n记录数: %d\n密钥指纹: %s\n", out, m.RecordCount, m.KeyFP)

	case "restore":
		if len(args) < 2 {
			fmt.Fprintln(os.Stderr, "用法: lokistack restore <备份包.tar.gz>")
			os.Exit(1)
		}
		f, err := os.Open(args[1])
		if err != nil {
			fmt.Fprintf(os.Stderr, "打开备份失败: %v\n", err)
			os.Exit(1)
		}
		defer f.Close()
		res, err := services.ValidateAndStageBackup(f)
		if err != nil {
			fmt.Fprintf(os.Stderr, "校验失败: %v\n", err)
			os.Exit(1)
		}
		if res.Warning != "" {
			fmt.Fprintf(os.Stderr, "⚠️  %s\n", res.Warning)
		}
		path, err := services.ApplyStagedBackup()
		if err != nil {
			fmt.Fprintf(os.Stderr, "应用失败: %v\n", err)
			os.Exit(1)
		}
		fmt.Printf("恢复完成: %s（备份 %d 条记录）\n请重启平台进程使新数据生效。\n", path, res.RowsInBackup)

	case "self-check":
		b, _ := json.MarshalIndent(services.SelfCheck(), "", "  ")
		fmt.Println(string(b))

	case "version":
		fmt.Println(version)

	default:
		fmt.Fprintf(os.Stderr, "未知命令 %q。可用: backup / restore / self-check / version\n", args[0])
		os.Exit(1)
	}
}

func main() {
	// 运维命令行：支持在容器/CronJob 里不经 Web 完成备份/恢复/自检
	if len(os.Args) > 1 {
		runCLI(os.Args[1:])
		return
	}

	cfg := config.GetConfig()

	// gin 默认是 debug 模式：每次启动会把上百行路由表 dump 到日志里，
	// 且每个请求有额外开销。生产默认切 release，需要调试时显式设
	// DEBUG=true 或 GIN_MODE=debug 覆盖。
	if os.Getenv("GIN_MODE") == "" && !strings.EqualFold(os.Getenv("DEBUG"), "true") {
		gin.SetMode(gin.ReleaseMode)
	}

	log.Println("Initializing database...")
	if err := repository.InitDB(); err != nil {
		log.Fatalf("Failed to init DB: %v", err)
	}

	// 指标注册要早于任何可能埋点的组件（middleware / alert engine）
	services.InitMetrics(version)
	log.Println("Metrics enabled at /api/v1/metrics")

	log.Println("Seeding default data...")
	seed.Seed()

	// 安全迁移：历史明文 kubeconfig → 加密存储，并清理明文文件
	services.MigrateClusterKubeconfigs()

	// 孤儿部署状态恢复:进程重启会杀掉进行中的部署 goroutine,
	// 残留的 "deploying" 状态将永远卡住(UI 无停止入口),启动时统一标记为失败。
	repository.GetDB().RecoverOrphanDeployments()

	services.StartAlertEngine()
	// 每日日志量落库：Loki retention 只有 24h，历史趋势必须自己留
	services.StartDailyVolumeCapture()
	defer services.StopAlertEngine()

	r := routers.SetupRouter()
	srv := &http.Server{
		Addr:         ":" + cfg.ServerPort,
		Handler:      r,
		ReadTimeout:  30 * time.Second,
		WriteTimeout: 130 * time.Second,
		IdleTimeout:  120 * time.Second,
	}

	// 监听中断/终止信号，实现优雅关停：先停 HTTP（拒绝新请求、处理完在途请求），
	// 再 flush 数据库，避免防抖/异步保存丢失最近写入。
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	go func() {
		log.Printf("lokistack starting on port %s", cfg.ServerPort)
		log.Printf("API: http://localhost:%s/api/v1", cfg.ServerPort)
		// 初始密码只在**真正首次部署**时提示：
		// 旧实现无条件打印这句话，于是每次重启日志里都出现"首次部署"，
		// 排查问题时极易被误导成"平台重新初始化了、数据是不是丢了"。
		// 判据用密码文件是否存在（admin 改密后该文件会被删除）。
		if pwdFile := filepath.Join(cfg.DataDir, "initial-admin-password.txt"); fileExists(pwdFile) {
			log.Printf("首次部署：admin 初始密码见 %s（改密后该文件会自动删除）", pwdFile)
		}
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Fatalf("Server failed: %v", err)
		}
	}()

	<-ctx.Done()
	log.Println("Shutdown signal received, draining connections...")

	shutdownCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	if err := srv.Shutdown(shutdownCtx); err != nil {
		log.Printf("HTTP shutdown error: %v", err)
	}

	log.Println("Flushing database...")
	if err := repository.Shutdown(); err != nil {
		log.Printf("DB shutdown error: %v", err)
	}
	log.Println("lokistack stopped")
}