TsDB2/DB/core/transaction.go

package core

import (
	"fmt"
)

// Transaction 事务管理器
//
// 核心思想：
// - 事务内所有修改只落在 BufferPool —— autoCommit 被 currentTxn 拦住, 中途不会刷盘,
//   所以磁盘始终是事务开始前的状态, 不需要任何「原始数据快照」
// - Commit 时把脏页 + Header 整体写入 WAL, 保证原子
// - Rollback 只需丢弃本事务碰过的页缓存(并从磁盘重载 TableInfo), 后续读取自然回到旧值
type Transaction struct {
	db      *DB
	touched map[uint32]struct{} // 本事务修改过的页 ID（Rollback 时丢弃其缓存）
	// 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 {
	if db.errIfReadOnly() != nil {
		panic(ErrReadOnly) // 事务无 error 返回, 抢救模式下显式拒绝
	}
	txn := &Transaction{
		db:         db,
		touched:    make(map[uint32]struct{}),
		committed:  false,
		rolledBack: false,
	}

	// 设置当前活跃事务
	db.currentTxn.Store(txn)

	return txn
}

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

// 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.Store(nil)
		txn.touched = nil
		txn.pendingFree = nil
		txn.pendingFreeChain = nil
		txn.allocated = 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.Store(nil)
	txn.touched = nil
	txn.pendingFree = nil
	txn.pendingFreeChain = nil
	txn.allocated = nil

	return nil
}

// 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.Store(nil)
	txn.touched = nil

	return nil
}

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

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