You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

524 lines
15 KiB

package task
import (
"fmt"
"testing"
"github.com/seaweedfs/seaweedfs/weed/worker/types"
)
func TestAdminServer_TaskAssignmentWithStateManagement(t *testing.T) {
// Test the core functionality: accurate task assignment based on comprehensive state
adminServer := NewAdminServer(DefaultAdminConfig(), nil)
// Initialize components
adminServer.workerRegistry = NewWorkerRegistry()
adminServer.taskQueue = NewPriorityTaskQueue()
adminServer.volumeStateManager = NewVolumeStateManager(nil)
adminServer.taskScheduler = NewTaskScheduler(adminServer.workerRegistry, adminServer.taskQueue)
adminServer.running = true // Mark as running for test
// Setup test worker
worker := &types.Worker{
ID: "test_worker_1",
Address: "server1:8080",
Capabilities: []types.TaskType{types.TaskTypeErasureCoding, types.TaskTypeVacuum},
MaxConcurrent: 2,
Status: "active",
CurrentLoad: 0,
}
adminServer.workerRegistry.RegisterWorker(worker)
// Setup volume state
volumeID := uint32(1)
adminServer.volumeStateManager.volumes[volumeID] = &VolumeState{
VolumeID: volumeID,
CurrentState: &VolumeInfo{
ID: volumeID,
Size: 28 * 1024 * 1024 * 1024, // 28GB - good for EC
Server: "server1",
},
InProgressTasks: []*TaskImpact{},
PlannedChanges: []*PlannedOperation{},
}
// Setup server capacity
adminServer.volumeStateManager.capacityCache["server1"] = &CapacityInfo{
Server: "server1",
TotalCapacity: 100 * 1024 * 1024 * 1024, // 100GB
UsedCapacity: 50 * 1024 * 1024 * 1024, // 50GB used
PredictedUsage: 50 * 1024 * 1024 * 1024, // Initially same as used
}
// Create EC task
task := &types.Task{
ID: "ec_task_1",
Type: types.TaskTypeErasureCoding,
VolumeID: volumeID,
Server: "server1",
Priority: types.TaskPriorityNormal,
}
// Test task assignment
adminServer.taskQueue.Push(task)
assignedTask, err := adminServer.RequestTask("test_worker_1", []types.TaskType{types.TaskTypeErasureCoding})
if err != nil {
t.Errorf("Task assignment failed: %v", err)
}
if assignedTask == nil {
t.Fatal("Expected task to be assigned, got nil")
}
if assignedTask.ID != "ec_task_1" {
t.Errorf("Expected task ec_task_1, got %s", assignedTask.ID)
}
// Verify state manager was updated
if len(adminServer.volumeStateManager.inProgressTasks) != 1 {
t.Errorf("Expected 1 in-progress task in state manager, got %d", len(adminServer.volumeStateManager.inProgressTasks))
}
// Verify capacity reservation
capacity := adminServer.volumeStateManager.GetAccurateCapacity("server1")
if capacity.ReservedCapacity <= 0 {
t.Error("Expected capacity to be reserved for EC task")
}
t.Log("✅ Task assignment with state management test passed")
}
func TestAdminServer_CanAssignTask(t *testing.T) {
adminServer := NewAdminServer(DefaultAdminConfig(), nil)
adminServer.volumeStateManager = NewVolumeStateManager(nil)
adminServer.inProgressTasks = make(map[string]*InProgressTask)
// Setup volume state
volumeID := uint32(1)
adminServer.volumeStateManager.volumes[volumeID] = &VolumeState{
VolumeID: volumeID,
CurrentState: &VolumeInfo{
ID: volumeID,
Size: 25 * 1024 * 1024 * 1024, // 25GB
},
}
// Setup server capacity - limited space
serverID := "server1"
adminServer.volumeStateManager.capacityCache[serverID] = &CapacityInfo{
Server: serverID,
TotalCapacity: 30 * 1024 * 1024 * 1024, // 30GB total
UsedCapacity: 20 * 1024 * 1024 * 1024, // 20GB used
PredictedUsage: 20 * 1024 * 1024 * 1024, // 10GB available
}
worker := &types.Worker{
ID: "worker1",
Address: serverID,
}
tests := []struct {
name string
taskType types.TaskType
expected bool
desc string
}{
{
name: "EC task fits",
taskType: types.TaskTypeErasureCoding,
expected: false, // 25GB * 1.4 = 35GB needed, but only 10GB available
desc: "EC task should not fit due to insufficient capacity",
},
{
name: "Vacuum task fits",
taskType: types.TaskTypeVacuum,
expected: true,
desc: "Vacuum task should fit (no capacity increase)",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
task := &types.Task{
ID: "test_task",
Type: tt.taskType,
VolumeID: volumeID,
Server: serverID,
}
result := adminServer.canAssignTask(task, worker)
if result != tt.expected {
t.Errorf("canAssignTask() = %v, want %v. %s", result, tt.expected, tt.desc)
}
})
}
}
func TestAdminServer_CreateTaskImpact(t *testing.T) {
adminServer := NewAdminServer(DefaultAdminConfig(), nil)
adminServer.volumeStateManager = NewVolumeStateManager(nil)
// Setup volume state for EC task
volumeID := uint32(1)
adminServer.volumeStateManager.volumes[volumeID] = &VolumeState{
VolumeID: volumeID,
CurrentState: &VolumeInfo{
ID: volumeID,
Size: 25 * 1024 * 1024 * 1024, // 25GB
},
}
task := &types.Task{
ID: "ec_task_1",
Type: types.TaskTypeErasureCoding,
VolumeID: volumeID,
Server: "server1",
}
impact := adminServer.createTaskImpact(task, "worker1")
// Verify impact structure
if impact.TaskID != "ec_task_1" {
t.Errorf("Expected task ID ec_task_1, got %s", impact.TaskID)
}
if impact.TaskType != types.TaskTypeErasureCoding {
t.Errorf("Expected task type %v, got %v", types.TaskTypeErasureCoding, impact.TaskType)
}
// Verify volume changes for EC task
if !impact.VolumeChanges.WillBecomeReadOnly {
t.Error("Expected volume to become read-only after EC")
}
// Verify capacity delta (EC should require ~40% more space)
expectedCapacity := int64(float64(25*1024*1024*1024) * 1.4) // ~35GB
actualCapacity := impact.CapacityDelta["server1"]
if actualCapacity != expectedCapacity {
t.Errorf("Expected capacity delta %d, got %d", expectedCapacity, actualCapacity)
}
// Verify shard changes (should plan 14 shards)
if len(impact.ShardChanges) != 14 {
t.Errorf("Expected 14 shard changes, got %d", len(impact.ShardChanges))
}
for i := 0; i < 14; i++ {
shardChange := impact.ShardChanges[i]
if shardChange == nil {
t.Errorf("Missing shard change for shard %d", i)
continue
}
if !shardChange.WillBeCreated {
t.Errorf("Shard %d should be marked for creation", i)
}
}
t.Log("✅ Task impact creation test passed")
}
func TestAdminServer_TaskCompletionStateCleanup(t *testing.T) {
adminServer := NewAdminServer(DefaultAdminConfig(), nil)
adminServer.workerRegistry = NewWorkerRegistry()
adminServer.volumeStateManager = NewVolumeStateManager(nil)
adminServer.inProgressTasks = make(map[string]*InProgressTask)
// Setup worker
worker := &types.Worker{
ID: "worker1",
CurrentLoad: 1, // Has 1 task assigned
}
adminServer.workerRegistry.RegisterWorker(worker)
// Setup in-progress task
task := &types.Task{
ID: "test_task_1",
Type: types.TaskTypeVacuum,
VolumeID: 1,
}
inProgressTask := &InProgressTask{
Task: task,
WorkerID: "worker1",
VolumeReserved: true,
}
adminServer.inProgressTasks["test_task_1"] = inProgressTask
// Register impact in state manager
impact := &TaskImpact{
TaskID: "test_task_1",
VolumeID: 1,
CapacityDelta: map[string]int64{"server1": -100 * 1024 * 1024}, // 100MB savings
}
adminServer.volumeStateManager.RegisterTaskImpact("test_task_1", impact)
// Complete the task
err := adminServer.CompleteTask("test_task_1", true, "")
if err != nil {
t.Errorf("Task completion failed: %v", err)
}
// Verify cleanup
if len(adminServer.inProgressTasks) != 0 {
t.Errorf("Expected 0 in-progress tasks after completion, got %d", len(adminServer.inProgressTasks))
}
// Verify worker load updated
updatedWorker, _ := adminServer.workerRegistry.GetWorker("worker1")
if updatedWorker.CurrentLoad != 0 {
t.Errorf("Expected worker load 0 after task completion, got %d", updatedWorker.CurrentLoad)
}
// Verify state manager cleaned up
if len(adminServer.volumeStateManager.inProgressTasks) != 0 {
t.Errorf("Expected 0 tasks in state manager after completion, got %d", len(adminServer.volumeStateManager.inProgressTasks))
}
t.Log("✅ Task completion state cleanup test passed")
}
func TestAdminServer_PreventDuplicateTaskAssignment(t *testing.T) {
adminServer := NewAdminServer(DefaultAdminConfig(), nil)
adminServer.workerRegistry = NewWorkerRegistry()
adminServer.taskQueue = NewPriorityTaskQueue()
adminServer.volumeStateManager = NewVolumeStateManager(nil)
adminServer.inProgressTasks = make(map[string]*InProgressTask)
// Setup worker
worker := &types.Worker{
ID: "worker1",
Capabilities: []types.TaskType{types.TaskTypeVacuum},
MaxConcurrent: 2,
Status: "active",
CurrentLoad: 0,
}
adminServer.workerRegistry.RegisterWorker(worker)
// Setup volume state
volumeID := uint32(1)
adminServer.volumeStateManager.volumes[volumeID] = &VolumeState{
VolumeID: volumeID,
CurrentState: &VolumeInfo{ID: volumeID, Size: 1024 * 1024 * 1024},
}
// Create first task and assign it
task1 := &types.Task{
ID: "vacuum_task_1",
Type: types.TaskTypeVacuum,
VolumeID: volumeID,
Priority: types.TaskPriorityNormal,
}
adminServer.taskQueue.Push(task1)
assignedTask1, err := adminServer.RequestTask("worker1", []types.TaskType{types.TaskTypeVacuum})
if err != nil || assignedTask1 == nil {
t.Fatal("First task assignment failed")
}
// Try to assign another vacuum task for the same volume
task2 := &types.Task{
ID: "vacuum_task_2",
Type: types.TaskTypeVacuum,
VolumeID: volumeID, // Same volume!
Priority: types.TaskPriorityNormal,
}
adminServer.taskQueue.Push(task2)
assignedTask2, err := adminServer.RequestTask("worker1", []types.TaskType{types.TaskTypeVacuum})
// Should not assign duplicate task
if assignedTask2 != nil {
t.Error("Should not assign duplicate vacuum task for same volume")
}
t.Log("✅ Duplicate task prevention test passed")
}
func TestAdminServer_SystemStats(t *testing.T) {
adminServer := NewAdminServer(DefaultAdminConfig(), nil)
adminServer.workerRegistry = NewWorkerRegistry()
adminServer.taskQueue = NewPriorityTaskQueue()
adminServer.volumeStateManager = NewVolumeStateManager(nil)
adminServer.inProgressTasks = make(map[string]*InProgressTask)
adminServer.running = true
// Add some test data
worker := &types.Worker{ID: "worker1", Status: "active"}
adminServer.workerRegistry.RegisterWorker(worker)
task := &types.Task{ID: "task1", Type: types.TaskTypeErasureCoding}
adminServer.taskQueue.Push(task)
inProgressTask := &InProgressTask{
Task: &types.Task{ID: "task2", Type: types.TaskTypeVacuum},
}
adminServer.inProgressTasks["task2"] = inProgressTask
// Get system stats
stats := adminServer.GetSystemStats()
// Verify stats structure
if !stats["running"].(bool) {
t.Error("Expected running to be true")
}
if stats["in_progress_tasks"].(int) != 1 {
t.Errorf("Expected 1 in-progress task, got %d", stats["in_progress_tasks"].(int))
}
if stats["queued_tasks"].(int) != 1 {
t.Errorf("Expected 1 queued task, got %d", stats["queued_tasks"].(int))
}
// Check task breakdown
tasksByType := stats["tasks_by_type"].(map[types.TaskType]int)
if tasksByType[types.TaskTypeVacuum] != 1 {
t.Errorf("Expected 1 vacuum task, got %d", tasksByType[types.TaskTypeVacuum])
}
t.Log("✅ System stats test passed")
}
func TestAdminServer_VolumeStateIntegration(t *testing.T) {
// Integration test: Verify admin server correctly uses volume state for decisions
adminServer := NewAdminServer(DefaultAdminConfig(), nil)
adminServer.workerRegistry = NewWorkerRegistry()
adminServer.taskQueue = NewPriorityTaskQueue()
adminServer.volumeStateManager = NewVolumeStateManager(nil)
adminServer.inProgressTasks = make(map[string]*InProgressTask)
// Setup worker
worker := &types.Worker{
ID: "worker1",
Address: "server1",
Capabilities: []types.TaskType{types.TaskTypeErasureCoding},
MaxConcurrent: 1,
Status: "active",
CurrentLoad: 0,
}
adminServer.workerRegistry.RegisterWorker(worker)
// Setup volume and capacity that would normally allow EC
volumeID := uint32(1)
adminServer.volumeStateManager.volumes[volumeID] = &VolumeState{
VolumeID: volumeID,
CurrentState: &VolumeInfo{
ID: volumeID,
Size: 25 * 1024 * 1024 * 1024, // 25GB
Server: "server1",
},
}
adminServer.volumeStateManager.capacityCache["server1"] = &CapacityInfo{
Server: "server1",
TotalCapacity: 100 * 1024 * 1024 * 1024, // 100GB
UsedCapacity: 20 * 1024 * 1024 * 1024, // 20GB used
PredictedUsage: 20 * 1024 * 1024 * 1024, // 80GB available
}
// Create EC task
task := &types.Task{
ID: "ec_task_1",
Type: types.TaskTypeErasureCoding,
VolumeID: volumeID,
Server: "server1",
}
adminServer.taskQueue.Push(task)
// First assignment should work
assignedTask1, err := adminServer.RequestTask("worker1", []types.TaskType{types.TaskTypeErasureCoding})
if err != nil || assignedTask1 == nil {
t.Fatal("First EC task assignment should succeed")
}
// Verify capacity is now reserved
capacity := adminServer.volumeStateManager.GetAccurateCapacity("server1")
if capacity.ReservedCapacity <= 0 {
t.Error("Expected capacity to be reserved for first EC task")
}
// Try to assign another large EC task - should fail due to capacity
volumeID2 := uint32(2)
adminServer.volumeStateManager.volumes[volumeID2] = &VolumeState{
VolumeID: volumeID2,
CurrentState: &VolumeInfo{
ID: volumeID2,
Size: 30 * 1024 * 1024 * 1024, // 30GB - would need 42GB for EC
Server: "server1",
},
}
task2 := &types.Task{
ID: "ec_task_2",
Type: types.TaskTypeErasureCoding,
VolumeID: volumeID2,
Server: "server1",
}
adminServer.taskQueue.Push(task2)
// Add another worker to test capacity-based rejection
worker2 := &types.Worker{
ID: "worker2",
Address: "server1",
Capabilities: []types.TaskType{types.TaskTypeErasureCoding},
MaxConcurrent: 1,
Status: "active",
CurrentLoad: 0,
}
adminServer.workerRegistry.RegisterWorker(worker2)
assignedTask2, err := adminServer.RequestTask("worker2", []types.TaskType{types.TaskTypeErasureCoding})
// Should not assign due to insufficient capacity
if assignedTask2 != nil {
t.Error("Should not assign second EC task due to insufficient server capacity")
}
t.Log("✅ Volume state integration test passed")
t.Log("✅ Admin server correctly uses comprehensive state for task assignment decisions")
}
// Benchmark for task assignment performance
func BenchmarkAdminServer_RequestTask(b *testing.B) {
adminServer := NewAdminServer(DefaultAdminConfig(), nil)
adminServer.workerRegistry = NewWorkerRegistry()
adminServer.taskQueue = NewPriorityTaskQueue()
adminServer.volumeStateManager = NewVolumeStateManager(nil)
adminServer.inProgressTasks = make(map[string]*InProgressTask)
// Setup worker
worker := &types.Worker{
ID: "bench_worker",
Capabilities: []types.TaskType{types.TaskTypeVacuum},
MaxConcurrent: 1000, // High limit for benchmark
Status: "active",
CurrentLoad: 0,
}
adminServer.workerRegistry.RegisterWorker(worker)
// Setup many tasks
for i := 0; i < 1000; i++ {
volumeID := uint32(i + 1)
adminServer.volumeStateManager.volumes[volumeID] = &VolumeState{
VolumeID: volumeID,
CurrentState: &VolumeInfo{ID: volumeID, Size: 1024 * 1024 * 1024},
}
task := &types.Task{
ID: fmt.Sprintf("task_%d", i),
Type: types.TaskTypeVacuum,
VolumeID: volumeID,
}
adminServer.taskQueue.Push(task)
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
adminServer.RequestTask("bench_worker", []types.TaskType{types.TaskTypeVacuum})
}
}