TsDB2/DB/core/blob_commit.go

package core

import (
	"bytes"
	"errors"
	"io"

	"github.com/golang/snappy"

	"TsDB2/DB/core/types"
)

// BlobCommit 把 io.Reader 流式写入 blob 链, 一遍跑完即得到行内引用.
//
// 调用约定 (用户必须遵守):
//  1. 必须在 Transaction 上下文中调用, 即 tx := db.BeginTransaction();
//  2. 流可被任意次 Read 消费; 库内部 io.Copy(buf=32KB) 循环, 全程峰值 ~128KB;
//  3. Snappy 压缩按块 (snappy.NewBufferedWriter) 自动分段, 调用方无感;
//  4. 调用方必须保证 compress 与将来的 BlobRead 一致 (库不持久化压缩标志).
//
// 中断保护:
//   - 流中途断 (conn closed, EOF 提前, Read 返错): 函数返 err, 此时 NewPage 出来的页
//     已在 BufferPool (脏页不刷盘) + tx.allocated 列表里;
//   - 调用方应 tx.Rollback(): 整链随 allocated 走"事务回滚归还空闲链"路径回收;
//   - 若 tx.Commit() 已被调用, 则链随提交落盘, 调用方需自管生命周期 (BlobFree).
//
// 内存峰值固定 (与 blob 大小无关):
//
//	io.Copy buf 32KB + snappy 内部缓冲 ~64KB + 链适配器临时页 buf 16KB ≈ 128KB.
//
// 容量上限:
//
//	BlobChainMax = 256K 页 ≈ 4 GB (原始数据); 超出返 ErrBlobTooLarge.
func (tx *Transaction) BlobCommit(r io.Reader, compress types.FieldCompress) (BlobRef, error) {
	if tx == nil || tx.db == nil {
		return BlobRef{}, errors.New("BlobCommit: tx is nil")
	}
	return tx.db.blobCommit(r, compress)
}

// blobCommit 是 BlobCommit 的无事务版本: 同样的流式写入, 但不要求调用方持有 Transaction.
//
// 为什么需要它: FieldType_Blob 列的行内只有 8 字节 BlobRef, 而用户往这一列写的是
// []byte / io.Reader 这种"内容本身" (见 Table.blobFieldRef). core 的 InsertRow / UpdateRow
// 在无显式事务时也必须能落地 —— 否则"写一行 Blob 列"就得先手动开事务, 与
// types/blob.go 里"引擎自动走链页"的设计目标背离.
//
// 与 tx.BlobCommit 的唯一差别在失败回收: 有 currentTxn 时新页登记进 tx.allocated, 回滚即回收;
// 无事务时没有 allocated 可回滚 —— 函数返回 err 后, 已分配的页只能由调用方保证不再复用
// (core 的列内联写入在出错时整行写入中止, 这批页成为孤儿, 与无事务下任何写入失败的语义一致).
func (db *DB) blobCommit(r io.Reader, compress types.FieldCompress) (BlobRef, error) {
	if db == nil {
		return BlobRef{}, errors.New("BlobCommit: db is nil")
	}
	if !compress.IsValid() {
		return BlobRef{}, errors.New("BlobCommit: unknown compress option")
	}
	if r == nil {
		// 与"空流"等价: 给空 blob
		return BlobRef{FirstPageID: 0, TotalLen: 0}, nil
	}

	// 0. 先探一个字节: 空输入连第一页都不分配 (否则空 blob 会白占一页 —— 提前开页是为了接住
	//    snappy 的 stream header, 但空输入压根没有 header 要接).
	var head [1]byte
	n, rerr := r.Read(head[:])
	if rerr != nil && rerr != io.EOF {
		return BlobRef{}, rerr
	}
	if n == 0 {
		return BlobRef{FirstPageID: 0, TotalLen: 0}, nil
	}
	src := io.MultiReader(bytes.NewReader(head[:n]), r)

	// 1. 预开第一页 (snappy.Writer 第一次 Write 会先写 stream header, 必须有页接住)
	w, err := newBlobChainWriter(db)
	if err != nil {
		return BlobRef{}, err
	}

	// 2. io.Copy 链: r -> [snappy?] -> blobChainWriter
	//    io.Copy 内部 32KB buf, snappy.NewBufferedWriter 内部 ~64KB, 都是固定大小.
	var totalRaw int64
	if compress == types.CompressSnappy {
		sw := snappy.NewBufferedWriter(w)
		totalRaw, err = io.Copy(sw, src)
		if cerr := sw.Close(); cerr != nil && err == nil {
			err = cerr // Close 才能 flush snappy 末尾
		}
	} else {
		totalRaw, err = io.Copy(w, src)
	}

	if err != nil {
		// 流中断 / io 错误: 链已分配但未 Commit, 由 tx.Rollback() 走 allocated 路径回收
		return BlobRef{}, err
	}

	if totalRaw == 0 {
		// 有输入却一个字节都没落地 (压缩后为空的极端情况): 把预开的页还回去, 给空引用.
		db.deferFreeBlobChain(w.firstPage, w.tailPage, w.chainLen)
		return BlobRef{FirstPageID: 0, TotalLen: 0}, nil
	}

	if uint64(totalRaw) > uint64(^uint32(0)) {
		return BlobRef{}, ErrBlobTooLarge // 单 blob 上限 4GB
	}
	return BlobRef{
		FirstPageID: w.firstPage,
		TotalLen:    uint32(totalRaw),
	}, nil
}