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 / buildTime / gitCommit 由构建注入（build.sh / Dockerfile 的 -ldflags -X main.*），
// 未注入时为 dev，避免与前端版本号互相矛盾
var (
	version   = "dev"
	buildTime = "unknown"
	gitCommit = "unknown"
)

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

// runCLI 提供不依赖 Web 的运维命令（备份 CronJob、恢复、自检都走它）。
//
//	lokistack self-check                    # 打印 JSON 自检结果

func strDefault(v, d string) string {
	if strings.TrimSpace(v) == "" {
		return d
	}
	return v
}

func runCLI(args []string) {
	switch args[0] {
	case "self-check":
		out, _ := json.MarshalIndent(services.SelfCheck(), "", "  ")
		fmt.Println(string(out))
	case "license":
		// 发行方工具:签发 Pro 许可码(需发行密钥构建;默认密钥仅供开发自测)
		flags := map[string]string{}
		for _, a := range args[1:] {
			if rest, ok := strings.CutPrefix(a, "--"); ok {
				if k, v, ok2 := strings.Cut(rest, "="); ok2 {
					flags[k] = v
				}
			}
		}
		getInt := func(k string, d int) int {
			if v := flags[k]; v != "" {
				n := 0
				for _, ch := range v {
					if ch < '0' || ch > '9' {
						n = 0
						break
					}
					n = n*10 + int(ch-'0')
				}
				if n > 0 {
					return n
				}
			}
			return d
		}
		days := getInt("days", 0)
		key, err := services.GenerateLicenseKey(
			"pro",
			strDefault(flags["edition"], "personal"),
			strDefault(flags["customer"], "unnamed"),
			getInt("seats", 0), getInt("instances", 1), days,
		)
		if err != nil {
			fmt.Fprintln(os.Stderr, "签发失败:", err)
			os.Exit(1)
		}
		expiry := "永久"
		if days > 0 {
			expiry = fmt.Sprintf("%d 天", days)
		}
		fmt.Printf("许可码(%s / %s / %s):\n%s\n", strDefault(flags["edition"], "personal"), strDefault(flags["customer"], "unnamed"), expiry, key)
	case "version":
		fmt.Printf("%s (commit %s, built %s)\n", version, gitCommit, buildTime)
	default:
		fmt.Fprintf(os.Stderr, "未知命令 %q。可用: 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）
	routers.Version = version
	services.InitMetrics(version)
	log.Println("Metrics enabled at /api/v1/metrics")

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

	// 清理旧版本删除部署后残留的托管数据源（幂等）
	services.CleanupOrphanManagedLogSources()
	// 异步重新探测运行中部署的数据源连接方式（平台/Loki 可能在主机或集群内）
	go services.EnsureLogSourcesForRunningDeployments()
	// 按应用日志量的**小时级预聚合**：每小时抓上一个完整小时，
	// 让 6h/24h 这类长范围的"应用排行"变成读库（毫秒级），不扫日志内容。
	go services.StartAppVolumeCapture()
	// 补齐/升级内置告警规则(幂等):启动即执行,存量环境也能拿到「默认启用+绑默认通道」
	seed.SeedAlertRules()
	// 14 天 Pro 试用:首次启动落起点(幂等,含防删库重装指纹)
	services.EnsureTrialStarted()

	// 安全迁移：历史明文 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")
}