LokiStack/backend/repository/db.go

package repository

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"lokistack/config"
	"lokistack/models"
	"lokistack/pkg/metrics"
)

// saveRequest 表示一个保存请求，通过channel序列化
type saveRequest struct {
	done chan error // 用于同步等待保存完成
}

// JsonDB 使用基于channel的序列化保存机制，防止并发数据丢失。
// 该包只负责持久化（落盘 + 并发安全），领域实体定义在 models 包。
type JsonDB struct {
	mu              sync.RWMutex
	dataFile        string
	Users           []models.User           `json:"users"`
	Clusters        []models.Cluster        `json:"clusters"`
	Hosts           []models.Host           `json:"hosts"`
	Storages        []models.StorageConfig  `json:"storages"`
	Deployments     []models.LokiDeployment `json:"deployments"`
	Collectors      []models.Collector      `json:"collectors"`
	Logs            []models.DeploymentLog  `json:"logs"`
	Settings        []models.SystemSetting  `json:"settings"`
	LogSources      []models.LogSource      `json:"logSources"`
	AlertRules      []models.AlertRule      `json:"alertRules"`
	AlertDeliveries []models.AlertDelivery  `json:"alertDeliveries"`
	NotifyChannels  []models.NotifyChannel  `json:"notifyChannels"`
	LogPolicyGroups []models.LogPolicyGroup `json:"logPolicyGroups"`
	// AuditLogs 操作审计日志（保留期 90 天 / 上限 2 万条，见 services.auditRetention）
	AuditLogs []models.AuditLog `json:"auditLogs"`

	// SchemaVersion 数据结构版本号。
	// 之前加字段全靠 JSON 反序列化的零值兜底，字段一多必然出事
	// （例如新增 ScrapeJobs 时老数据要手写兼容）。有了版本号，
	// 以后的结构变更可以在 InitDB 里按版本号顺序迁移，可测、可回滚。
	SchemaVersion int `json:"schemaVersion"`
	sequences       map[string]uint

	// 防抖保存机制
	saveTimer *time.Timer
	saveMu    sync.Mutex
	saveChan  chan saveRequest
	stopChan  chan struct{}
	wg        sync.WaitGroup
}

var db *JsonDB

func GetDB() *JsonDB {
	return db
}

// CurrentSchemaVersion 当前数据结构版本。
// 每次做**不兼容**的结构变更时 +1，并在 migrateSchema 里补一段迁移逻辑。
//
//	v0/v1（无字段）→ 1：新增 auditLogs / schemaVersion（老库走零值兜底即可，无需处理）
//	v1 → 2：预留示例（如把 ScrapeJobs 从 logPath 拆分时，需按旧结构生成默认 job）
const CurrentSchemaVersion = 1

// migrateSchema 按版本号顺序迁移老库。返回是否发生了迁移（需要回写）。
func (d *JsonDB) migrateSchema() bool {
	if d.SchemaVersion >= CurrentSchemaVersion {
		return false
	}
	from := d.SchemaVersion

	// 逐版本递进，每段只处理"上一版 → 下一版"的差异，
	// 不要写成 if version < X 的叠加判断，否则分支会指数增长。
	for d.SchemaVersion < CurrentSchemaVersion {
		switch d.SchemaVersion {
		case 0:
			// v0 → v1：补审计表；老库没有 auditLogs 字段，反序列化后为 nil，
			// 这里显式初始化，避免后续 append 到 nil 之外还留着 nil 判断。
			if d.AuditLogs == nil {
				d.AuditLogs = []models.AuditLog{}
			}
			d.SchemaVersion = 1
		default:
			// 防御：出现未知版本（比当前代码还新）时不乱动，避免降级把数据写坏
			log.Printf("[schema] 数据库版本 %d 高于当前代码支持版本 %d，跳过迁移", d.SchemaVersion, CurrentSchemaVersion)
			return false
		}
	}
	log.Printf("[schema] 数据结构已从 v%d 迁移到 v%d", from, d.SchemaVersion)
	return true
}

func InitDB() error {
	cfg := config.GetConfig()
	_ = os.MkdirAll(cfg.DataDir, 0700)

	d := &JsonDB{
		dataFile:        cfg.DatabasePath,
		Users:           []models.User{},
		Clusters:        []models.Cluster{},
		Hosts:           []models.Host{},
		Storages:        []models.StorageConfig{},
		Deployments:     []models.LokiDeployment{},
		Collectors:      []models.Collector{},
		Logs:            []models.DeploymentLog{},
		Settings:        []models.SystemSetting{},
		LogSources:      []models.LogSource{},
		AlertRules:      []models.AlertRule{},
		AlertDeliveries: []models.AlertDelivery{},
		NotifyChannels:  []models.NotifyChannel{},
		LogPolicyGroups: []models.LogPolicyGroup{},
		AuditLogs:       []models.AuditLog{},
		SchemaVersion:   CurrentSchemaVersion,
		sequences: map[string]uint{
			"users":           0,
			"clusters":        0,
			"hosts":           0,
			"storages":        0,
			"deployments":     0,
			"collectors":      0,
			"logs":            0,
			"logPolicyGroups": 0,
		},
		saveChan: make(chan saveRequest, 10),
		stopChan: make(chan struct{}),
	}

	if data, err := os.ReadFile(d.dataFile); err == nil {
		if err := json.Unmarshal(data, d); err == nil {
			d.sequences = map[string]uint{}
			for _, u := range d.Users {
				if u.ID > d.sequences["users"] {
					d.sequences["users"] = u.ID
				}
			}
			for _, c := range d.Clusters {
				if c.ID > d.sequences["clusters"] {
					d.sequences["clusters"] = c.ID
				}
			}
			for _, h := range d.Hosts {
				if h.ID > d.sequences["hosts"] {
					d.sequences["hosts"] = h.ID
				}
			}
			for _, s := range d.Storages {
				if s.ID > d.sequences["storages"] {
					d.sequences["storages"] = s.ID
				}
			}
			for _, de := range d.Deployments {
				if de.ID > d.sequences["deployments"] {
					d.sequences["deployments"] = de.ID
				}
			}
			for _, col := range d.Collectors {
				if col.ID > d.sequences["collectors"] {
					d.sequences["collectors"] = col.ID
				}
			}
			for _, l := range d.Logs {
				if l.ID > d.sequences["logs"] {
					d.sequences["logs"] = l.ID
				}
			}
			for _, group := range d.LogPolicyGroups {
				if group.ID > d.sequences["logPolicyGroups"] {
					d.sequences["logPolicyGroups"] = group.ID
				}
			}
		}
	}

	migratedLogAccess := migrateLogAccessToGroups(d)
	// 结构版本迁移（审计表、采集项等新字段的兼容处理）
	migratedSchema := d.migrateSchema()
	// 启动序列化保存goroutine
	d.wg.Add(1)
	go d.saveWorker()

	db = d
	if migratedLogAccess || migratedSchema {
		d.RequestSave()
	}
	return nil
}

// Legacy per-user grants are grouped by identical source and label scope.
func migrateLogAccessToGroups(d *JsonDB) bool {
	changed := false
	for i := range d.LogSources {
		source := &d.LogSources[i]
		if len(source.Access) == 0 {
			continue
		}
		byLabels := map[string]*models.LogPolicyGroup{}
		for _, rule := range source.Access {
			validUser := false
			for _, user := range d.Users {
				if user.ID == rule.UserID && user.Role != "admin" {
					validUser = true
					break
				}
			}
			if !validUser {
				continue
			}
			alreadyGrouped := false
			for _, group := range d.LogPolicyGroups {
				if group.SourceID != source.ID {
					continue
				}
				for _, id := range group.UserIDs {
					if id == rule.UserID {
						alreadyGrouped = true
						break
					}
				}
			}
			if alreadyGrouped {
				continue
			}
			keyBytes, _ := json.Marshal(rule.Labels)
			key := string(keyBytes)
			group := byLabels[key]
			if group == nil {
				d.sequences["logPolicyGroups"]++
				id := d.sequences["logPolicyGroups"]
				group = &models.LogPolicyGroup{ID: id, Name: fmt.Sprintf("旧授权 · %s · %d", source.Name, id), SourceID: source.ID, Labels: rule.Labels, CreatedAt: time.Now(), UpdatedAt: time.Now()}
				byLabels[key] = group
			}
			group.UserIDs = append(group.UserIDs, rule.UserID)
		}
		for _, group := range byLabels {
			d.LogPolicyGroups = append(d.LogPolicyGroups, *group)
		}
		source.Access = nil
		changed = true
	}
	return changed
}

// Shutdown 优雅关闭数据库，确保所有数据已落盘。
// 顺序很关键：先停防抖定时器（不再产生新请求）→ 关闭 worker 并等其退出
// （worker 可能正在处理已入队的请求）→ 此时已无并发，做最后一次 doSave 把
// 最新内存状态原子写入。注意绝不能在 worker 仍在运行时直接 Save()，否则
// 两个 doSave 并发改写同一文件会发生"后写覆盖先写"的丢更新。
func Shutdown() error {
	if db == nil {
		return nil
	}

	// 1. 取消防抖定时器，阻止新的保存请求进入 channel
	db.saveMu.Lock()
	if db.saveTimer != nil {
		db.saveTimer.Stop()
		db.saveTimer = nil
	}
	db.saveMu.Unlock()

	// 2. 通知 worker 停止并等待其排空 saveChan、完成进行中的 doSave 后退出
	close(db.stopChan)
	db.wg.Wait()

	// 3. worker 已退出，无并发，落盘最终状态
	if err := db.doSave(); err != nil {
		return fmt.Errorf("final save failed: %w", err)
	}
	return nil
}

// saveWorker 通过channel序列化所有保存操作，防止并发写入冲突
func (d *JsonDB) saveWorker() {
	defer d.wg.Done()
	for {
		select {
		case req, ok := <-d.saveChan:
			if !ok {
				return
			}
			err := d.doSave()
			if req.done != nil {
				req.done <- err
			}
		case <-d.stopChan:
			return
		}
	}
}

// doSave 执行实际的原子文件写入
func (d *JsonDB) doSave() error {
	start := time.Now()
	err := d.doSaveInner()
	// JSON DB 是全量写盘，落盘耗时是平台最主要的延迟来源之一，必须可观测
	metrics.ObserveDBSave(time.Since(start).Seconds(), err == nil)
	return err
}

func (d *JsonDB) doSaveInner() error {
	d.mu.RLock()
	data, err := json.MarshalIndent(d, "", "  ")
	d.mu.RUnlock()
	if err != nil {
		return err
	}

	// 原子写入：先写入临时文件，再重命名
	dir := filepath.Dir(d.dataFile)
	tmpFile := filepath.Join(dir, fmt.Sprintf(".lokistack-%d.tmp", time.Now().UnixNano()))

	if err := os.WriteFile(tmpFile, data, 0600); err != nil {
		return err
	}

	// 确保数据落盘
	f, err := os.Open(tmpFile)
	if err != nil {
		os.Remove(tmpFile)
		return err
	}
	_ = f.Sync()
	f.Close()

	if err := os.Rename(tmpFile, d.dataFile); err != nil {
		os.Remove(tmpFile)
		return err
	}
	metrics.DBSizeBytes.Set(float64(len(data)))

	return nil
}

// Save 同步保存数据库到文件（原子写入）
func (d *JsonDB) Save() error {
	return d.doSave()
}

// RequestSave 请求异步保存（防抖机制+序列化channel）
// 多次快速调用会在500ms内合并为一次保存
func (d *JsonDB) RequestSave() {
	d.saveMu.Lock()
	defer d.saveMu.Unlock()

	// 如果已有待触发的定时器，重置它（防抖）
	if d.saveTimer != nil {
		d.saveTimer.Stop()
	}

	d.saveTimer = time.AfterFunc(500*time.Millisecond, func() {
		d.saveMu.Lock()
		d.saveTimer = nil
		d.saveMu.Unlock()

		// 通过channel发送到保存worker，确保序列化
		select {
		case d.saveChan <- saveRequest{}:
		case <-d.stopChan:
		}
	})
}

// RequestSaveSync 请求同步保存，等待保存完成（用于shutdown场景）
func (d *JsonDB) RequestSaveSync() error {
	req := saveRequest{done: make(chan error, 1)}
	select {
	case d.saveChan <- req:
	case <-d.stopChan:
		return d.Save() // worker已停止，直接同步保存
	}
	return <-req.done
}

func (d *JsonDB) NextID(table string) uint {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.sequences[table]++
	return d.sequences[table]
}

func (d *JsonDB) Lock() {
	d.mu.Lock()
}

func (d *JsonDB) Unlock() {
	d.mu.Unlock()
}

func (d *JsonDB) RLock() {
	d.mu.RLock()
}

func (d *JsonDB) RUnlock() {
	d.mu.RUnlock()
}

// CreateUser 创建用户并触发防抖保存
func (d *JsonDB) CreateUser(user *models.User) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	if user.ID == 0 {
		user.ID = d.sequences["users"] + 1
		d.sequences["users"] = user.ID
	}
	if user.CreatedAt.IsZero() {
		now := time.Now()
		user.CreatedAt = now
		user.UpdatedAt = now
	}
	d.Users = append(d.Users, *user)
	d.RequestSave()
	return nil
}

// RecoverOrphanDeployments 启动时把残留的 deploying 状态部署标记为失败:
// 部署 goroutine 随进程消亡,不恢复会永久卡在「部署中」(UI 无停止入口)。
func (d *JsonDB) RecoverOrphanDeployments() {
	d.Lock()
	defer d.Unlock()
	recovered := 0
	for i := range d.Deployments {
		if d.Deployments[i].Status == "deploying" {
			d.Deployments[i].Status = "failed"
			d.Deployments[i].StatusMessage = "服务重启导致部署中断，请重新部署"
			d.Deployments[i].UpdatedAt = time.Now()
			recovered++
		}
	}
	if recovered > 0 {
		d.RequestSave()
	}
}