Files

208 lines
4.2 KiB
Go

package queue
import (
"context"
"errors"
"sort"
"sync"
"time"
"git.casaderoll.de/michael/urbm/internal/model"
)
var ErrDuplicate = errors.New("task already queued or running")
type Handler func(context.Context, model.Run) model.Run
type OnChange func([]model.Run)
type Queue struct {
mu sync.Mutex
cond *sync.Cond
pending []model.Run
history []model.Run
active *model.Run
cancel context.CancelFunc
stopping bool
handler Handler
onChange OnChange
done chan struct{}
doneOnce sync.Once
}
func New(handler Handler, onChange OnChange) *Queue {
q := &Queue{handler: handler, onChange: onChange, done: make(chan struct{})}
q.cond = sync.NewCond(&q.mu)
return q
}
func (q *Queue) RestoreHistory(runs []model.Run) {
q.mu.Lock()
defer q.mu.Unlock()
for _, run := range runs {
if run.Status == "queued" || run.Status == "running" || run.Status == "paused" {
run.Status = "failed"
run.ErrorCode = "environment"
run.Message = "daemon restarted while task was active"
now := time.Now().UTC()
run.FinishedAt = &now
}
q.history = append(q.history, run)
}
}
func (q *Queue) Enqueue(run model.Run) error {
q.mu.Lock()
defer q.mu.Unlock()
if q.stopping {
return errors.New("queue is stopping")
}
if q.active != nil && run.JobID != "" && q.active.JobID == run.JobID {
return ErrDuplicate
}
for _, item := range q.pending {
if run.JobID != "" && item.JobID == run.JobID {
return ErrDuplicate
}
}
run.Status = "queued"
q.pending = append(q.pending, run)
sort.SliceStable(q.pending, func(i, j int) bool { return q.pending[i].Priority > q.pending[j].Priority })
q.emitLocked()
q.cond.Signal()
return nil
}
func (q *Queue) Run(ctx context.Context) {
defer q.doneOnce.Do(func() { close(q.done) })
for {
q.mu.Lock()
for len(q.pending) == 0 && !q.stopping {
q.cond.Wait()
}
if q.stopping {
q.mu.Unlock()
return
}
run := q.pending[0]
q.pending = q.pending[1:]
now := time.Now().UTC()
run.Status, run.StartedAt = "running", &now
workCtx, cancel := context.WithCancel(ctx)
q.cancel, q.active = cancel, &run
q.emitLocked()
q.mu.Unlock()
result := q.handler(workCtx, run)
finished := time.Now().UTC()
result.FinishedAt = &finished
q.mu.Lock()
q.active, q.cancel = nil, nil
q.history = append(q.history, result)
if len(q.history) > 500 {
q.history = q.history[len(q.history)-500:]
}
q.emitLocked()
q.mu.Unlock()
}
}
func (q *Queue) Wait(ctx context.Context) error {
select {
case <-q.done:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
func (q *Queue) Cancel(id string) bool {
q.mu.Lock()
defer q.mu.Unlock()
if q.active != nil && q.active.ID == id && q.cancel != nil {
q.cancel()
return true
}
for i := range q.pending {
if q.pending[i].ID == id {
run := q.pending[i]
q.pending = append(q.pending[:i], q.pending[i+1:]...)
now := time.Now().UTC()
run.Status, run.ErrorCode, run.FinishedAt = "cancelled", "cancelled", &now
q.history = append(q.history, run)
q.emitLocked()
return true
}
}
return false
}
func (q *Queue) SetPaused(id string, paused bool) bool {
q.mu.Lock()
defer q.mu.Unlock()
if q.active == nil || q.active.ID != id {
return false
}
expected := "running"
next := "paused"
if !paused {
expected, next = "paused", "running"
}
if q.active.Status != expected {
return false
}
q.active.Status = next
q.emitLocked()
return true
}
func (q *Queue) Stop() {
q.mu.Lock()
q.stopping = true
if q.cancel != nil {
q.cancel()
}
q.cond.Broadcast()
q.mu.Unlock()
}
func (q *Queue) Snapshot() []model.Run {
q.mu.Lock()
defer q.mu.Unlock()
return q.snapshotLocked()
}
func (q *Queue) UpdateActive(id string, update func(*model.Run)) bool {
q.mu.Lock()
defer q.mu.Unlock()
if q.active == nil || q.active.ID != id {
return false
}
update(q.active)
return true
}
func (q *Queue) ClearHistory() int {
q.mu.Lock()
defer q.mu.Unlock()
count := len(q.history)
q.history = nil
q.emitLocked()
return count
}
func (q *Queue) snapshotLocked() []model.Run {
result := append([]model.Run{}, q.history...)
if q.active != nil {
result = append(result, *q.active)
}
result = append(result, q.pending...)
return result
}
func (q *Queue) emitLocked() {
if q.onChange != nil {
q.onChange(q.snapshotLocked())
}
}