TsDB2/DB/core/wal.go

package core

import (
	"TsDB2/DB/core/page"
	"encoding/binary"
	"hash/crc32"
	"unsafe"
)

// =============================================================================
// WAL (Write-Ahead Log) 实现
//
// WAL 页布局 (16KB):
//   首页: [PageHdr(20B)] [TxnID(8B)] [TotalEntries(4B)] [State(1B)]
//         [NextWalPage(4B)] [EntryCount(2B)] [WALEntry...]
//   续页: [PageHdr(20B)] [NextWalPage(4B)] [EntryCount(2B)] [WALEntry...]
//
// WALEntry: [TargetPageID(4B)] [DataLen(4B)] [Checksum(4B)] [Data(DataLen)]
//
// WAL 状态机 (State 字段):
//   0 = Mem (内存中，未写盘)
//   1 = Written (已 write()，未 fsync())
//   2 = Persisted (已 fsync()，但未标记提交)
//   3 = Committed (已提交，可安全恢复)
//
// 每个 WAL 页存 1 条完整 16KB 页镜像 (因为 entry 12B + 16KB > 16KB - 39B)
// =============================================================================

const (
	walPHDR      = uint32(unsafe.Sizeof(page.PageHdr{})) // 20
	walFirstHdr  = walPHDR + 8 + 4 + 1 + 4 + 2           // 39: PageHdr + TxnID + TotalEntries + State + NextWalPage + EntryCount
	walContHdr   = walPHDR + 4 + 2                       // 26: PageHdr + NextWalPage + EntryCount
	walEntryMeta = 4 + 4 + 4                             // 12: TargetPageID + DataLen + Checksum
)

// WAL 状态常量 (对应 State 字段)
const (
	WALStateMem       = 0 // 内存中，未写盘
	WALStateWritten   = 1 // 已 write()，未 fsync()
	WALStatePersisted = 2 // 已 fsync()，但未标记提交
	WALStateCommitted = 3 // 已提交，可安全恢复
)

// WALWriter 事务 WAL 写入器
type WALWriter struct {
	txnID    uint64
	entries  []walEntry
	file     *File
	header   *HeaderPage
	readPage func(uint32, []byte) error // 用于分配页时读取空闲链表
	state    byte                       // 当前 WAL 状态 (Mem/Written/Persisted/Committed)
}

type walEntry struct {
	pageID         uint32
	data           []byte // 完整页镜像（原有逻辑，用于 Commit）
	originalData   []byte // 原始行数据（事务回滚用）← 新增
	newData        []byte // 新行数据（事务提交用）← 新增
	isRowLevel     bool   // 是否为行级记录（true=存数据，false=存整页）← 新增
}

// BeginTxn 开始新事务
func (db *DB) BeginTxn() *WALWriter {
	txnID := db.header.NextTxnID
	db.header.NextTxnID++
	return &WALWriter{
		txnID:    txnID,
		entries:  make([]walEntry, 0, 16),
		file:     db.file,
		header:   db.header,
		readPage: db.file.readPageDisk, // 必须直接读盘! 走 BufferPool 会被旧缓存干扰空闲链表
		state:    WALStateMem,          // 初始状态: Mem (内存中)
	}
}

// RecordPage 记录一个被修改的页 (同一 pageID 只保留最新镜像)
func (w *WALWriter) RecordPage(pageID uint32, data []byte) {
	d := make([]byte, len(data))
	copy(d, data)
	// 同一 pageID 只保留最新镜像
	for i := range w.entries {
		if w.entries[i].pageID == pageID {
			w.entries[i].data = d
			return
		}
	}
	w.entries = append(w.entries, walEntry{
		pageID:     pageID,
		data:       d,
		isRowLevel: false, // 原有逻辑：存整页
	})
}

// RecordRowLevel 记录行级数据（用于事务支持）
// 存储原始数据和新数据，而非整页，节省空间
func (w *WALWriter) RecordRowLevel(pageID uint32, originalData []byte, newData []byte) {
	orig := make([]byte, len(originalData))
	copy(orig, originalData)
	
	newD := make([]byte, len(newData))
	copy(newD, newData)
	
	// 同一 pageID 只保留最新
	for i := range w.entries {
		if w.entries[i].pageID == pageID && w.entries[i].isRowLevel {
			w.entries[i].originalData = orig
			w.entries[i].newData = newD
			return
		}
	}
	
	w.entries = append(w.entries, walEntry{
		pageID:       pageID,
		originalData: orig,
		newData:      newD,
		isRowLevel:   true, // 行级记录
	})
}

// Rollback 回滚: 丢弃所有条目
func (w *WALWriter) Rollback() {
	w.entries = nil
	w.state = WALStateMem // 重置状态
}

// Commit 原子提交流程:
//  1. 将 entries 写入 WAL 页链 (State: Mem → Written)
//  2. fsync WAL 页 (State: Written → Persisted)
//  3. 回写实际页 (从 WAL 拷贝到真实位置)
//  4. 标记 WAL 首页 Committed = 1 (State: Persisted → Committed)
//  5. 回收所有 WAL 页
func (w *WALWriter) Commit() error {
	if len(w.entries) == 0 {
		return nil
	}

	// === 步骤 1: 写入 WAL 页链 (State: Mem → Written) ===
	walPageIDs, err := w.writeWALPages()
	if err != nil {
		return err
	}

	// === 步骤 2: fsync WAL 页 (State: Written → Persisted) ===
	if err := w.file.sync(); err != nil {
		return err
	}
	// 状态转换: Written → Persisted (已落盘)
	w.state = WALStatePersisted

	// === 步骤 3: 回写实际页 (直接写盘, 绕过 BufferPool) ===
	for _, e := range w.entries {
		if e.pageID == 0 {
			continue
		}
		if err := w.file.WritePageDirect(e.pageID, e.data); err != nil {
			return err
		}
	}
	w.file.ClearCache()

	// === 步骤 4: 标记 WAL 首页为已提交 (State: Persisted → Committed) ===
	if err := w.markCommitted(walPageIDs[0]); err != nil {
		return err
	}
	if err := w.file.sync(); err != nil {
		return err
	}

	// === 步骤 5: 回收所有 WAL 页 (直接写盘 + 同步 BufferPool 缓存) ===
	for _, pid := range walPageIDs {
		w.header.recyclePage(pid, w.file.WritePageDirect)
		// 同步 BufferPool 缓存: recyclePage 通过 WritePageDirect 写盘, 绕过了 BufferPool
		// 必须从磁盘读回并更新缓存, 否则 newPageWithReader 从缓存读到旧数据导致页面无法回收
		syncBuf := make([]byte, PageSize)
		_ = w.file.readPageDisk(pid, syncBuf)
		w.file.SyncCache(pid, syncBuf)
	}
	w.header.FirstWALPage = 0
	w.header.LastWALPage = 0
	if err := w.file.WriteHeader(w.header); err != nil {
		return err
	}

	w.entries = nil
	return nil
}

// writeWALPages 将所有 entry 写入 WAL 页链, 返回分配的页 ID 列表
// 状态转换: Mem → Written
func (w *WALWriter) writeWALPages() ([]uint32, error) {
	var pageIDs []uint32
	entryIdx := 0

	for entryIdx < len(w.entries) {
		// 分配新 WAL 页 (必须用 newPageWithReader 推进空闲链表)
		pid := w.header.newPageWithReader(w.readPage)
		pageIDs = append(pageIDs, pid)

		buf := make([]byte, PageSize)
		hdr := (*page.PageHdr)(unsafe.Pointer(&buf[0]))
		hdr.Type = page.WALIndex
		hdr.PageID = pid
		if len(pageIDs) > 1 {
			hdr.PrePage = pageIDs[len(pageIDs)-2]
		}

		off := uint32(0)

		if entryIdx == 0 {
			// === 首页: 写入 TxnID + TotalEntries + State + NextWalPage + EntryCount ===
			binary.LittleEndian.PutUint64(buf[walPHDR:], w.txnID)
			binary.LittleEndian.PutUint32(buf[walPHDR+8:], uint32(len(w.entries)))
			buf[walPHDR+12] = WALStateWritten                  // State = Written (已写盘，未 fsync)
			binary.LittleEndian.PutUint32(buf[walPHDR+13:], 0) // NextWalPage (稍后填)
			binary.LittleEndian.PutUint16(buf[walPHDR+17:], 0) // EntryCount (稍后填)

			off = walFirstHdr
		} else {
			// === 续页: 写入 NextWalPage + EntryCount ===
			binary.LittleEndian.PutUint32(buf[walPHDR:], 0)   // NextWalPage (稍后填)
			binary.LittleEndian.PutUint16(buf[walPHDR+4:], 0) // EntryCount (稍后填)
			off = walContHdr
		}

		// 写入当前 entry
		e := w.entries[entryIdx]
		
		if e.isRowLevel {
			// === 行级记录: [PageID][OrigLen][NewLen][Checksum][OriginalData][NewData] ===
			origLen := uint32(len(e.originalData))
			newLen := uint32(len(e.newData))
			totalLen := origLen + newLen
			
			// 计算 checksum (覆盖 originalData + newData)
			combined := make([]byte, 0, totalLen)
			combined = append(combined, e.originalData...)
			combined = append(combined, e.newData...)
			checksum := crc32.ChecksumIEEE(combined)
			
			// 写入元数据
			binary.LittleEndian.PutUint32(buf[off:], e.pageID)
			binary.LittleEndian.PutUint32(buf[off+4:], origLen)    // 用 origLen 替代 storeLen
			binary.LittleEndian.PutUint32(buf[off+8:], newLen)     // 新增: newLen
			binary.LittleEndian.PutUint32(buf[off+12:], checksum)
			off += 16 // PageID + OrigLen + NewLen + Checksum
			
			// 写入数据
			copy(buf[off:], e.originalData)
			off += origLen
			copy(buf[off:], e.newData)
			off += newLen
		} else {
			// === 整页记录: [PageID][DataLen][Checksum][Data] (原有逻辑) ===
			storeLen := uint32(len(e.data)) // 默认存全部
			if len(e.data) >= 20 {
				useSize := binary.LittleEndian.Uint32(e.data[16:20])
				if useSize > 0 && useSize < uint32(len(e.data)) {
					storeLen = useSize
				}
				// HeaderPage (pageID=0): 至少存 PageSize4K
				if e.pageID == 0 && storeLen < PageSize4K {
					storeLen = PageSize4K
				}
			}
			checksum := crc32.ChecksumIEEE(e.data[:storeLen])
			binary.LittleEndian.PutUint32(buf[off:], e.pageID)
			binary.LittleEndian.PutUint32(buf[off+4:], storeLen)
			binary.LittleEndian.PutUint32(buf[off+8:], checksum)
			copy(buf[off+12:], e.data[:storeLen])
			// storeLen 可能 < len(e.data), 剩余字节保持零填充
			off += walEntryMeta + storeLen
		}

		hdr.UseSize = off
		entryIdx++

		// 立即设置当前页的 EntryCount = 1 (不依赖后续链接代码)
		if entryIdx == 1 {
			// 首页: EntryCount 在 walPHDR+17
			binary.LittleEndian.PutUint16(buf[walPHDR+17:], 1)
		} else {
			// 续页: EntryCount 在 walPHDR+4
			binary.LittleEndian.PutUint16(buf[walPHDR+4:], 1)
		}

		// 链接上一页的 NextWalPage
		if len(pageIDs) > 1 {
			prevPID := pageIDs[len(pageIDs)-2]
			prevBuf := make([]byte, PageSize)
			if err := w.file.ReadPage(prevPID, prevBuf); err == nil {
				if isFirstWALPage(prevBuf) {
					binary.LittleEndian.PutUint32(prevBuf[walPHDR+13:], pid)
					binary.LittleEndian.PutUint16(prevBuf[walPHDR+17:], 1)
				} else {
					binary.LittleEndian.PutUint32(prevBuf[walPHDR:], pid)
					binary.LittleEndian.PutUint16(prevBuf[walPHDR+4:], 1)
				}
				_ = w.file.WritePageDirect(prevPID, prevBuf)
			}
		}

		// 写入当前页
		if err := w.file.WritePageDirect(pid, buf); err != nil {
			return nil, err
		}
	}

	// 更新 HeaderPage 的 WAL 链指针
	if len(pageIDs) > 0 {
		if w.header.FirstWALPage == 0 {
			w.header.FirstWALPage = pageIDs[0]
		}
		w.header.LastWALPage = pageIDs[len(pageIDs)-1]
		_ = w.file.WriteHeader(w.header)
	}

	// 状态转换: Mem → Written (已写盘，未 fsync)
	w.state = WALStateWritten

	return pageIDs, nil
}

// markCommitted 标记 WAL 首页为已提交 (State: Persisted → Committed)
func (w *WALWriter) markCommitted(firstPageID uint32) error {
	buf := make([]byte, PageSize)
	if err := w.file.ReadPage(firstPageID, buf); err != nil {
		return err
	}
	buf[walPHDR+12] = WALStateCommitted // State = Committed
	if err := w.file.WritePageDirect(firstPageID, buf); err != nil {
		return err
	}
	// 状态转换: Persisted → Committed
	w.state = WALStateCommitted
	return nil
}

// isFirstWALPage 判断 WAL 页是否为事务首页
func isFirstWALPage(buf []byte) bool {
	// 首页在 walPHDR 处有 TxnID (非零), 续页没有
	txnID := binary.LittleEndian.Uint64(buf[walPHDR:])
	return txnID != 0
}

// sync 刷盘 (fsync)
func (ff *File) sync() error {
	return ff.f.Sync()
}