From e3f9b347bb79e415883664d3e98cb500baa285f1 Mon Sep 17 00:00:00 2001
From: Raphael Simon <simon.raphael@gmail.com>
Date: Thu, 10 Oct 2024 16:28:00 -0700
Subject: [PATCH] Stop workers before requeuing their jobs (#22)

When the node is closed.
---
 pool/node.go | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/pool/node.go b/pool/node.go
index 71acdd5..f72a642 100644
--- a/pool/node.go
+++ b/pool/node.go
@@ -357,9 +357,19 @@ func (node *Node) Close(ctx context.Context) error {
 	}
 	node.logger.Info("closing")
 	node.closing = true
+
+	// Need to stop workers before requeueing jobs to prevent
+	// requeued jobs from being handled by this node.
+	var wg sync.WaitGroup
 	for _, w := range node.localWorkers {
-		go w.stopAndWait(ctx)
+		wg.Add(1)
+		go func(w *Worker) {
+			defer wg.Done()
+			w.stopAndWait(ctx)
+		}(w)
 	}
+	wg.Wait()
+
 	for _, w := range node.localWorkers {
 		w.requeueJobs(ctx)
 	}