TsDB2/demo/blob_stream/main.go

package main

import (
	"bytes"
	"crypto/rand"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"git.ycl.pub/yl32940/TsDB2/DB"
	"git.ycl.pub/yl32940/TsDB2/DB/core"
	"git.ycl.pub/yl32940/TsDB2/DB/core/page"
	"git.ycl.pub/yl32940/TsDB2/DB/core/types"
)

const dbPath = "blob_demo.db"

// brokenStream 模拟"前 N 字节正常, 之后 io.ErrUnexpectedEOF"(网络中断 / 文件被截断).
type brokenStream struct {
	consumed int
	failAt   int
}

func (b *brokenStream) Read(p []byte) (int, error) {
	if b.consumed >= b.failAt {
		return 0, io.ErrUnexpectedEOF
	}
	n := b.failAt - b.consumed
	if n > len(p) {
		n = len(p)
	}
	got, _ := rand.Read(p[:n])
	b.consumed += got
	return got, nil
}

func must(err error) {
	if err != nil {
		fmt.Println("ERR:", err)
		os.Exit(1)
	}
}

// text 把定长字符串列读回的原样字节转成可打印文本 (定长列会补零到列宽).
func text(v interface{}) string {
	b, ok := v.([]byte)
	if !ok {
		return fmt.Sprintf("%v", v)
	}
	return strings.TrimRight(string(b), "\x00")
}

// fieldValue 从读回的行里按列名取值 (找不到返 nil).
func fieldValue(fvs []types.FieldValue, name string) interface{} {
	for _, fv := range fvs {
		if string(fv.FieldName) == name {
			return fv.Value
		}
	}
	return nil
}

// blobRef 取出一行的 blob 列引用, 并按 compress 把链页内容读回来。
func blobRef(t *core.Table, fvs []types.FieldValue, name string, compress types.FieldCompress) (types.BlobRef, []byte, error) {
	ref, ok := fieldValue(fvs, name).(types.BlobRef)
	if !ok {
		return types.BlobRef{}, nil, fmt.Errorf("列 %s 读出的不是 types.BlobRef (内容被内联进行里了?)", name)
	}
	got, err := t.BlobRead(ref, compress)
	return ref, got, err
}

func main() {
	os.Remove(dbPath)

	fmt.Println("=== TsDB2 Blob 流式存储 demo ===")
	fmt.Println("blob 是【普通的一列】: 与 name/size 同在一行, 内容单独走链页。")
	fmt.Println("写入给 []byte / io.Reader / types.Blob 都行, 引擎自动流式落链, 不用手动开事务。")
	fmt.Println("行号 (pk) 由引擎自增分配, InsertRow 的返回值就是它 —— 不用自己造 id 列。")
	fmt.Println()

	db, err := DB.Open(dbPath)
	must(err)
	db.SetAutoCommitDelay(0) // 关掉防抖自动提交: 每步手动 Commit, 便于对照

	// 两列普通列 + 一列 blob。blob 列在行内只占 8 字节引用, 与其它列平级。
	tbl, err := db.CoreDB().CreateTable("files", []types.FieldInfo{
		{FieldType: types.FieldType_String, FieldName: []byte("name"), Flag: types.WidthCodeFor(32)},
		{FieldType: types.FieldType_Int64, FieldName: []byte("size")},
		{FieldType: types.FieldType_Blob, FieldName: []byte("data"), Flag: types.FlagNullable},
	})
	must(err)
	// 主键 B+ 树的键始终是行号 pkID (与索引声明在哪个字段上无关, 见 PrimaryHandler),
	// 建了它 FindByPK / UpdateRowByPK / DeleteRowByPK 才能按行号定位。
	_, err = tbl.CreateIndex("pk", page.IndexTypePrimary, []byte{0})
	must(err)

	// ------------------------------------------------------------------
	fmt.Println("[1] 插一行: name + size + data([]byte) —— 多列共存, 行号自增")
	payload := bytes.Repeat([]byte("hello blob "), 512) // ~5.5 KB
	pk1, err := tbl.InsertRow([]types.FieldValue{
		{FieldName: []byte("name"), Value: "readme.txt"},
		{FieldName: []byte("size"), Value: int64(len(payload))},
		{FieldName: []byte("data"), Value: payload},
	})
	must(err)
	must(db.Commit())

	fvs, err := tbl.FindByPK(pk1)
	must(err)
	ref1, got1, err := blobRef(tbl, fvs, "data", types.CompressNone)
	must(err)
	fmt.Printf("    行号 pk=%d; 读回整行: name=%q size=%d data=%d 字节\n",
		pk1, text(fieldValue(fvs, "name")), fieldValue(fvs, "size").(int64), len(got1))
	fmt.Printf("    行内只放 8B 引用: ref=%+v; 内容逐字节一致: %v\n", ref1, bytes.Equal(got1, payload))
	fmt.Println()

	// ------------------------------------------------------------------
	fmt.Println("[2] 流式插入: 给 io.Reader (超过 Dynamicbyte 的单块上限, 链能跨页)")
	streamLen := int(core.MaxBlobSize)*2 + 7
	pk2, err := tbl.InsertRow([]types.FieldValue{
		{FieldName: []byte("name"), Value: "random.bin"},
		{FieldName: []byte("size"), Value: int64(streamLen)},
		// 流在这里被引擎 io.Copy (32KB buf) 消费完, 内存峰值与 blob 大小无关。
		{FieldName: []byte("data"), Value: io.LimitReader(rand.Reader, int64(streamLen))},
	})
	must(err)
	must(db.Commit())

	fvs, err = tbl.FindByPK(pk2)
	must(err)
	ref2, got2, err := blobRef(tbl, fvs, "data", types.CompressNone)
	must(err)
	fmt.Printf("    pk=%d 写入 %d 字节 (单块堆页上限 %d) -> ref.TotalLen=%d 首页=%d, 读回 %d 字节\n",
		pk2, streamLen, core.MaxBlobSize, ref2.TotalLen, ref2.FirstPageID, len(got2))
	fmt.Println()

	// ------------------------------------------------------------------
	fmt.Println("[3] 更新同一行的 blob 列: 换内容时旧链自动归还空闲链")
	next := []byte("second version, 换了内容")
	err = tbl.UpdateRowByPK(pk2, []types.FieldValue{
		// 只传 data: 其它列未提供 = 沿用旧行原值 (不会清空)
		{FieldName: []byte("data"), Value: next},
	})
	must(err)
	must(db.Commit())

	fvs, err = tbl.FindByPK(pk2)
	must(err)
	ref3, got3, err := blobRef(tbl, fvs, "data", types.CompressNone)
	must(err)
	fmt.Printf("    链换了: 更新前 ref=%+v / 更新后 ref=%+v\n", ref2, ref3)
	fmt.Printf("    内容: %q; 其它列仍在: name=%q size=%d\n",
		string(got3), text(fieldValue(fvs, "name")), fieldValue(fvs, "size").(int64))
	fmt.Println()

	// ------------------------------------------------------------------
	fmt.Println("[4] 流中断: 100KB 流读到 50KB 断 —— 插入失败 + Rollback, 链页归还")
	before := tbl.TotalRows()
	tx := db.BeginTransaction()
	_, err = tbl.InsertRow([]types.FieldValue{
		{FieldName: []byte("name"), Value: "broken.bin"},
		{FieldName: []byte("size"), Value: int64(100_000)},
		{FieldName: []byte("data"), Value: &brokenStream{failAt: 50_000}},
	})
	if err == nil {
		fmt.Println("    ERR: 期望流错误, 实际成功")
		os.Exit(1)
	}
	fmt.Printf("    插入失败: %v\n", err)
	must(tx.Rollback())
	fmt.Printf("    回滚后行数仍是 %d (失败的行未生效, 已分配的链页随事务归还)\n", tbl.TotalRows())
	if tbl.TotalRows() != before {
		fmt.Println("    ERR: 行数被改动了")
		os.Exit(1)
	}
	if _, err := tbl.InsertRow([]types.FieldValue{
		{FieldName: []byte("name"), Value: "after-rollback.txt"},
		{FieldName: []byte("size"), Value: int64(4)},
		{FieldName: []byte("data"), Value: []byte("ok")},
	}); err != nil {
		fmt.Println("    ERR: 回滚后库不可用:", err)
		os.Exit(1)
	}
	must(db.Commit())
	fmt.Println("    回滚后照常插入新行: 成功")
	fmt.Println()

	// ------------------------------------------------------------------
	fmt.Println("[5] 100 MB 随机流 + Snappy: 内存峰值固定 (~128KB, 与 blob 大小无关)")
	const MB = 1024 * 1024
	t0 := time.Now()
	pkBig, err := tbl.InsertRow([]types.FieldValue{
		{FieldName: []byte("name"), Value: "big.snappy"},
		{FieldName: []byte("size"), Value: int64(100 * MB)},
		// 压缩策略是运行时参数 (不落盘), 读侧必须给同一个 compress。
		{FieldName: []byte("data"), Value: types.Blob{
			Source:   io.LimitReader(rand.Reader, 100*MB),
			Compress: types.CompressSnappy,
		}},
	})
	must(err)
	must(db.Commit())

	fvs, err = tbl.FindByPK(pkBig)
	must(err)
	ref4, got4, err := blobRef(tbl, fvs, "data", types.CompressSnappy)
	must(err)
	fmt.Printf("    pk=%d 写入耗时 %.1fs; ref=%+v; 读回 %d 字节 (链上 %d 字节)\n",
		pkBig, time.Since(t0).Seconds(), ref4, len(got4), ref4.TotalLen)
	fmt.Println()

	// ------------------------------------------------------------------
	fmt.Println("[6] 关库重开再读: 表是从盘上加载的, blob 仍读得回来")
	must(db.Close())
	db2, err := DB.Open(dbPath)
	must(err)
	db2.SetAutoCommitDelay(0)
	tbl2, ok := db2.CoreDB().FindTable("files")
	if !ok {
		fmt.Println("    ERR: 重开后找不到表 files")
		os.Exit(1)
	}
	fvs, err = tbl2.FindByPK(pk1)
	must(err)
	_, got6, err := blobRef(tbl2, fvs, "data", types.CompressNone)
	must(err)
	fmt.Printf("    pk=%d 内容与写入一致: %v (前 16 字节 %q)\n",
		pk1, bytes.Equal(got6, payload), string(got6[:min(16, len(got6))]))
	must(db2.Close())

	fmt.Println()
	fmt.Println("=== 全部场景通过 ===")
}