TsDB2/DB/core/transaction.go

package core

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

// Transaction 事务管理器
//
// 核心思想：
// - 事务内所有修改只落在 BufferPool —— autoCommit 被 currentTxn 拦住, 中途不会刷盘,
//   所以磁盘始终是事务开始前的状态, 不需要任何「原始数据快照」
// - Commit 时把脏页 + Header 整体写入 WAL, 保证原子
// - Rollback 只需丢弃本事务碰过的页缓存(并从磁盘重载 TableInfo), 后续读取自然回到旧值
type Transaction struct {
	db         *DB
	touched    map[uint32]struct{}                       // 本事务修改过的页 ID（Rollback 时丢弃其缓存）
	rowChanges map[uint32]struct{ original, new []byte } // pageID → {原始行, 新行}（用于 WAL）
	// pendingFree 本事务登记待回收的页 (堆页摘链 / 批量删除后作废的旧 DataPage)。
	// 不在事务中途直写盘: 提交成功后才真正还链, 回滚时整体丢弃
	// (回滚后链指针回到事务前, 这些页仍被引用, 绝不能还链)。
	pendingFree []uint32
	// pendingFreeChain 本事务登记待回收的整链 (如哈希索引整体重建后的旧桶链)。
	// 与 pendingFree 同一时序约定, 但一条链只登记一次, 提交后一次 FreeChain 还链。
	pendingFreeChain []pendingChain
	// allocated 本事务新分配的页 ID (从空闲链摘出, 或空闲链耗尽时由 NextFreePage 递增)。
	// 分配在 newPageWithReader 里就地推进了 header 的游离指针并直写盘, 事务中途无法回退;
	// 因此 Rollback 必须把这些页还给空闲链 —— 否则它们既不被任何表/索引引用、也不在
	// 空闲链上, 就成了永久泄漏的孤儿页 (文件只涨不缩)。
	// Commit 成功时这些页已被数据引用, 直接丢弃本表即可。
	allocated  []uint32
	committed  bool
	rolledBack bool
}

// 事务内被直接 FreePage(立即还链) 的页, 从 allocated 中撤销登记, 避免 Rollback 二次归还。
func (txn *Transaction) undoAllocated(pageID uint32) {
	for i, id := range txn.allocated {
		if id == pageID {
			txn.allocated = append(txn.allocated[:i], txn.allocated[i+1:]...)
			return
		}
	}
}

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

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

	return txn
}

// markTouched 登记本事务修改过的页 (由 BufferPool 的 onPageWrite 回调触发)。
// 只记页 ID(4B/页), 不存页数据 —— 回滚靠「丢弃缓存 + 从磁盘重读」实现。
func (txn *Transaction) markTouched(pageID uint32) {
	txn.touched[pageID] = struct{}{}
}

// 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 提交事务
// 只写「整页镜像」条目: 恢复所有被修改页的结构与数据
// (B+树节点, IndexMeta, HeaderPage, DataPage 一次到位)。
//
// 不再额外写行级条目: 整页镜像已完整包含行数据, 再叠加行级记录会在恢复时被
// 重复应用 (INSERT 变成两行), 见 wal_recovery 的 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")
	}

	// 与 DB.Commit 共用同一把提交锁: 事务提交和自动提交(定时器 / 脏页阈值)也可能撞
	// 在一起, 两笔提交同时在飞会把页的归属搞乱。详见 DB.Commit 的说明。
	txn.db.commitMu.Lock()
	defer txn.db.commitMu.Unlock()

	// 1. 收集 BufferPool 脏页
	dirty := txn.db.file.CollectDirty()
	if len(dirty) == 0 {
		// 无脏页 = 修改已持久, 本事务登记的待回收页 + 全局积压都可以安全还链
		txn.db.freeNow(txn.pendingFree)
		txn.db.freeChainNow(txn.pendingFreeChain)
		txn.db.drainPendingFree()
		txn.committed = true
		txn.db.currentTxn = nil
		txn.touched = nil
		txn.pendingFree = nil
		txn.pendingFreeChain = nil
		txn.allocated = nil
		txn.rowChanges = nil
		return nil
	}

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

	// 3. 写整页条目: 所有被修改的页 (保证页结构完整恢复)
	//    数据源就是 BufferPool 脏页本身(最新数据), 事务期间没有中途刷盘,
	//    因此不存在"写回事务前旧快照覆盖新数据"的风险。
	for _, d := range dirty {
		wal.RecordPage(d.PageID, d.Data)
	}

	// 4. HeaderPage 也纳入
	//    marshalHeader 读整张头页, 含提交路径正在改的 WAL 链指针, 必须在 headerMu 内拍快照。
	txn.db.file.withHeader(func() {
		wal.RecordPage(0, marshalHeader(txn.db.header))
	})

	// 5. ⭐ 关键：Commit 会触发 fsync，确保落盘 ⭐
	if err := wal.Commit(); err != nil {
		// 提交失败: 这批页继续留在缓存里当脏页(CollectDirty 不再移出缓存),
		// RestoreDirty 只做兜底, 见 DB.Commit 里的说明
		txn.db.file.RestoreDirty(dirty)
		return fmt.Errorf("WAL commit failed: %v", err)
	}

	// 6. 提交成功 → 把这一批页收成干净页 (按快照版本号, 提交窗口里被改过的页继续留脏)
	txn.db.file.MarkClean(dirty)

	// 7. 提交成功 → 堆链/索引指针已持久化, 本事务登记摘链的页现在可以真正还链。
	//    必须在 wal.Commit() 之后: 它内部已完成 fsync + 回写真实页。
	txn.db.freeNow(txn.pendingFree)
	txn.db.freeChainNow(txn.pendingFreeChain)
	txn.db.drainPendingFree()

	// 8. 清理事务状态
	txn.committed = true
	txn.db.currentTxn = nil
	txn.touched = nil
	txn.pendingFree = nil
	txn.pendingFreeChain = nil
	txn.allocated = nil
	txn.rowChanges = nil

	return nil
}

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

	type rowInfo struct {
		storeRow []byte
		deleted  bool
	}

	// 解析 DataPage 中所有行
	// 布局: [PageHdr(20B)][Count(4B)][Min(8B)][Max(8B)] = 40B, 行区从 40 开始
	parseRows := func(buf []byte) map[uint64]rowInfo {
		rows := make(map[uint64]rowInfo)
		if len(buf) < int(page.DataPageHdrSize) {
			return rows
		}
		dh := page.DataHdrOf(buf)
		count := dh.Count
		dataEnd := dh.UseSize
		if dataEnd < page.DataPageHdrSize || dataEnd > uint32(len(buf)) {
			return rows
		}

		off := page.DataPageHdrSize
		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 回滚事务: 丢弃本事务修改过的页缓存, 从磁盘重新加载。
//
// 事务期间 autoCommit 被 currentTxn 拦住, 没有任何脏页被刷到磁盘, 磁盘上仍是事务
// 开始前的状态 —— 所以"丢掉缓存"就等价于"撤销内存修改", 不需要任何页级快照。
// 只丢 txn.touched 里的页: 事务开始前遗留的脏页不属于本事务, 必须原样保留。
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")
	}

	// 1. 丢弃本事务碰过的页缓存 (不落盘), 后续读取自然回落到磁盘旧值
	txn.db.file.DiscardPages(txn.touched)

	// 2. Table.info 是独立于页缓存的内存副本, 必须从磁盘重新读回,
	//    否则 TotalRows / FirstData 等仍停留在事务中的值
	txn.db.tables.ReloadInfos(txn.touched)

	// 3. 本事务登记待回收的页/链必须整体丢弃: 回滚后堆链指针回到事务前,
	//    这些页重新被引用, 绝不能还链 (全局 pendingFree 不属于本事务, 保持不动)
	txn.pendingFree = nil
	txn.pendingFreeChain = nil

	// 4. 归还本事务新分配的页 (见 allocated 字段注释): 分配时 header 的游离指针
	//    已被就地推进并直写盘, 回滚不补回这些页就会产生孤儿页。
	//    放在 DiscardPages 之后: 池里那些未落盘的脏内容已丢弃, 这里用干净的空闲页头覆盖。
	txn.db.freeNow(txn.allocated)
	txn.allocated = nil

	// 清理事务状态
	txn.rolledBack = true
	txn.db.currentTxn = nil
	txn.touched = nil
	txn.rowChanges = nil

	return nil
}

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

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