TsDB2/demo/bench_ts/main.go

package main

// =============================================================================
// bench_ts — 时序表基准: 时序索引 (IndexTypeTime) 行存储的写入/范围扫描/点查
//
// 场景仿照典型时序负载: 递增时间戳 + 浮点数值 (+设备标签), 时间非递减写入。
//
//   TSDB_BENCH_TS_ROWS    行数        (默认 100 万)
//   TSDB_BENCH_TS_BATCH   每事务行数  (默认 1000, 控制 fsync 频率)
//   TSDB_BENCH_TS_QUERIES 查询采样数  (默认 200)
//   TSDB_BENCH_TS_SPAN    范围窗口大小(默认 1 万行)
//
// 输出: 写入行/秒、写入 MB/s、落盘文件大小、范围扫描行/秒、点查 QPS。
// 最后做一轮数据校验 (TotalRows 与行级抽查), 防止"测了个错的"。
// =============================================================================

import (
	"fmt"
	"math/rand"
	"os"
	"runtime/pprof"
	"time"

	"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 = "bench_ts.db"

func envInt(key string, def int) int {
	if v := os.Getenv(key); v != "" {
		var n int
		if _, err := fmt.Sscanf(v, "%d", &n); err == nil && n > 0 {
			return n
		}
	}
	return def
}

func main() {
	rows := envInt("TSDB_BENCH_TS_ROWS", 1_000_000)
	batch := envInt("TSDB_BENCH_TS_BATCH", 1000)
	queries := envInt("TSDB_BENCH_TS_QUERIES", 200)
	span := int64(envInt("TSDB_BENCH_TS_SPAN", 10000))

	_ = os.Remove(dbPath)
	db, err := core.Open(dbPath)
	if err != nil {
		fmt.Println("open:", err)
		os.Exit(1)
	}
	db.SetAutoCommitDelay(0)

	tbl, err := db.CreateTable("ticks", []types.FieldInfo{
		{FieldType: types.FieldType_Uint64, FieldName: []byte("ts")},
		{FieldType: types.FieldType_Float64, FieldName: []byte("price")},
		{FieldType: types.FieldType_String, Flag: types.WidthCodeFor(16), FieldName: []byte("dev")},
	})
	if err != nil {
		fmt.Println("create table:", err)
		os.Exit(1)
	}
	if _, err := tbl.CreateIndex("ts", page.IndexTypeTime, []byte{0}); err != nil {
		fmt.Println("create time index:", err)
		os.Exit(1)
	}

	const baseTS = 1_700_000_000_000 // 固定起点, 便于复现
	devs := []string{"alpha", "bravo", "charlie", "delta"}

	// ============================ 写入 ============================
	if pf := os.Getenv("TSDB_BENCH_TS_CPUPROFILE"); pf != "" {
		f, err := os.Create(pf)
		if err == nil {
			_ = pprof.StartCPUProfile(f)
			defer func() { pprof.StopCPUProfile(); _ = f.Close() }()
		}
	}
	fmt.Printf("写入阶段: %d 行, 每 %d 行一个事务...\n", rows, batch)
	start := time.Now()
	for i := int64(0); i < int64(rows); i++ {
		if _, err := tbl.InsertRowWithID(uint64(i+1), []types.FieldValue{
			{FieldName: []byte("ts"), Value: uint64(baseTS + i)},
			{FieldName: []byte("price"), Value: float64(i) * 1.5},
			{FieldName: []byte("dev"), Value: devs[i%4]},
		}); err != nil {
			fmt.Println("insert", i, ":", err)
			os.Exit(1)
		}
		if (i+1)%int64(batch) == 0 {
			if err := db.Commit(); err != nil {
				fmt.Println("commit:", err)
				os.Exit(1)
			}
		}
	}
	if err := db.Commit(); err != nil {
		fmt.Println("final commit:", err)
		os.Exit(1)
	}
	writeDur := time.Since(start)

	st, _ := os.Stat(dbPath)
	fileMB := float64(st.Size()) / 1024 / 1024
	fmt.Printf("  写入: %d 行 / %v = %s 行/秒, %.1f MB (%.2f MB/秒)\n",
		rows, writeDur.Round(time.Millisecond), comma(int(float64(rows)/writeDur.Seconds())), fileMB, fileMB/writeDur.Seconds())

	if total := tbl.TotalRows(); total != uint64(rows) {
		fmt.Printf("  ❌ TotalRows=%d, 想要 %d —— 数据校验失败\n", total, rows)
		os.Exit(1)
	}

	// ============================ 范围扫描 ============================
	rng := rand.New(rand.NewSource(42))
	warm := time.Now()
	scanOnce := func(from int64) (int, time.Duration) {
		n := 0
		s := time.Now()
		err := tbl.Query().Where(core.And(
			core.Gt("ts", uint64(from), types.FieldType_Uint64),
			core.Lt("ts", uint64(from+span), types.FieldType_Uint64),
		)).ScanWithPKID(func(pkID uint64, row []types.FieldValue) bool {
			n++
			return true
		})
		if err != nil {
			fmt.Println("range scan:", err)
			os.Exit(1)
		}
		return n, time.Since(s)
	}
	// 预热一次 (页缓存)
	if n, _ := scanOnce(baseTS); n != int(span)-1 {
		fmt.Printf("  ⚠️ 预热范围扫描返回 %d 行, 期望 %d\n", n, span-1)
	}
	_ = warm

	var scanRows int
	var scanDur time.Duration
	for q := 0; q < queries; q++ {
		from := baseTS + rng.Int63n(int64(rows)-span)
		n, d := scanOnce(from)
		scanRows += n
		scanDur += d
	}
	fmt.Printf("  范围扫描 (%d 次 × ~%s 行窗口): %s 行/秒, 平均 %v/次\n",
		queries, comma(int(span)), comma(int(float64(scanRows)/scanDur.Seconds())), (scanDur / time.Duration(queries)).Round(time.Microsecond))

	// ============================ 点查 ============================
	hit := 0
	ptStart := time.Now()
	for q := 0; q < queries; q++ {
		ts := uint64(baseTS + rng.Int63n(int64(rows)))
		vals, err := tbl.Query().Where(core.Eq("ts", ts, types.FieldType_Uint64)).Find()
		if err != nil {
			fmt.Println("point query:", err)
			os.Exit(1)
		}
		if len(vals) == 1 {
			hit++
		}
	}
	ptDur := time.Since(ptStart)
	fmt.Printf("  点查 (ts=随机): %d/%d 命中, QPS=%s, 平均 %v/次\n",
		hit, queries, comma(int(float64(queries)/ptDur.Seconds())), (ptDur / time.Duration(queries)).Round(time.Microsecond))

	fmt.Println("✅ 时序基准完成 (含 TotalRows 与范围行数校验)")
}

func comma(n int) string {
	s := fmt.Sprintf("%d", n)
	out := make([]byte, 0, len(s)+len(s)/3)
	for i, c := range s {
		if i > 0 && (len(s)-i)%3 == 0 {
			out = append(out, ',')
		}
		out = append(out, byte(c))
	}
	return string(out)
}