TsDB2/DB/core/vacuum.go

package core

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"TsDB2/DB/core/page"
	"TsDB2/DB/core/types"
)

// ============================================================================
// VACUUM 整库重写 —— 回收整库碎片, 让文件真正缩小
//
// 工作原理:
//   扫描旧库 schema + 全部行 → 在临时文件 <path>.vacuum.tmp 上重建 →
//   原子 rename 替换旧文件 → 重新打开。
//
// 与自动 compact (40% 失效字节触发, 详见 datapage.go) 的区别:
//   - 自动 compact: 页内墓碑回收, 文件大小几乎不变
//   - VACUUM:       整库重写, 文件真正缩小, 归还所有空闲页给 OS
//
// 失败回滚: 任何阶段失败 → 删除临时文件, 旧库保持不变。
// ============================================================================

// VACUUMMode VACUUM 工作模式
type VACUUMMode int

const (
	// VACUUMModeSync 同步模式: 锁写后单步完成, 阻塞所有读写。
	VACUUMModeSync VACUUMMode = iota

	// VACUUMModeOnline 在线模式: 扫描期允许读 (但不允许写),
	// 切换期短暂阻塞。当前 MVP 实现退化为 Sync (Close-Open),
	// 保留 API 形状供后续接入增量捕获。
	VACUUMModeOnline
)

// VACUUMProgress 进度回调类型。
//   phase:  "scan" | "build" | "index" | "swap" | "done" | "fail"
//   copied: 已复制行数 (scan/build/index 阶段有效)
//   total:   总行数 (估计, 0 表示未知)
type VACUUMProgress func(phase string, copied, total int64)

// VACUUMOptions VACUUM 选项
type VACUUMOptions struct {
	// Progress 进度回调 (nil = 不调用)
	Progress VACUUMProgress
	// Mode 工作模式, 默认 VACUUMModeSync
	Mode VACUUMMode
	// KeepTempOnError 失败时保留临时文件供分析, 默认 false
	KeepTempOnError bool
}

// VACUUMResult VACUUM 结果
type VACUUMResult struct {
	OldFileSize  int64
	NewFileSize  int64
	SavedBytes   int64
	SavedPercent float64
	Tables       int
	RowsCopied   int64
	Duration     time.Duration
	Mode         VACUUMMode
}

// vacuumSrcRow 源行快照 (deserializeRow 的产物, 跨 Copy 安全)
type vacuumSrcRow struct {
	pkID uint64
	vals []types.FieldValue
}

// vacuumSrcIndex 源索引定义
type vacuumSrcIndex struct {
	name   string
	typ    page.IndexType
	fields []byte
	unique bool
}

// vacuumSrcTable 源表信息 (schema + 索引 + 行快照)
type vacuumSrcTable struct {
	name    string
	fields  []types.FieldInfo
	indexes []vacuumSrcIndex
	rows    []vacuumSrcRow
}

// ============================================================================
// 主入口
// ============================================================================

// VACUUM 整库重写, 让数据库文件真正缩小。
//
// 拒绝条件:
//   - 当前有活跃事务 (db.currentTxn != nil)
//   - db 路径不可用
//
// 调用后 db 句柄会被关闭并重新打开, 调用者继续使用即可。
func (db *DB) VACUUM(opts ...VACUUMOptions) (*VACUUMResult, error) {
	opt := VACUUMOptions{Mode: VACUUMModeSync}
	if len(opts) > 0 {
		opt = opts[0]
	}

	if db.currentTxn != nil {
		return nil, fmt.Errorf("%w: VACUUM", ErrTxnActive)
	}
	if db.path == "" || db.path == os.DevNull {
		return nil, fmt.Errorf("%w (%q)", ErrVACUUMNoPath, db.path)
	}

	progress := opt.Progress
	if progress == nil {
		progress = func(string, int64, int64) {}
	}

	start := time.Now()

	// 1. 把 WAL / 脏页全部落盘, 拿到一致快照
	if err := db.Commit(); err != nil {
		return nil, fmt.Errorf("VACUUM: 提交源库失败: %w", err)
	}
	srcStat, _ := os.Stat(db.path)
	var oldSize int64
	if srcStat != nil {
		oldSize = srcStat.Size()
	}

	// 2. 扫描旧库
	progress("scan", 0, 0)
	sources, totalRows, err := db.vacuumScan(progress)
	if err != nil {
		return nil, fmt.Errorf("VACUUM: 扫描源库失败: %w", err)
	}

	// 3. 关闭旧 db (确保所有页都落盘)
	if err := db.Close(); err != nil {
		_ = os.Remove(db.path + ".vacuum.tmp")
		return nil, fmt.Errorf("VACUUM: 关闭源库失败: %w", err)
	}

	// 4. 在临时路径上重建
	tmpPath := db.path + ".vacuum.tmp"
	_ = os.Remove(tmpPath) // 清理残留

	progress("build", 0, totalRows)
	newDB, err := Open(tmpPath)
	if err != nil {
		return db.reopenAfterVacuum(nil, fmt.Errorf("VACUUM: 打开临时库失败: %w", err))
	}

	// 5. 重建表 + 索引 + 插入行
	var rowsCopied int64
	for _, src := range sources {
		// 5a. 建表
		newTbl, err := newDB.CreateTable(src.name, src.fields)
		if err != nil {
			_ = newDB.Close()
			_ = os.Remove(tmpPath)
			return db.reopenAfterVacuum(nil, fmt.Errorf("VACUUM: 重建表 %q 失败: %w", src.name, err))
		}
		// CreateTable 不自动建主键索引, 必须显式建; 否则 InsertRowWithID 找不到行存储索引,
		// resolveIndexes() 返回空, 插入路径无副作用通过, 但 TotalRows++ 谎报数据, Count() 拿到 0。
		if _, err := newTbl.CreateIndex("pk", page.IndexTypePrimary, []byte{0}); err != nil {
			_ = newDB.Close()
			_ = os.Remove(tmpPath)
			return db.reopenAfterVacuum(nil, fmt.Errorf("VACUUM: 建主键索引 %q 失败: %w", src.name, err))
		}

		// 5b. 插入全部行 (走主键索引)
		for ri, row := range src.rows {
			if _, err := newTbl.InsertRowWithID(row.pkID, row.vals); err != nil {
				// 单行失败: 跳过 (deserializeRow 失败/罕见约束冲突)
				continue
			}
			rowsCopied++
			if ri%1000 == 0 {
				progress("build", rowsCopied, totalRows)
			}
		}
		// 周期 commit, 减少 WAL 大小
		if err := newDB.Commit(); err != nil {
			_ = newDB.Close()
			_ = os.Remove(tmpPath)
			return db.reopenAfterVacuum(nil, fmt.Errorf("VACUUM: 提交 %q 失败: %w", src.name, err))
		}
		progress("build", rowsCopied, totalRows)

		// 5c. 重建二级索引 (自动批量回填已存在的行)
		progress("index", rowsCopied, totalRows)
		for _, idx := range src.indexes {
			if isRowStoreType(idx.typ) {
				continue // 行存储索引已在 CreateTable 时建好
			}
			if _, err := newTbl.CreateIndex(idx.name, idx.typ, idx.fields, idx.unique); err != nil {
				_ = newDB.Close()
				_ = os.Remove(tmpPath)
				return db.reopenAfterVacuum(nil, fmt.Errorf("VACUUM: 重建索引 %q 失败: %w", idx.name, err))
			}
		}
		progress("index", rowsCopied, totalRows)
	}

	// 6. 关闭新 db
	if err := newDB.Commit(); err != nil {
		_ = newDB.Close()
		_ = os.Remove(tmpPath)
		return db.reopenAfterVacuum(nil, fmt.Errorf("VACUUM: 最终提交失败: %w", err))
	}
	if err := newDB.Close(); err != nil {
		_ = os.Remove(tmpPath)
		return db.reopenAfterVacuum(nil, fmt.Errorf("VACUUM: 关闭新库失败: %w", err))
	}

	// 7. fsync 父目录 (确保 rename 的目录项落盘)
	if dir := filepath.Dir(db.path); dir != "" {
		if d, derr := os.Open(dir); derr == nil {
			_ = d.Sync()
			_ = d.Close()
		}
	}

	// 8. 原子替换
	progress("swap", 0, 0)
	if err := os.Rename(tmpPath, db.path); err != nil {
		if !opt.KeepTempOnError {
			_ = os.Remove(tmpPath)
		}
		return db.reopenAfterVacuum(nil, fmt.Errorf("VACUUM: 替换文件失败: %w", err))
	}

	// 9. 重新打开
	progress("done", 0, 0)
	newDB2, err := Open(db.path)
	if err != nil {
		return nil, fmt.Errorf("VACUUM: 重新打开失败: %w", err)
	}
	*db = *newDB2

	// 10. 统计结果
	newStat, _ := os.Stat(db.path)
	var newSize int64
	if newStat != nil {
		newSize = newStat.Size()
	}
	result := &VACUUMResult{
		OldFileSize: oldSize,
		NewFileSize: newSize,
		SavedBytes:  oldSize - newSize,
		Tables:      len(sources),
		RowsCopied:  rowsCopied,
		Duration:    time.Since(start),
		Mode:        opt.Mode,
	}
	if oldSize > 0 {
		result.SavedPercent = float64(oldSize-newSize) * 100 / float64(oldSize)
	}
	return result, nil
}

// ============================================================================
// 扫描源库
// ============================================================================

// vacuumScan 扫描源库, 拿到 schema + 索引定义 + 行快照
func (db *DB) vacuumScan(progress VACUUMProgress) ([]vacuumSrcTable, int64, error) {
	tables := db.GetAllTables()
	sources := make([]vacuumSrcTable, 0, len(tables))

	// 先统计总行数 (用于进度回调的 total)
	totalRows := int64(0)
	for _, meta := range tables {
		tbl, ok := db.FindTable(meta.Name)
		if !ok {
			continue
		}
		totalRows += int64(tbl.TotalRows())
	}
	progress("scan", 0, totalRows)

	var scanned int64
	for _, meta := range tables {
		tbl, ok := db.FindTable(meta.Name)
		if !ok {
			continue
		}

		// 字段定义 (剔除已 DROP 的列)
		fields := tbl.GetFields()
		activeFields := make([]types.FieldInfo, 0, len(fields))
		for _, f := range fields {
			if f.IsDropped() {
				continue
			}
			activeFields = append(activeFields, f)
		}

		// 索引定义
		var srcIndexes []vacuumSrcIndex
		for _, im := range tbl.IndexMeta().GetIndexes() {
			srcIndexes = append(srcIndexes, vacuumSrcIndex{
				name:   im.Name(),
				typ:    im.IndexType,
				fields: append([]byte(nil), im.Fields()...),
				unique: im.Unique(),
			})
		}

		// 行快照
		rows := make([]vacuumSrcRow, 0, tbl.TotalRows())
		q := tbl.Query()
		scanErr := q.fullTableScanRaw(func(pkID uint64, rowData []byte) bool {
			vals := tbl.deserializeRow(rowData)
			if vals == nil {
				return true // 解析失败, 跳过
			}
			rows = append(rows, vacuumSrcRow{pkID: pkID, vals: vals})
			scanned++
			if scanned%10000 == 0 {
				progress("scan", scanned, totalRows)
			}
			return true
		})
		if scanErr != nil {
			return nil, 0, fmt.Errorf("扫描表 %q: %w", meta.Name, scanErr)
		}

		sources = append(sources, vacuumSrcTable{
			name:    meta.Name,
			fields:  activeFields,
			indexes: srcIndexes,
			rows:    rows,
		})
	}

	progress("scan", totalRows, totalRows)
	return sources, totalRows, nil
}

// ============================================================================
// 自动调度
// ============================================================================

// SetAutoVACUUM 开启自动 VACUUM, 周期触发。
//
// interval: 两次 VACUUM 的最小间隔 (例如 1*time.Hour)
// options:  每次 VACUUM 使用的选项
//
// 关闭: SetAutoVACUUM(0)
func (db *DB) SetAutoVACUUM(interval time.Duration, opts ...VACUUMOptions) {
	db.autoVacuumMu.Lock()
	defer db.autoVacuumMu.Unlock()

	if interval <= 0 {
		if db.autoVacuumTimer != nil {
			db.autoVacuumTimer.Stop()
			db.autoVacuumTimer = nil
		}
		db.autoVacuumEnabled = false
		return
	}

	opt := VACUUMOptions{Mode: VACUUMModeSync}
	if len(opts) > 0 {
		opt = opts[0]
	}
	db.autoVacuumOpts = opt
	db.autoVacuumEnabled = true
	db.lastVACUUM = time.Now()

	run := func() {
		db.autoVacuumMu.Lock()
		enabled := db.autoVacuumEnabled
		curOpt := db.autoVacuumOpts
		db.autoVacuumMu.Unlock()
		if !enabled {
			return
		}
		// 自动调度: 失败也吞掉 (用户可通过 Progress 回调处理)
		_, _ = db.VACUUM(curOpt)

		db.autoVacuumMu.Lock()
		db.lastVACUUM = time.Now()
		stillEnabled := db.autoVacuumEnabled
		db.autoVacuumMu.Unlock()

		if stillEnabled && db.autoVacuumTimer != nil {
			db.autoVacuumTimer.Reset(interval)
		}
	}

	db.autoVacuumTimer = time.AfterFunc(interval, run)
}

// LastVACUUM 返回上次自动 VACUUM 的时间
func (db *DB) LastVACUUM() time.Time {
	db.autoVacuumMu.Lock()
	defer db.autoVacuumMu.Unlock()
	return db.lastVACUUM
}

// ============================================================================
// 辅助
// ============================================================================

// reopenAfterVacuum VACUUM 失败时尝试重新打开 db.path, 让 db 句柄继续可用
//
// 返回 (*VACUUMResult, error): 即使失败也返回 nil result + 错误,
// 调用方处理逻辑一致。
func (db *DB) reopenAfterVacuum(_ *VACUUMResult, origErr error) (*VACUUMResult, error) {
	newDB, err := Open(db.path)
	if err != nil {
		return nil, fmt.Errorf("%w; 且重新打开源库也失败: %v", origErr, err)
	}
	*db = *newDB
	return nil, origErr
}

// isRowStoreType 索引是否为行存储索引 (主键 / 时序)
//
// 行存储索引在 step 5a 已由 VACUUM 自己手动建好, 重建二级索引时跳过, 避免重复。
func isRowStoreType(typ page.IndexType) bool {
	return typ == page.IndexTypePrimary || typ == page.IndexTypeTime
}

// 编译期检查 (用到的 import 别漏)
var (
	_ = errors.New
	_ sync.Mutex
)