URBM/internal/queue/queue.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())
	}
}