KyotoCabinet TreeDB大规模数据写入性能优化与基准测试策略

package main

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

    kc "github.com/vmihailenco/kyotocabinet" // 假设使用一个Go语言KyotoCabinet绑定库
)

// Pair 结构体用于存储键值对
type Pair struct {
    Key   string
    Value string
}

// genRandomString 生成指定长度的随机字符串
func genRandomString(length int) string {
    const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
    b := make([]byte, length)
    for i := range b {
        b[i] = charset[rand.Intn(len(charset))]
    }
    return string(b)
}

// setupRandomPairs 预生成指定数量的随机键值对
func setupRandomPairs(count int, keyLenRange, valLenRange int) []Pair {
    rand.Seed(time.Now().UnixNano()) // 初始化随机数种子
    pairs := make([]Pair, count)
    for i := 0; i < count; i++ {
        key := genRandomString(rand.Intn(keyLenRange) + 1) // 1到keyLenRange
        value := genRandomString(rand.Intn(valLenRange) + 1) // 1到valLenRange
        pairs[i] = Pair{Key: key, Value: value}
    }
    return pairs
}

// setupSequentialPairs 预生成指定数量的顺序递增键值对
func setupSequentialPairs(count int, valLen int) []Pair {
    pairs := make([]Pair, count)
    for i := 0; i < count; i++ {
        key := fmt.Sprintf("key%d", i)
        value := genRandomString(valLen) // 值仍然随机
        pairs[i] = Pair{Key: key, Value: value}
    }
    return pairs
}

func main() {
    const recordCount = 1000000 // 示例记录数，可根据需要调整
    const dbPath = "test.kct"

    // ---------------------------------------------------------------------
    // 阶段1: 数据预生成 (在计时前完成)
    fmt.Printf("Generating %d key-value pairs...\n", recordCount)
    // 可以选择生成随机键值对
    // allPairs := setupRandomPairs(recordCount, 1024, 1024) 
    // 或者生成顺序递增键值对进行对比测试
    allPairs := setupSequentialPairs(recordCount, 1024) 
    fmt.Printf("Data generation complete.\n")

    // ---------------------------------------------------------------------
    // 阶段2: 数据库初始化 (在计时前完成)
    // 清理旧的数据库文件，确保测试环境纯净
    os.Remove(dbPath) 

    // 打开TreeDB数据库
    db, err := kc.NewTreeDB()
    if err != nil {
        fmt.Printf("Failed to create TreeDB: %v\n", err)
        return
    }
    // 使用defer确保数据库最终关闭，即使程序发生错误
    defer func() {
        if db != nil {
            db.Close()
        }
        os.Remove(dbPath) // 清理数据库文件
        fmt.Printf("Database file '%s' cleaned up.\n", dbPath)
    }()

    // 配置并打开数据库，例如启用自动事务或设置缓存
    // kc.OWRITER | kc.OCREATE | kc.OTRUNCATE: 以写入模式打开，如果不存在则创建，如果存在则截断
    // kc.OAUTOTRAN: 启用自动事务，可以提高批量写入性能
    if !db.Open(dbPath, kc.OWRITER|kc.OCREATE|kc.OTRUNCATE|kc.OAUTOTRAN) { 
        fmt.Printf("Failed to open TreeDB: %s\n", db.Error().Error())
        return
    }

    // ---------------------------------------------------------------------
    // 阶段3: 核心写入操作计时
    fmt.Printf("Starting database write benchmark for %d records...\n", recordCount)
    startTime := time.Now()

    // 批量事务处理，每隔一定数量的写入提交一次事务，减少磁盘同步开销
    // 如果db.Open时使用了OAUTOTRAN，则可以省略手动事务管理
    // 如果没有使用OAUTOTRAN，则需要手动BeginTran/EndTran
    // const transactionBatchSize = 50000 
    // db.BeginTran() // 开始第一个事务

    for i, pair := range allPairs {
        // if i > 0 && i%transactionBatchSize == 0 {
        //  if !db.EndTran(true) { // 提交事务
        //      fmt.Printf("Failed to commit transaction at %d: %s\n", i, db.Error().Error())
        //      break
        //  }
        //  if !db.BeginTran() { // 开始新事务
        //      fmt.Printf("Failed to begin transaction at %d: %s\n", i, db.Error().Error())
        //      break
        //  }
        // }

        if !db.Set(pair.Key, pair.Value) {
            fmt.Printf("Failed to set key '%s': %s\n", pair.Key, db.Error().Error())
            break
        }
    }
    // if !db.EndTran(true) { // 提交最后一个事务（如果存在未提交的）
    //