srdb/compaction_stage_test.go

package srdb

import (
	"testing"
)

// TestPickerStageRotation 测试 Picker 的阶段轮换机制
func TestPickerStageRotation(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
	if stage := picker.GetCurrentStage(); stage != 0 {
		t.Errorf("Initial stage should be 0 (L0), got %d", stage)
	}

	// 添加 L0 文件（触发 L0 compaction）
	edit := NewVersionEdit()
	for i := 0; i < 10; i++ {
		edit.AddFile(&FileMetadata{
			FileNumber: int64(i + 1),
			Level:      0,
			FileSize:   10 * 1024 * 1024, // 10MB each
			MinKey:     int64(i * 100),
			MaxKey:     int64((i+1)*100 - 1),
			RowCount:   100,
		})
	}
	edit.SetNextFileNumber(11)
	err = versionSet.LogAndApply(edit)
	if err != nil {
		t.Fatal(err)
	}

	version := versionSet.GetCurrent()

	// 第1次调用：应该返回 L0 任务，然后推进到 L1
	t.Log("=== 第1次调用 PickCompaction ===")
	tasks1 := picker.PickCompaction(version)
	if len(tasks1) == 0 {
		t.Error("Expected L0 tasks on first call")
	}
	for _, task := range tasks1 {
		if task.Level != 0 {
			t.Errorf("Expected L0 task, got L%d", task.Level)
		}
	}
	if stage := picker.GetCurrentStage(); stage != 1 {
		t.Errorf("After L0 tasks, stage should be 1 (L1), got %d", stage)
	}
	t.Logf("✓ Returned %d L0 tasks, stage advanced to L1", len(tasks1))

	// 第2次调用：应该尝试 Stage 1 (L0-upgrade，没有大文件)
	t.Log("=== 第2次调用 PickCompaction ===")
	tasks2 := picker.PickCompaction(version)
	if len(tasks2) == 0 {
		t.Log("✓ Stage 1 (L0-upgrade) has no tasks")
	}
	// 此时 stage 应该已经循环（尝试了 Stage 1→2→3→0...）
	if stage := picker.GetCurrentStage(); stage >= 0 {
		t.Logf("After trying, current stage is %d", stage)
	}

	// 现在添加 L1 文件
	edit2 := NewVersionEdit()
	for i := 0; i < 20; i++ {
		edit2.AddFile(&FileMetadata{
			FileNumber: int64(100 + i + 1),
			Level:      1,
			FileSize:   20 * 1024 * 1024, // 20MB each
			MinKey:     int64(i * 200),
			MaxKey:     int64((i+1)*200 - 1),
			RowCount:   200,
		})
	}
	edit2.SetNextFileNumber(121)
	err = versionSet.LogAndApply(edit2)
	if err != nil {
		t.Fatal(err)
	}

	version2 := versionSet.GetCurrent()

	// 现在可能需要多次调用才能到达 Stage 2 (L1-upgrade)
	// 因为要经过 Stage 1 (L0-upgrade) 和 Stage 0 (L0-merge)
	t.Log("=== 多次调用 PickCompaction 直到找到 L1 任务 ===")
	var tasks3 []*CompactionTask
	for i := 0; i < 8; i++ { // 最多尝试两轮（4个阶段×2）
		tasks3 = picker.PickCompaction(version2)
		if len(tasks3) > 0 && tasks3[0].Level == 1 {
			t.Logf("✓ Found %d L1 tasks after %d attempts", len(tasks3), i+1)
			break
		}
	}
	if len(tasks3) == 0 || tasks3[0].Level != 1 {
		t.Error("Expected to find L1 tasks within 8 attempts")
	}

	t.Log("=== Stage rotation test passed ===")
}

// TestPickerStageWithMultipleLevels 测试多层级同时有任务时的阶段轮换
func TestPickerStageWithMultipleLevels(t *testing.T) {
	tmpDir := t.TempDir()
	manifestDir := tmpDir

	versionSet, err := NewVersionSet(manifestDir)
	if err != nil {
		t.Fatal(err)
	}
	defer versionSet.Close()

	picker := NewPicker()

	// 同时添加 L0、L1、L2 文件
	edit := NewVersionEdit()

	// L0 小文件: 5 files × 10MB = 50MB (应该触发 Stage 0: L0-merge)
	for i := 0; i < 5; i++ {
		edit.AddFile(&FileMetadata{
			FileNumber: int64(i + 1),
			Level:      0,
			FileSize:   10 * 1024 * 1024,
			MinKey:     int64(i * 100),
			MaxKey:     int64((i+1)*100 - 1),
			RowCount:   100,
		})
	}

	// L0 大文件: 5 files × 40MB = 200MB (应该触发 Stage 1: L0-upgrade)
	for i := 0; i < 5; i++ {
		edit.AddFile(&FileMetadata{
			FileNumber: int64(10 + i + 1),
			Level:      0,
			FileSize:   40 * 1024 * 1024,
			MinKey:     int64((i+5) * 100),
			MaxKey:     int64((i+6)*100 - 1),
			RowCount:   100,
		})
	}

	// L1: 20 files × 20MB = 400MB (应该触发 Stage 2: L1-upgrade，256MB阈值)
	for i := 0; i < 20; i++ {
		edit.AddFile(&FileMetadata{
			FileNumber: int64(100 + i + 1),
			Level:      1,
			FileSize:   20 * 1024 * 1024,
			MinKey:     int64(i * 200),
			MaxKey:     int64((i+1)*200 - 1),
			RowCount:   200,
		})
	}

	// L2: 10 files × 150MB = 1500MB (应该触发 Stage 3: L2-upgrade，1GB阈值)
	for i := 0; i < 10; i++ {
		edit.AddFile(&FileMetadata{
			FileNumber: int64(200 + i + 1),
			Level:      2,
			FileSize:   150 * 1024 * 1024,
			MinKey:     int64(i * 300),
			MaxKey:     int64((i+1)*300 - 1),
			RowCount:   300,
		})
	}

	edit.SetNextFileNumber(301)
	err = versionSet.LogAndApply(edit)
	if err != nil {
		t.Fatal(err)
	}

	version := versionSet.GetCurrent()

	// 验证阶段按顺序执行：Stage 0→1→2→3→0→1→2→3
	expectedStages := []struct {
		stage int
		name  string
		level int
	}{
		{0, "L0-merge", 0},
		{1, "L0-upgrade", 0},
		{2, "L1-upgrade", 1},
		{3, "L2-upgrade", 2},
		{0, "L0-merge", 0},
		{1, "L0-upgrade", 0},
		{2, "L1-upgrade", 1},
		{3, "L2-upgrade", 2},
	}

	for i, expected := range expectedStages {
		t.Logf("=== 第%d次调用 PickCompaction (期望 Stage %d: %s) ===", i+1, expected.stage, expected.name)
		tasks := picker.PickCompaction(version)

		if len(tasks) == 0 {
			t.Errorf("Call %d: Expected tasks from Stage %d (%s), got no tasks", i+1, expected.stage, expected.name)
			continue
		}

		actualLevel := tasks[0].Level
		if actualLevel != expected.level {
			t.Errorf("Call %d: Expected L%d tasks, got L%d tasks", i+1, expected.level, actualLevel)
		} else {
			t.Logf("✓ Call %d: Got %d tasks from L%d (Stage %d: %s) as expected",
				i+1, len(tasks), actualLevel, expected.stage, expected.name)
		}
	}

	t.Log("=== Multi-level stage rotation test passed ===")
}