worker: don't leak task goroutines on forced shutdown (#10062)

* worker: don't leak task goroutines on forced shutdown Stop() drains in-flight tasks for 30s, then terminates the manager loop. A task still running past that deadline later reports completion through w.cmds - getAdmin in completeTask, the ActionIncTask* send, removeTask - but with the loop gone and cmds unbuffered, those sends and the response reads behind getAdmin/getTaskLoad block forever, leaking the goroutine. Close a done channel when the loop exits and route the task-goroutine sends through it so they abort and return zero values instead of blocking. getAdmin can now return nil mid-shutdown, so collapse its double-call sites to a single nil-checked call to avoid a deref. * worker: abort remaining manager-loop sends after shutdown Extend the post-shutdown abort to the sends that still blocked: Stop()'s own ActionStop (so a second Stop, e.g. an admin-shutdown timer racing an explicit one, doesn't hang), setTask, and handleTaskCancellation. Route them through w.done so they return instead of blocking when the loop is gone. Stop is now idempotent.
2026-07-19 06:22:31 +00:00 · 2026-06-23 10:06:59 -07:00
parent faa8c3963b
commit b0bad761ff
2 changed files with 147 additions and 56 deletions
@@ -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
@@ -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)
+	}
+}