diff --git a/weed/worker/worker.go b/weed/worker/worker.go index e361b30e0..536cd9b86 100644 --- a/weed/worker/worker.go +++ b/weed/worker/worker.go @@ -26,6 +26,7 @@ type Worker struct { config *types.WorkerConfig registry *tasks.TaskRegistry cmds chan workerCommand + done chan struct{} // closed when managerLoop exits, so senders stop blocking state *workerState taskLogHandler *tasks.TaskLogHandler } @@ -193,6 +194,7 @@ func NewWorker(config *types.WorkerConfig) (*Worker, error) { registry: registry, taskLogHandler: taskLogHandler, cmds: make(chan workerCommand), + done: make(chan struct{}), } glog.V(1).Infof("Worker created with %d registered task types", len(registry.GetAll())) @@ -292,24 +294,35 @@ out: } } + // The loop no longer receives. Wake any task goroutine still trying to + // report into w.cmds so it returns instead of blocking forever. + close(w.done) +} + +// dispatch delivers a fire-and-forget command to the manager loop, aborting +// if the worker has stopped. Returns false when the loop is gone, so task +// goroutines that outlive a forced shutdown don't block on it forever. +func (w *Worker) dispatch(cmd workerCommand) bool { + select { + case w.cmds <- cmd: + return true + case <-w.done: + return false + } } func (w *Worker) getTaskLoad() int { respCh := make(chan int, 1) - w.cmds <- workerCommand{ - action: ActionGetTaskLoad, - data: respCh, - resp: nil, + if !w.dispatch(workerCommand{action: ActionGetTaskLoad, data: respCh}) { + return 0 } return <-respCh } func (w *Worker) setTask(task *types.TaskInput) error { resp := make(chan error) - w.cmds <- workerCommand{ - action: ActionSetTask, - data: task, - resp: resp, + if !w.dispatch(workerCommand{action: ActionSetTask, data: task, resp: resp}) { + return fmt.Errorf("worker is stopping") } if err := <-resp; err != nil { glog.Errorf("TASK REJECTED: Worker %s at capacity (%d/%d) - rejecting task %s", @@ -324,18 +337,16 @@ func (w *Worker) setTask(task *types.TaskInput) error { } func (w *Worker) removeTask(task *types.TaskInput) int { - w.cmds <- workerCommand{ - action: ActionRemoveTask, - data: task.ID, + if !w.dispatch(workerCommand{action: ActionRemoveTask, data: task.ID}) { + return 0 } return w.getTaskLoad() } func (w *Worker) getAdmin() AdminClient { respCh := make(chan AdminClient, 1) - w.cmds <- workerCommand{ - action: ActionGetAdmin, - data: respCh, + if !w.dispatch(workerCommand{action: ActionGetAdmin, data: respCh}) { + return nil } return <-respCh } @@ -466,8 +477,8 @@ func (w *Worker) handleStart(cmd workerCommand) { func (w *Worker) Stop() error { // Snapshot the admin client BEFORE ActionStop terminates the manager loop. - // Once the loop exits, any further w.cmds send (getTaskLoad, getAdmin, etc.) - // would deadlock because no one reads w.cmds anymore. + // Once the loop exits, getAdmin/getTaskLoad return zero values (they abort + // on w.done) rather than the real client, so grab it while it's still live. adminClient := w.getAdmin() // Best-effort task drain: wait up to 30s for in-flight tasks to finish. @@ -485,14 +496,16 @@ out: } } - // Terminate the manager loop. + // Terminate the manager loop. A second Stop (e.g. an admin-shutdown timer + // racing an explicit Stop) finds the loop already gone, so abort on w.done + // instead of blocking on a send no one will receive. resp := make(chan error) - w.cmds <- workerCommand{ - action: ActionStop, - resp: resp, - } - if err := <-resp; err != nil { - return err + select { + case w.cmds <- workerCommand{action: ActionStop, resp: resp}: + if err := <-resp; err != nil { + return err + } + case <-w.done: } // Disconnect from admin server. @@ -565,8 +578,10 @@ func (w *Worker) executeTask(task *types.TaskInput) { w.id, task.ID, task.Type, task.VolumeID, task.Server, task.Collection, startTime.Format(time.RFC3339)) // Report task start to admin server - if err := w.getAdmin().UpdateTaskProgress(task.ID, 0.0); err != nil { - glog.V(1).Infof("Failed to report task start to admin: %v", err) + if admin := w.getAdmin(); admin != nil { + if err := admin.UpdateTaskProgress(task.ID, 0.0); err != nil { + glog.V(1).Infof("Failed to report task start to admin: %v", err) + } } // Determine task-specific working directory (BaseWorkingDir is guaranteed to be non-empty) @@ -639,8 +654,10 @@ func (w *Worker) executeTask(task *types.TaskInput) { taskInstance.SetProgressCallback(func(progress float64, stage string) { // Report progress updates to admin server glog.V(2).Infof("Task %s progress: %.1f%% - %s", task.ID, progress, stage) - if err := w.getAdmin().UpdateTaskProgress(task.ID, progress); err != nil { - glog.V(1).Infof("Failed to report task progress to admin: %v", err) + if admin := w.getAdmin(); admin != nil { + if err := admin.UpdateTaskProgress(task.ID, progress); err != nil { + glog.V(1).Infof("Failed to report task progress to admin: %v", err) + } } if fileLogger != nil { // Use meaningful stage description or fallback to generic message @@ -662,9 +679,7 @@ func (w *Worker) executeTask(task *types.TaskInput) { } if err != nil { w.completeTask(task.ID, false, err.Error()) - w.cmds <- workerCommand{ - action: ActionIncTaskFail, - } + w.dispatch(workerCommand{action: ActionIncTaskFail}) glog.Errorf("Worker %s failed to execute task %s: %v", w.id, task.ID, err) if fileLogger != nil { fileLogger.LogStatus("failed", err.Error()) @@ -676,19 +691,19 @@ func (w *Worker) executeTask(task *types.TaskInput) { fileLogger.Sync() } w.completeTask(task.ID, true, "") - w.cmds <- workerCommand{ - action: ActionIncTaskComplete, - } + w.dispatch(workerCommand{action: ActionIncTaskComplete}) glog.Infof("Worker %s completed task %s successfully", w.id, task.ID) } } // completeTask reports task completion to admin server func (w *Worker) completeTask(taskID string, success bool, errorMsg string) { - if w.getAdmin() != nil { - if err := w.getAdmin().CompleteTask(taskID, success, errorMsg); err != nil { - glog.Errorf("Failed to report task completion: %v", err) - } + admin := w.getAdmin() + if admin == nil { + return + } + if err := admin.CompleteTask(taskID, success, errorMsg); err != nil { + glog.Errorf("Failed to report task completion: %v", err) } } @@ -724,17 +739,19 @@ func (w *Worker) taskRequestLoop() { // sendHeartbeat sends heartbeat to admin server func (w *Worker) sendHeartbeat() { - if w.getAdmin() != nil { - if err := w.getAdmin().SendHeartbeat(w.id, &types.WorkerStatus{ - WorkerID: w.id, - Status: "active", - Capabilities: w.config.Capabilities, - MaxConcurrent: w.config.MaxConcurrent, - CurrentLoad: w.getTaskLoad(), - LastHeartbeat: time.Now(), - }); err != nil { - glog.Warningf("Failed to send heartbeat: %v", err) - } + admin := w.getAdmin() + if admin == nil { + return + } + if err := admin.SendHeartbeat(w.id, &types.WorkerStatus{ + WorkerID: w.id, + Status: "active", + Capabilities: w.config.Capabilities, + MaxConcurrent: w.config.MaxConcurrent, + CurrentLoad: w.getTaskLoad(), + LastHeartbeat: time.Now(), + }); err != nil { + glog.Warningf("Failed to send heartbeat: %v", err) } } @@ -748,11 +765,12 @@ func (w *Worker) requestTasks() { return // Already at capacity } - if w.getAdmin() != nil { + admin := w.getAdmin() + if admin != nil { glog.V(4).Infof("REQUESTING TASK: Worker %s requesting task from admin server (current load: %d/%d, capabilities: %v)", w.id, currentLoad, w.config.MaxConcurrent, w.config.Capabilities) - task, err := w.getAdmin().RequestTask(w.id, w.config.Capabilities) + task, err := admin.RequestTask(w.id, w.config.Capabilities) if err != nil { glog.V(2).Infof("TASK REQUEST FAILED: Worker %s failed to request task: %v", w.id, err) return @@ -784,7 +802,8 @@ func (w *Worker) connectionMonitorLoop() { return case <-ticker.C: // Monitor connection status and log changes - currentConnectionStatus := w.getAdmin() != nil && w.getAdmin().IsConnected() + admin := w.getAdmin() + currentConnectionStatus := admin != nil && admin.IsConnected() if currentConnectionStatus != lastConnectionStatus { if currentConnectionStatus { @@ -918,11 +937,7 @@ func (w *Worker) handleTaskLogRequest(request *worker_pb.TaskLogRequest) { func (w *Worker) handleTaskCancellation(cancellation *worker_pb.TaskCancellation) { glog.Infof("Worker %s received task cancellation for task %s", w.id, cancellation.TaskId) - w.cmds <- workerCommand{ - action: ActionCancelTask, - data: cancellation.TaskId, - resp: nil, - } + w.dispatch(workerCommand{action: ActionCancelTask, data: cancellation.TaskId}) } // handleAdminShutdown processes admin shutdown notifications diff --git a/weed/worker/worker_shutdown_test.go b/weed/worker/worker_shutdown_test.go new file mode 100644 index 000000000..8c7597578 --- /dev/null +++ b/weed/worker/worker_shutdown_test.go @@ -0,0 +1,76 @@ +package worker + +import ( + "testing" + "time" + + "github.com/seaweedfs/seaweedfs/weed/pb/worker_pb" + "github.com/seaweedfs/seaweedfs/weed/worker/types" +) + +// A task goroutine can outlive a forced shutdown: Stop() only drains for 30s, +// then terminates the manager loop while the task is still running. When the +// task finally reports completion it goes through w.cmds (getAdmin in +// completeTask, the ActionIncTask* send, removeTask). With the loop gone and an +// unbuffered channel, those sends - and the response reads behind getAdmin / +// getTaskLoad - block forever, leaking the goroutine. They must abort instead. +func TestManagerLoopSendersAbortAfterStop(t *testing.T) { + w, err := NewWorker(&types.WorkerConfig{ + BaseWorkingDir: t.TempDir(), + MaxConcurrent: 2, + }) + if err != nil { + t.Fatalf("NewWorker: %v", err) + } + + // No Start(): the drain loop sees zero load and ActionStop terminates the + // manager loop immediately, modelling the post-shutdown state a lingering + // task goroutine would race into. + if err := w.Stop(); err != nil { + t.Fatalf("Stop: %v", err) + } + + // Every interaction with the manager loop - whether from a lingering + // executeTask goroutine, the task-assignment path, or a second Stop - must + // return promptly now that the loop no longer receives. + mustNotBlock(t, "getAdmin", func() { w.getAdmin() }) + mustNotBlock(t, "getTaskLoad", func() { w.getTaskLoad() }) + mustNotBlock(t, "removeTask", func() { w.removeTask(&types.TaskInput{ID: "t1"}) }) + mustNotBlock(t, "completeTask", func() { w.completeTask("t1", true, "") }) + mustNotBlock(t, "dispatch", func() { w.dispatch(workerCommand{action: ActionIncTaskComplete}) }) + mustNotBlock(t, "setTask", func() { w.setTask(&types.TaskInput{ID: "t1"}) }) + mustNotBlock(t, "handleTaskCancellation", func() { w.handleTaskCancellation(&worker_pb.TaskCancellation{TaskId: "t1"}) }) + mustNotBlock(t, "Stop again", func() { w.Stop() }) + + // The abort path returns zero values rather than blocking on real state. + if got := w.getAdmin(); got != nil { + t.Errorf("getAdmin after stop = %v, want nil", got) + } + if got := w.getTaskLoad(); got != 0 { + t.Errorf("getTaskLoad after stop = %d, want 0", got) + } + if got := w.removeTask(&types.TaskInput{ID: "t1"}); got != 0 { + t.Errorf("removeTask after stop = %d, want 0", got) + } + if err := w.setTask(&types.TaskInput{ID: "t1"}); err == nil { + t.Errorf("setTask after stop = nil error, want error") + } + if w.dispatch(workerCommand{action: ActionIncTaskComplete}) { + t.Errorf("dispatch after stop = true, want false") + } +} + +// mustNotBlock fails the test if fn does not return within the timeout. +func mustNotBlock(t *testing.T, name string, fn func()) { + t.Helper() + done := make(chan struct{}) + go func() { + fn() + close(done) + }() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatalf("%s blocked after manager loop stopped (goroutine leak)", name) + } +}