LokiStack/backend/repository/db.go

package repository

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

	"lokistack/config"
	"lokistack/models"
)

// 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"`
	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
}

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{},
		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)
	// 启动序列化保存goroutine
	d.wg.Add(1)
	go d.saveWorker()

	db = d
	if migratedLogAccess {
		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 {
	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
	}

	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()
	}
}