LokiStack/backend/repository/db.go

package repository

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

	"lokideploy/config"
	"lokideploy/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"`
	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, 0755)

	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{},
		sequences: map[string]uint{
			"users":       0,
			"clusters":    0,
			"hosts":       0,
			"storages":    0,
			"deployments": 0,
			"collectors":  0,
			"logs":        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
				}
			}
		}
	}

	// 启动序列化保存goroutine
	d.wg.Add(1)
	go d.saveWorker()

	db = d
	return nil
}

// 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(".lokideploy-%d.tmp", time.Now().UnixNano()))

	if err := os.WriteFile(tmpFile, data, 0644); 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
	}
	d.Users = append(d.Users, *user)
	d.RequestSave()
	return nil
}