From 8059e66a0dae874b5964deaf1cea3daf1524a15f Mon Sep 17 00:00:00 2001 From: Jeroen Mulder Date: Mon, 10 Aug 2026 11:28:08 +0200 Subject: [PATCH 1/2] Update handle_message.go Actually shut down workers --- worker/handle_message.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/worker/handle_message.go b/worker/handle_message.go index 4a84226..79f586a 100644 --- a/worker/handle_message.go +++ b/worker/handle_message.go @@ -100,7 +100,14 @@ func (w *Worker) handleSystem(msg *contract.SystemMessage) { switch msg.Command { case "drain": w.state.Store(contract.WorkerStateDraining) - // TODO: maybe it would help if we actually drained here.... + w.tasksMu.Lock() + remaining := len(w.tasks) + + if remaining == 0 { + w.state.Store(contract.WorkerStateShuttingDown) + w.cfg.Logger.Info("[Worker][handleSystem] Drain completed (no active tasks. Closing connection...") + w.conn.Close() + } case "shutdown": if w.state.Transition(contract.WorkerStateOnline, contract.WorkerStateShuttingDown) || w.state.Transition(contract.WorkerStateDraining, contract.WorkerStateShuttingDown) { close(w.done) From 84a2ad8cb9a1f4de6f5b9744986b2cc1b44a78e5 Mon Sep 17 00:00:00 2001 From: Jeroen Mulder Date: Mon, 10 Aug 2026 11:35:18 +0200 Subject: [PATCH 2/2] Update handle_message.go Let's unlock it too --- worker/handle_message.go | 1 + 1 file changed, 1 insertion(+) diff --git a/worker/handle_message.go b/worker/handle_message.go index 79f586a..b0f59d8 100644 --- a/worker/handle_message.go +++ b/worker/handle_message.go @@ -102,6 +102,7 @@ func (w *Worker) handleSystem(msg *contract.SystemMessage) { w.state.Store(contract.WorkerStateDraining) w.tasksMu.Lock() remaining := len(w.tasks) + w.tasksMu.Unlock() if remaining == 0 { w.state.Store(contract.WorkerStateShuttingDown)