package srdb import ( "fmt" "os" "path/filepath" "testing" "time" ) func TestCompactionBasic(t *testing.T) { // 创建临时目录 tmpDir := t.TempDir() sstDir := filepath.Join(tmpDir, "sst") manifestDir := tmpDir err := os.MkdirAll(sstDir, 0755) if err != nil { t.Fatal(err) } // 创建 VersionSet versionSet, err := NewVersionSet(manifestDir) if err != nil { t.Fatal(err) } defer versionSet.Close() // 创建 SST Manager sstMgr, err := NewSSTableManager(sstDir) if err != nil { t.Fatal(err) } defer sstMgr.Close() // 创建测试数据 rows1 := make([]*SSTableRow, 100) for i := range 100 { rows1[i] = &SSTableRow{ Seq: int64(i), Time: 1000, Data: map[string]any{"value": i}, } } // 创建第一个 SST 文件 reader1, err := sstMgr.CreateSST(1, rows1) if err != nil { t.Fatal(err) } // 添加到 Version edit1 := NewVersionEdit() edit1.AddFile(&FileMetadata{ FileNumber: 1, Level: 0, FileSize: 1024, MinKey: 0, MaxKey: 99, RowCount: 100, }) nextFileNum := int64(2) edit1.SetNextFileNumber(nextFileNum) err = versionSet.LogAndApply(edit1) if err != nil { t.Fatal(err) } // 验证 Version version := versionSet.GetCurrent() if version.GetLevelFileCount(0) != 1 { t.Errorf("Expected 1 file in L0, got %d", version.GetLevelFileCount(0)) } // 创建 Compaction Manager compactionMgr := NewCompactionManager(sstDir, versionSet, sstMgr) // 创建更多文件触发 Compaction for i := 1; i < 5; i++ { rows := make([]*SSTableRow, 50) for j := range 50 { rows[j] = &SSTableRow{ Seq: int64(i*100 + j), Time: int64(1000 + i), Data: map[string]any{"value": i*100 + j}, } } _, err := sstMgr.CreateSST(int64(i+1), rows) if err != nil { t.Fatal(err) } edit := NewVersionEdit() edit.AddFile(&FileMetadata{ FileNumber: int64(i + 1), Level: 0, FileSize: 512, MinKey: int64(i * 100), MaxKey: int64(i*100 + 49), RowCount: 50, }) nextFileNum := int64(i + 2) edit.SetNextFileNumber(nextFileNum) err = versionSet.LogAndApply(edit) if err != nil { t.Fatal(err) } } // 验证 L0 有 5 个文件 version = versionSet.GetCurrent() if version.GetLevelFileCount(0) != 5 { t.Errorf("Expected 5 files in L0, got %d", version.GetLevelFileCount(0)) } // 检查是否需要 Compaction picker := compactionMgr.GetPicker() if !picker.ShouldCompact(version) { t.Error("Expected compaction to be needed") } // 获取 Compaction 任务 tasks := picker.PickCompaction(version) if len(tasks) == 0 { t.Fatal("Expected compaction task") } task := tasks[0] // 获取第一个任务(优先级最高) if task.Level != 0 { t.Errorf("Expected L0 compaction, got L%d", task.Level) } if task.OutputLevel != 1 { t.Errorf("Expected output to L1, got L%d", task.OutputLevel) } t.Logf("Found %d compaction tasks", len(tasks)) t.Logf("First task: L%d -> L%d, %d files", task.Level, task.OutputLevel, len(task.InputFiles)) // 清理 reader1.Close() } func TestPickerLevelScore(t *testing.T) { // 创建临时目录 tmpDir := t.TempDir() manifestDir := tmpDir // 创建 VersionSet versionSet, err := NewVersionSet(manifestDir) if err != nil { t.Fatal(err) } defer versionSet.Close() // 创建 Picker picker := NewPicker() // 添加一些文件到 L0 edit := NewVersionEdit() for i := range 3 { edit.AddFile(&FileMetadata{ FileNumber: int64(i + 1), Level: 0, FileSize: 1024 * 1024, // 1MB MinKey: int64(i * 100), MaxKey: int64((i+1)*100 - 1), RowCount: 100, }) } nextFileNum := int64(4) edit.SetNextFileNumber(nextFileNum) err = versionSet.LogAndApply(edit) if err != nil { t.Fatal(err) } version := versionSet.GetCurrent() // 计算 L0 的得分 score := picker.GetLevelScore(version, 0) t.Logf("L0 score: %.2f (files: %d, limit: %d)", score, version.GetLevelFileCount(0), picker.levelFileLimits[0]) // L0 有 3 个文件,限制是 4,得分应该是 0.75 expectedScore := 3.0 / 4.0 if score != expectedScore { t.Errorf("Expected L0 score %.2f, got %.2f", expectedScore, score) } } func TestCompactionMerge(t *testing.T) { // 创建临时目录 tmpDir := t.TempDir() sstDir := filepath.Join(tmpDir, "sst") manifestDir := tmpDir err := os.MkdirAll(sstDir, 0755) if err != nil { t.Fatal(err) } // 创建 VersionSet versionSet, err := NewVersionSet(manifestDir) if err != nil { t.Fatal(err) } defer versionSet.Close() // 创建 SST Manager sstMgr, err := NewSSTableManager(sstDir) if err != nil { t.Fatal(err) } defer sstMgr.Close() // 创建两个有重叠 key 的 SST 文件 rows1 := []*SSTableRow{ {Seq: 1, Time: 1000, Data: map[string]any{"value": "old"}}, {Seq: 2, Time: 1000, Data: map[string]any{"value": "old"}}, } rows2 := []*SSTableRow{ {Seq: 1, Time: 2000, Data: map[string]any{"value": "new"}}, // 更新 {Seq: 3, Time: 2000, Data: map[string]any{"value": "new"}}, } reader1, err := sstMgr.CreateSST(1, rows1) if err != nil { t.Fatal(err) } defer reader1.Close() reader2, err := sstMgr.CreateSST(2, rows2) if err != nil { t.Fatal(err) } defer reader2.Close() // 添加到 Version edit := NewVersionEdit() edit.AddFile(&FileMetadata{ FileNumber: 1, Level: 0, FileSize: 512, MinKey: 1, MaxKey: 2, RowCount: 2, }) edit.AddFile(&FileMetadata{ FileNumber: 2, Level: 0, FileSize: 512, MinKey: 1, MaxKey: 3, RowCount: 2, }) nextFileNum := int64(3) edit.SetNextFileNumber(nextFileNum) err = versionSet.LogAndApply(edit) if err != nil { t.Fatal(err) } // 创建 Compactor compactor := NewCompactor(sstDir, versionSet) // 创建 Compaction 任务 version := versionSet.GetCurrent() task := &CompactionTask{ Level: 0, InputFiles: version.GetLevel(0), OutputLevel: 1, } // 执行 Compaction resultEdit, err := compactor.DoCompaction(task, version) if err != nil { t.Fatal(err) } // 验证结果 if len(resultEdit.DeletedFiles) != 2 { t.Errorf("Expected 2 deleted files, got %d", len(resultEdit.DeletedFiles)) } if len(resultEdit.AddedFiles) == 0 { t.Error("Expected at least 1 new file") } t.Logf("Compaction result: deleted %d files, added %d files", len(resultEdit.DeletedFiles), len(resultEdit.AddedFiles)) // 验证新文件在 L1 for _, file := range resultEdit.AddedFiles { if file.Level != 1 { t.Errorf("Expected new file in L1, got L%d", file.Level) } t.Logf("New file: %d, L%d, rows: %d, key range: [%d, %d]", file.FileNumber, file.Level, file.RowCount, file.MinKey, file.MaxKey) } } func BenchmarkCompaction(b *testing.B) { // 创建临时目录 tmpDir := b.TempDir() sstDir := filepath.Join(tmpDir, "sst") manifestDir := tmpDir err := os.MkdirAll(sstDir, 0755) if err != nil { b.Fatal(err) } // 创建 VersionSet versionSet, err := NewVersionSet(manifestDir) if err != nil { b.Fatal(err) } defer versionSet.Close() // 创建 SST Manager sstMgr, err := NewSSTableManager(sstDir) if err != nil { b.Fatal(err) } defer sstMgr.Close() // 创建测试数据 const numFiles = 5 const rowsPerFile = 1000 for i := range numFiles { rows := make([]*SSTableRow, rowsPerFile) for j := range rowsPerFile { rows[j] = &SSTableRow{ Seq: int64(i*rowsPerFile + j), Time: int64(1000 + i), Data: map[string]any{ "value": fmt.Sprintf("data-%d-%d", i, j), }, } } reader, err := sstMgr.CreateSST(int64(i+1), rows) if err != nil { b.Fatal(err) } reader.Close() edit := NewVersionEdit() edit.AddFile(&FileMetadata{ FileNumber: int64(i + 1), Level: 0, FileSize: 10240, MinKey: int64(i * rowsPerFile), MaxKey: int64((i+1)*rowsPerFile - 1), RowCount: rowsPerFile, }) nextFileNum := int64(i + 2) edit.SetNextFileNumber(nextFileNum) err = versionSet.LogAndApply(edit) if err != nil { b.Fatal(err) } } // 创建 Compactor compactor := NewCompactor(sstDir, versionSet) version := versionSet.GetCurrent() task := &CompactionTask{ Level: 0, InputFiles: version.GetLevel(0), OutputLevel: 1, } for b.Loop() { _, err := compactor.DoCompaction(task, version) if err != nil { b.Fatal(err) } } } // TestCompactionQueryOrder 测试 compaction 后查询结果的排序 func TestCompactionQueryOrder(t *testing.T) { // 创建临时目录 tmpDir := t.TempDir() // 创建 Schema - 包含多个字段以增加数据大小 schema := NewSchema("test", []Field{ {Name: "id", Type: FieldTypeInt64}, {Name: "name", Type: FieldTypeString}, {Name: "data", Type: FieldTypeString}, {Name: "timestamp", Type: FieldTypeInt64}, }) // 打开 Engine (使用较小的 MemTable 触发频繁 flush) engine, err := OpenEngine(&EngineOptions{ Dir: tmpDir, MemTableSize: 2 * 1024 * 1024, // 2MB MemTable Schema: schema, }) if err != nil { t.Fatal(err) } defer engine.Close() t.Logf("开始插入 4000 条数据...") // 插入 4000 条数据,每条数据大小在 2KB-1MB 之间 for i := range 4000 { // 生成 2KB 到 1MB 的随机数据 dataSize := 2*1024 + (i % (1024*1024 - 2*1024)) // 2KB ~ 1MB largeData := make([]byte, dataSize) for j := range largeData { largeData[j] = byte('A' + (j % 26)) } err := engine.Insert(map[string]any{ "id": int64(i), "name": fmt.Sprintf("user_%d", i), "data": string(largeData), "timestamp": int64(1000000 + i), }) if err != nil { t.Fatal(err) } if (i+1)%500 == 0 { t.Logf("已插入 %d 条数据", i+1) } } t.Logf("插入完成,等待后台 compaction...") // 等待一段时间让后台 compaction 有机会运行 // 后台 compaction 每 10 秒检查一次,所以需要等待至少 12 秒 time.Sleep(12 * time.Second) t.Logf("开始查询所有数据...") // 查询所有数据 rows, err := engine.Query().Rows() if err != nil { t.Fatal(err) } defer rows.Close() // 验证顺序和数据完整性 var lastSeq int64 = 0 count := 0 expectedIDs := make(map[int64]bool) // 用于验证所有 ID 都存在 for rows.Next() { row := rows.Row() data := row.Data() currentSeq := data["_seq"].(int64) // 验证顺序 if currentSeq <= lastSeq { t.Errorf("Query results NOT in order: got seq %d after seq %d", currentSeq, lastSeq) } // 验证数据完整性 id, ok := data["id"].(int64) if !ok { // 尝试其他类型 if idFloat, ok2 := data["id"].(float64); ok2 { id = int64(idFloat) expectedIDs[id] = true } else { t.Errorf("Seq %d: missing or invalid id field, actual type: %T, value: %v", currentSeq, data["id"], data["id"]) } } else { expectedIDs[id] = true } // 验证 name 字段 name, ok := data["name"].(string) if !ok || name != fmt.Sprintf("user_%d", id) { t.Errorf("Seq %d: invalid name field, expected 'user_%d', got '%v'", currentSeq, id, name) } // 验证 data 字段存在且不为空 dataStr, ok := data["data"].(string) if !ok || len(dataStr) < 2*1024 { t.Errorf("Seq %d: invalid data field size", currentSeq) } lastSeq = currentSeq count++ } if count != 4000 { t.Errorf("Expected 4000 rows, got %d", count) } // 验证所有 ID 都存在 for i := range int64(4000) { if !expectedIDs[i] { t.Errorf("Missing ID: %d", i) } } t.Logf("✓ 查询返回 %d 条记录,顺序正确 (seq 1→%d)", count, lastSeq) t.Logf("✓ 所有数据完整性验证通过") // 输出 compaction 统计信息 stats := engine.GetCompactionManager().GetLevelStats() t.Logf("Compaction 统计:") for _, levelStat := range stats { level := levelStat["level"].(int) fileCount := levelStat["file_count"].(int) totalSize := levelStat["total_size"].(int64) if fileCount > 0 { t.Logf(" L%d: %d 个文件, %.2f MB", level, fileCount, float64(totalSize)/(1024*1024)) } } }