[fix] Fix possibly deadlock (#818)

Signed-off-by: Pavel Okhlopkov <[email protected]>
Co-authored-by: Copilot <[email protected]>
Signed-off-by: Pavel Okhlopkov <[email protected]>

Author: ldmonster

pkg/shell-operator/combine_binding_context.go

@@ -38,7 +38,7 @@ func (op *ShellOperator) combineBindingContextForHook(tqs *queue.TaskQueueSet, q
 
 	otherTasks := make([]task.Task, 0)
 	stopIterate := false
-	q.Iterate(func(tsk task.Task) {
+	q.IterateSnapshot(func(tsk task.Task) {
 		if stopIterate {
 			return
 		}
@@ -91,7 +91,7 @@ func (op *ShellOperator) combineBindingContextForHook(tqs *queue.TaskQueueSet, q
 	}
 
 	// Delete tasks with false in tasksFilter map
-	tqs.GetByName(t.GetQueueName()).Filter(func(tsk task.Task) bool {
+	tqs.GetByName(t.GetQueueName()).DeleteFunc(func(tsk task.Task) bool {
 		if v, ok := tasksFilter[tsk.GetId()]; ok {
 			return v
 		}

pkg/shell-operator/manager_events_handler.go

@@ -88,7 +88,7 @@ func (m *ManagerEventsHandler) Start() {
 
 			m.taskQueues.DoWithLock(func(tqs *queue.TaskQueueSet) {
 				for _, resTask := range tailTasks {
-					if q := tqs.Queues[resTask.GetQueueName()]; q == nil {
+					if q, ok := tqs.Queues.Get(resTask.GetQueueName()); !ok {
 						log.Error("Possible bug!!! Got task for queue but queue is not created yet.",
 							slog.String("queueName", resTask.GetQueueName()),
 							slog.String("description", resTask.GetDescription()))

pkg/shell-operator/operator.go

@@ -731,7 +731,7 @@ func (op *ShellOperator) CombineBindingContextForHook(q *queue.TaskQueue, t task
 
 	otherTasks := make([]task.Task, 0)
 	stopIterate := false
-	q.Iterate(func(tsk task.Task) {
+	q.IterateSnapshot(func(tsk task.Task) {
 		if stopIterate {
 			return
 		}
@@ -782,7 +782,7 @@ func (op *ShellOperator) CombineBindingContextForHook(q *queue.TaskQueue, t task
 	}
 
 	// Delete tasks with false in tasksFilter map
-	op.TaskQueues.GetByName(t.GetQueueName()).Filter(func(tsk task.Task) bool {
+	op.TaskQueues.GetByName(t.GetQueueName()).DeleteFunc(func(tsk task.Task) bool {
 		if v, ok := tasksFilter[tsk.GetId()]; ok {
 			return v
 		}
@@ -949,7 +949,7 @@ func (op *ShellOperator) runMetrics() {
 	// task queue length
 	go func() {
 		for {
-			op.TaskQueues.Iterate(func(queue *queue.TaskQueue) {
+			op.TaskQueues.Iterate(context.TODO(), func(_ context.Context, queue *queue.TaskQueue) {
 				queueLen := float64(queue.Length())
 				op.MetricStorage.GaugeSet("{PREFIX}tasks_queue_length", queueLen, map[string]string{"queue": queue.Name})
 			})

pkg/shell-operator/operator_test.go

@@ -66,7 +66,7 @@ func Test_Operator_startup_tasks(t *testing.T) {
 	}
 
 	i := 0
-	op.TaskQueues.GetMain().Iterate(func(tsk task.Task) {
+	op.TaskQueues.GetMain().IterateSnapshot(func(tsk task.Task) {
 		// Stop checking if no expects left.
 		if i >= len(expectTasks) {
 			return

pkg/task/dump/dump.go

@@ -1,6 +1,7 @@
 package dump
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"sort"
@@ -78,7 +79,7 @@ func TaskQueues(tqs *queue.TaskQueueSet, format string, showEmpty bool) interfac
 	tasksCount := 0
 	mainTasksCount := 0
 
-	tqs.Iterate(func(queue *queue.TaskQueue) {
+	tqs.Iterate(context.TODO(), func(_ context.Context, queue *queue.TaskQueue) {
 		if queue == nil {
 			return
 		}
@@ -160,7 +161,7 @@ func getTasksForQueue(q *queue.TaskQueue) []dumpTask {
 	tasks := make([]dumpTask, 0, q.Length())
 
 	index := 1
-	q.Iterate(func(task task.Task) {
+	q.IterateSnapshot(func(task task.Task) {
 		tasks = append(tasks, dumpTask{
 			Index:       index,
 			Description: task.GetDescription(),

pkg/task/queue/queue_set.go

@@ -13,6 +13,58 @@ import (
 
 const MainQueueName = "main"
 
+// queueStorage is a thread-safe storage for task queues with basic Get/Set/Delete operations
+type queueStorage struct {
+	mu     sync.RWMutex
+	queues map[string]*TaskQueue
+}
+
+func newQueueStorage() *queueStorage {
+	return &queueStorage{
+		queues: make(map[string]*TaskQueue),
+	}
+}
+
+// Get retrieves a queue by name, returns nil if not found
+func (qs *queueStorage) Get(name string) (*TaskQueue, bool) {
+	qs.mu.RLock()
+	defer qs.mu.RUnlock()
+	queue, exists := qs.queues[name]
+	return queue, exists
+}
+
+// List retrieves all queues
+func (qs *queueStorage) List() []*TaskQueue {
+	qs.mu.RLock()
+	defer qs.mu.RUnlock()
+
+	queues := make([]*TaskQueue, 0, len(qs.queues))
+	for _, queue := range qs.queues {
+		queues = append(queues, queue)
+	}
+
+	return queues
+}
+
+// Set stores a queue with the given name
+func (qs *queueStorage) Set(name string, queue *TaskQueue) {
+	qs.mu.Lock()
+	defer qs.mu.Unlock()
+	qs.queues[name] = queue
+}
+
+// Delete removes a queue by name
+func (qs *queueStorage) Delete(name string) {
+	qs.mu.Lock()
+	defer qs.mu.Unlock()
+	delete(qs.queues, name)
+}
+
+// Len returns the number of tasks in a queue by name
+func (qs *queueStorage) Len() int {
+	return len(qs.queues)
+}
+
 // TaskQueueSet is a manager for a set of named queues
 type TaskQueueSet struct {
 	MainName string
@@ -23,12 +75,12 @@ type TaskQueueSet struct {
 	cancel context.CancelFunc
 
 	m      sync.RWMutex
-	Queues map[string]*TaskQueue
+	Queues *queueStorage
 }
 
 func NewTaskQueueSet() *TaskQueueSet {
 	return &TaskQueueSet{
-		Queues:   make(map[string]*TaskQueue),
+		Queues:   newQueueStorage(),
 		MainName: MainQueueName,
 	}
 }
@@ -61,18 +113,14 @@ func (tqs *TaskQueueSet) StartMain(ctx context.Context) {
 }
 
 func (tqs *TaskQueueSet) Start(ctx context.Context) {
-	tqs.m.RLock()
-	for _, q := range tqs.Queues {
-		q.Start(ctx)
-	}
-
-	tqs.m.RUnlock()
+	tqs.Iterate(ctx, func(ctx context.Context, queue *TaskQueue) {
+		queue.Start(ctx)
+	})
 }
 
+// Add register a new queue for TaskQueueSet.
 func (tqs *TaskQueueSet) Add(queue *TaskQueue) {
-	tqs.m.Lock()
-	tqs.Queues[queue.Name] = queue
-	tqs.m.Unlock()
+	tqs.Queues.Set(queue.Name, queue)
 }
 
 func (tqs *TaskQueueSet) NewNamedQueue(name string, handler func(ctx context.Context, t task.Task) TaskResult, opts ...TaskQueueOption) {
@@ -91,92 +139,88 @@ func (tqs *TaskQueueSet) NewNamedQueue(name string, handler func(ctx context.Con
 		q.logger = log.NewLogger().Named("task_queue")
 	}
 
-	tqs.m.Lock()
-	tqs.Queues[name] = q
-	tqs.m.Unlock()
+	tqs.Queues.Set(q.Name, q)
 }
 
 func (tqs *TaskQueueSet) GetByName(name string) *TaskQueue {
-	tqs.m.RLock()
-	defer tqs.m.RUnlock()
-	ts, exists := tqs.Queues[name]
-	if exists {
-		return ts
+	q, ok := tqs.Queues.Get(name)
+	if !ok {
+		return nil
 	}
-	return nil
+
+	return q
 }
 
 func (tqs *TaskQueueSet) GetMain() *TaskQueue {
 	return tqs.GetByName(tqs.MainName)
 }
 
-/*
-	taskQueueSet.DoWithLock(func(tqs *TaskQueueSet){
-	   tqs.GetMain().Pop()
-	})
-*/
 func (tqs *TaskQueueSet) DoWithLock(fn func(tqs *TaskQueueSet)) {
 	tqs.m.Lock()
 	defer tqs.m.Unlock()
+
 	if fn != nil {
 		fn(tqs)
 	}
 }
 
 // Iterate run doFn for every task.
-func (tqs *TaskQueueSet) Iterate(doFn func(queue *TaskQueue)) {
+func (tqs *TaskQueueSet) Iterate(ctx context.Context, doFn func(ctx context.Context, queue *TaskQueue)) {
 	if doFn == nil {
 		return
 	}
 
 	tqs.m.RLock()
 	defer tqs.m.RUnlock()
-	if len(tqs.Queues) == 0 {
+	if tqs.Queues.Len() == 0 {
 		return
 	}
 
 	main := tqs.GetMain()
 	if main != nil {
-		doFn(main)
+		doFn(ctx, main)
 	}
 	// TODO sort names
 
-	for _, q := range tqs.Queues {
+	for _, q := range tqs.Queues.List() {
 		if q.Name != tqs.MainName {
-			doFn(q)
+			doFn(ctx, q)
 		}
 	}
 }
 
 func (tqs *TaskQueueSet) Remove(name string) {
-	tqs.m.Lock()
-	ts, exists := tqs.Queues[name]
+	ts, exists := tqs.Queues.Get(name)
 	if exists {
 		ts.Stop()
 	}
 
-	delete(tqs.Queues, name)
-	tqs.m.Unlock()
+	tqs.Queues.Delete(name)
 }
 
 func (tqs *TaskQueueSet) WaitStopWithTimeout(timeout time.Duration) {
 	checkTick := time.NewTicker(time.Millisecond * 100)
 	defer checkTick.Stop()
+
 	timeoutTick := time.NewTicker(timeout)
 	defer timeoutTick.Stop()
 
 	for {
 		select {
 		case <-checkTick.C:
-			stopped := true
-			tqs.m.RLock()
-			for _, q := range tqs.Queues {
-				if q.Status != "stop" {
-					stopped = false
-					break
+			stopped := func() bool {
+				tqs.m.RLock()
+				defer tqs.m.RUnlock()
+
+				for _, q := range tqs.Queues.List() {
+					if q.GetStatusType() != QueueStatusStop {
+						return false
+					}
 				}
-			}
-			tqs.m.RUnlock()
+
+				return true
+			}()
+
 			if stopped {
 				return
 			}