tasks/service/history/queues/memory_scheduled_queue.go

package queues

import (
	"sync"
	"sync/atomic"
	"time"

	"github.com/hanzoai/tasks/common"
	"github.com/hanzoai/tasks/common/clock"
	"github.com/hanzoai/tasks/common/collection"
	"github.com/hanzoai/tasks/common/log"
	"github.com/hanzoai/tasks/common/log/tag"
	"github.com/hanzoai/tasks/common/metrics"
	ctasks "github.com/hanzoai/tasks/common/tasks"
)

const (
	workerBusyRescheduleDelay = 1 * time.Second
)

type (
	memoryScheduledQueue struct {
		scheduler ctasks.Scheduler[ctasks.Task]

		taskQueue     collection.Queue[Executable]
		nextTaskTimer *time.Timer
		newTaskCh     chan Executable

		timeSource     clock.TimeSource
		logger         log.Logger
		metricsHandler metrics.Handler

		status     int32
		shutdownCh chan struct{}
		shutdownWG sync.WaitGroup
	}
)

func newMemoryScheduledQueue(
	scheduler ctasks.Scheduler[ctasks.Task],
	timeSource clock.TimeSource,
	logger log.Logger,
	metricsHandler metrics.Handler,
) *memoryScheduledQueue {

	nextTaskTimer := time.NewTimer(0)
	if !nextTaskTimer.Stop() {
		<-nextTaskTimer.C
	}

	return &memoryScheduledQueue{
		taskQueue:     collection.NewPriorityQueue[Executable](executableVisibilityTimeCompareLess),
		nextTaskTimer: nextTaskTimer,
		newTaskCh:     make(chan Executable),

		timeSource:     timeSource,
		logger:         logger,
		metricsHandler: metricsHandler,

		status:     common.DaemonStatusInitialized,
		shutdownCh: make(chan struct{}),

		scheduler: scheduler,
	}
}

func executableVisibilityTimeCompareLess(
	this Executable,
	that Executable,
) bool {
	return this.GetVisibilityTime().Before(that.GetVisibilityTime())
}

func (q *memoryScheduledQueue) Start() {
	if !atomic.CompareAndSwapInt32(&q.status, common.DaemonStatusInitialized, common.DaemonStatusStarted) {
		return
	}

	q.logger.Info("", tag.LifeCycleStarting)
	defer q.logger.Info("", tag.LifeCycleStarted)

	q.shutdownWG.Add(1)
	go q.processQueueLoop()
}

func (q *memoryScheduledQueue) Stop() {
	if !atomic.CompareAndSwapInt32(&q.status, common.DaemonStatusStarted, common.DaemonStatusStopped) {
		return
	}

	q.logger.Info("", tag.LifeCycleStopping)
	defer q.logger.Info("", tag.LifeCycleStopped)

	close(q.shutdownCh)

	if success := common.AwaitWaitGroup(&q.shutdownWG, time.Minute); !success {
		q.logger.Warn("", tag.LifeCycleStopTimedout)
	}
}

func (q *memoryScheduledQueue) Add(task Executable) {
	q.newTaskCh <- task
}

//nolint:revive // cognitive complexity
func (q *memoryScheduledQueue) processQueueLoop() {
	defer q.shutdownWG.Done()

	for {
		select {
		case <-q.shutdownCh:
			return
		default:
		}

		select {
		case <-q.shutdownCh:
			return
		case newTask := <-q.newTaskCh:
			var nextTaskTime time.Time
			if !q.taskQueue.IsEmpty() {
				nextTaskTime = q.taskQueue.Peek().GetVisibilityTime()
			}
			q.taskQueue.Add(newTask)
			// If there is no timer set OR new time is earlier than the current one, then timer needs to be set to the new time.
			if nextTaskTime.IsZero() || newTask.GetVisibilityTime().Before(nextTaskTime) {
				// But before reset if timer is there it needs to be stopped.
				if !nextTaskTime.IsZero() {
					if !q.nextTaskTimer.Stop() {
						// Drain timer channel to prevent timer firing.
						<-q.nextTaskTimer.C
					}
				}
				q.nextTaskTimer.Reset(newTask.GetVisibilityTime().Sub(q.timeSource.Now()))
			}
			metrics.NewTimerNotifyCounter.With(q.metricsHandler).Record(1)
		case <-q.nextTaskTimer.C:
			taskToExecute := q.taskQueue.Remove()
			// Skip tasks which are already canceled. Majority of the tasks in the queue should be cancelled already.
			nextTask := q.purgeCanceledTasks()
			if nextTask != nil {
				q.nextTaskTimer.Reset(nextTask.GetVisibilityTime().Sub(q.timeSource.Now()))
			}
			// If current taskToExecute is also cancelled don't submit it to scheduler.
			if taskToExecute.State() == ctasks.TaskStateCancelled {
				continue
			}

			// For scheduler metrics to work properly.
			taskToExecute.SetScheduledTime(q.timeSource.Now())

			if !q.scheduler.TrySubmit(taskToExecute) {
				// If all workers are busy then put the task back to the queue.
				// Must be done in a separate goroutine to avoid deadlock.
				go func() {
					taskToExecute.SetVisibilityTime(taskToExecute.GetVisibilityTime().Add(workerBusyRescheduleDelay))
					q.Add(taskToExecute)
				}()
			}
		}
	}
}
func (q *memoryScheduledQueue) purgeCanceledTasks() (nextTask Executable) {
	if !q.taskQueue.IsEmpty() {
		nextTask = q.taskQueue.Peek()
	}
	for nextTask != nil {
		if nextTask.State() != ctasks.TaskStateCancelled {
			return
		}
		// Remove canceled task from queue.
		q.taskQueue.Remove()
		if q.taskQueue.IsEmpty() {
			return nil
		}
		nextTask = q.taskQueue.Peek()
	}
	return
}