TsDB2/DB/core/transaction.go

package core

import (
	"TsDB2/DB/core/page"
	"encoding/binary"
	"fmt"
)

// Transaction 事务管理器
//
// 核心思想：
// - 事务开始时，记录所有被修改行的原始数据和新数据
// - 事务内操作正常修改 BufferPool
// - Commit 时将行级数据写入 WAL（节省空间）
// - Rollback 时用原始数据覆盖 BufferPool 中的脏页（立即生效）
type Transaction struct {
	db            *DB
	originalPages map[uint32][]byte                         // pageID → 原始页数据快照（用于 Rollback）
	rowChanges    map[uint32]struct{ original, new []byte } // pageID → {原始行, 新行}（用于 WAL）
	committed     bool
	rolledBack    bool
}

// BeginTransaction 开启新事务
func (db *DB) BeginTransaction() *Transaction {
	txn := &Transaction{
		db:            db,
		originalPages: make(map[uint32][]byte),
		rowChanges:    make(map[uint32]struct{ original, new []byte }),
		committed:     false,
		rolledBack:    false,
	}

	// 设置当前活跃事务
	db.currentTxn = txn

	return txn
}

// cacheOriginalPage 缓存页的原始数据（首次修改时调用）
func (txn *Transaction) cacheOriginalPage(pageID uint32) error {
	// 如果已经缓存过，不需要重复缓存
	if _, exists := txn.originalPages[pageID]; exists {
		return nil
	}

	// 从磁盘读取原始数据
	buf := make([]byte, PageSize)
	if err := txn.db.file.readPageDisk(pageID, buf); err != nil {
		return fmt.Errorf("failed to read page %d for rollback cache: %v", pageID, err)
	}

	txn.originalPages[pageID] = buf
	return nil
}

// RecordRowChange 记录行级变化（用于 WAL）
// originalRow: 原始行数据
// newRow: 新行数据
func (txn *Transaction) RecordRowChange(pageID uint32, originalRow []byte, newRow []byte) {
	txn.rowChanges[pageID] = struct{ original, new []byte }{
		original: originalRow,
		new:      newRow,
	}
}

// Commit 提交事务
// 混合 WAL 策略:
//  1. 整页条目: 恢复所有被修改页的结构 (B+树节点, IndexMeta, HeaderPage, DataPage)
//  2. 行级条目: 补充 DataPage 的行数据变更 (用于崩溃恢复时精确重做)
//
// 恢复时先应用整页条目 (Phase 1), 再应用行级条目 (Phase 2)
func (txn *Transaction) Commit() error {
	if txn.committed {
		return fmt.Errorf("transaction already committed")
	}
	if txn.rolledBack {
		return fmt.Errorf("transaction already rolled back")
	}

	// 1. 收集 BufferPool 脏页
	dirty := txn.db.file.CollectDirty()
	if len(dirty) == 0 && len(txn.originalPages) == 0 {
		txn.committed = true
		txn.db.currentTxn = nil
		txn.originalPages = nil
		txn.rowChanges = nil
		return nil
	}

	// 2. 创建 WAL 事务
	wal := txn.db.BeginTxn()

	// 3. 写整页条目: 所有被修改的页 (保证页结构完整恢复)
	//    优先用 BufferPool 脏页 (最新数据), 补充 originalPages 中的页
	dirtySet := make(map[uint32]bool)
	for _, d := range dirty {
		wal.RecordPage(d.PageID, d.Data)
		dirtySet[d.PageID] = true
	}
	for pid, origData := range txn.originalPages {
		if !dirtySet[pid] {
			// 原始页被修改但不在脏页中 (理论上不会发生, 安全补充)
			wal.RecordPage(pid, origData)
		}
	}

	// 4. HeaderPage 也纳入
	wal.RecordPage(0, marshalHeader(txn.db.header))

	// 5. 写行级条目: 对 DataPage 提取行差异 (补充恢复信息)
	for _, d := range dirty {
		origData, hasOrig := txn.originalPages[d.PageID]
		if !hasOrig {
			continue // 新页无原始版本, 跳过行差异
		}
		// 只处理 DataPage (type=2)
		if len(origData) < 18 || origData[0] != byte(page.DataIndex) {
			continue
		}
		rowEntries := diffDataPages(origData, d.Data)
		for _, re := range rowEntries {
			wal.RecordRowChange(re.PageID, re.OpType, re.StoreRow)
		}
	}

	// 6. ⭐ 关键：Commit 会触发 fsync，确保落盘 ⭐
	if err := wal.Commit(); err != nil {
		return fmt.Errorf("WAL commit failed: %v", err)
	}

	// 7. 清理事务状态 + 释放原始数据快照
	txn.committed = true
	txn.db.currentTxn = nil
	txn.originalPages = nil
	txn.rowChanges = nil

	return nil
}

// =============================================================================
// diffDataPages 比较两个 DataPage 缓冲区, 提取行级变更
// 返回 INSERT/DELETE/UPDATE 的 RowEntry 列表
// =============================================================================
func diffDataPages(origBuf, curBuf []byte) []RowEntry {
	const hdrSize = 38 // PageHdr(18) + Count(4) + Min(8) + Max(8)
	const delFlag = uint32(0x80000000)

	type rowInfo struct {
		storeRow []byte
		deleted  bool
	}

	// 解析 DataPage 中所有行
	parseRows := func(buf []byte) map[uint64]rowInfo {
		rows := make(map[uint64]rowInfo)
		if len(buf) < hdrSize {
			return rows
		}
		count := binary.LittleEndian.Uint32(buf[18:])
		dataEnd := binary.LittleEndian.Uint32(buf[0:]) // UseSize (offset 14, but we use PageHdr layout)
		// UseSize is at offset 14 in PageHdr: Type(1) + _(3) + PageID(4) + PrePage(4) + NextPage(4) + UseSize(4)
		dataEnd = binary.LittleEndian.Uint32(buf[14:])

		off := uint32(hdrSize)
		for i := uint32(0); i < count; i++ {
			if off+4 > dataEnd || off+4 > uint32(len(buf)) {
				break
			}
			rawLen := binary.LittleEndian.Uint32(buf[off:])
			deleted := rawLen&delFlag != 0
			rowLen := rawLen & ^delFlag
			rowStart := off + 4
			if rowStart+rowLen > dataEnd || rowStart+rowLen > uint32(len(buf)) {
				break
			}
			if rowLen >= 8 {
				pkID := binary.BigEndian.Uint64(buf[rowStart : rowStart+8])
				row := make([]byte, rowLen)
				copy(row, buf[rowStart:rowStart+rowLen])
				rows[pkID] = rowInfo{storeRow: row, deleted: deleted}
			}
			off = rowStart + rowLen
		}
		return rows
	}

	origRows := parseRows(origBuf)
	curRows := parseRows(curBuf)

	var result []RowEntry

	// 检测 DELETE 和 UPDATE (存在于原始页中的行)
	for pkID, orig := range origRows {
		if orig.deleted {
			continue // 原始页中已删除的行不处理
		}
		cur, exists := curRows[pkID]
		if !exists || cur.deleted {
			// 行被删除
			pkBuf := make([]byte, 8)
			binary.BigEndian.PutUint64(pkBuf, pkID)
			result = append(result, RowEntry{OpType: WALOpDelete, StoreRow: pkBuf})
		} else {
			// 检查行数据是否变化
			if !bytesEqual(orig.storeRow, cur.storeRow) {
				result = append(result, RowEntry{OpType: WALOpUpdate, StoreRow: cur.storeRow})
			}
		}
	}

	// 检测 INSERT (存在于当前页但不在原始页中的行)
	for pkID, cur := range curRows {
		if cur.deleted {
			continue
		}
		if _, exists := origRows[pkID]; !exists {
			result = append(result, RowEntry{OpType: WALOpInsert, StoreRow: cur.storeRow})
		}
	}

	return result
}

func bytesEqual(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// Rollback 回滚事务
// 用原始数据快照覆盖 BufferPool 中的脏页（立即生效）
// 回滚后清理所有缓存数据（释放内存）
func (txn *Transaction) Rollback() error {
	if txn.committed {
		return fmt.Errorf("transaction already committed, cannot rollback")
	}
	if txn.rolledBack {
		return fmt.Errorf("transaction already rolled back")
	}

	// 用原始数据覆盖所有被修改的页
	for pageID, originalData := range txn.originalPages {
		// 直接写回 BufferPool（绕过 WAL）
		if err := txn.db.file.WritePage(pageID, originalData); err != nil {
			return fmt.Errorf("failed to rollback page %d: %v", pageID, err)
		}
	}

	// 清理事务状态 + 释放所有缓存数据
	txn.rolledBack = true
	txn.db.currentTxn = nil
	txn.originalPages = nil // ← 回滚成功，原始数据和修改数据都不再需要
	txn.rowChanges = nil    // ← 清理行级变化

	return nil
}

// IsCommitted 检查事务是否已提交
func (txn *Transaction) IsCommitted() bool {
	return txn.committed
}

// IsRolledBack 检查事务是否已回滚
func (txn *Transaction) IsRolledBack() bool {
	return txn.rolledBack
}