From 6029bae96b08e11f55767afe9383f5fe6ecf6608 Mon Sep 17 00:00:00 2001 From: H0llyW00dzZ Date: Sun, 31 Dec 2023 06:04:36 +0700 Subject: [PATCH] Refactor [Worker] [Captain] Fix Potential Issue of Deadlocks (#56) - [+] chore(captain.go): refactor shutdown function to use sync.Once for ensuring it is only called once --- worker/captain.go | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/worker/captain.go b/worker/captain.go index 2fafadd..e61694d 100644 --- a/worker/captain.go +++ b/worker/captain.go @@ -27,6 +27,7 @@ import ( func CaptainTellWorkers(ctx context.Context, clientset *kubernetes.Clientset, shipsNamespace string, tasks []configuration.Task, workerCount int) (<-chan string, func()) { results := make(chan string) var wg sync.WaitGroup + var once sync.Once // Use sync.Once to ensure shutdown is only called once taskStatus := NewTaskStatusMap() // Tracks the claiming of tasks to avoid duplication. shutdownCtx, cancelFunc := context.WithCancel(ctx) // Derived context to signal shutdown. @@ -42,13 +43,15 @@ func CaptainTellWorkers(ctx context.Context, clientset *kubernetes.Clientset, sh // shutdown is called to initiate a graceful shutdown of all workers. shutdown := func() { - cancelFunc() // Signal workers to stop by cancelling the context. - - // Ensure channel closure happens after all workers have finished. - go func() { - wg.Wait() // Wait for all workers to complete. - close(results) // Close the results channel safely. - }() + once.Do(func() { // Ensure this block only runs once + cancelFunc() // Signal workers to stop by cancelling the context. + + // Ensure channel closure happens after all workers have finished. + go func() { + wg.Wait() // Wait for all workers to complete. + close(results) // Close the results channel safely. + }() + }) } return results, shutdown