vibedit/jobs.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"sort"
	"sync"
	"time"
)

// Hub fans server-sent events out to every connected editor window.
type Hub struct {
	mu   sync.Mutex
	subs map[chan []byte]struct{}
}

func newHub() *Hub { return &Hub{subs: map[chan []byte]struct{}{}} }

func (h *Hub) Subscribe() chan []byte {
	ch := make(chan []byte, 256)
	h.mu.Lock()
	h.subs[ch] = struct{}{}
	h.mu.Unlock()
	return ch
}

func (h *Hub) Unsubscribe(ch chan []byte) {
	h.mu.Lock()
	delete(h.subs, ch)
	h.mu.Unlock()
}

func (h *Hub) Publish(event string, v any) {
	msg := sseMessage(event, v)
	h.mu.Lock()
	defer h.mu.Unlock()
	for ch := range h.subs {
		select {
		case ch <- msg:
		default: // slow client; it will catch up on the next update
		}
	}
}

func sseMessage(event string, v any) []byte {
	b, _ := json.Marshal(v)
	return []byte(fmt.Sprintf("event: %s\ndata: %s\n\n", event, b))
}

// Job is a long-running ffmpeg task (an export or a preview proxy).
type Job struct {
	ID       string   `json:"id"`
	Kind     string   `json:"kind"`
	Title    string   `json:"title"`
	SourceID string   `json:"sourceId,omitempty"`
	State    string   `json:"state"` // running, done, error, canceled
	Progress float64  `json:"progress"`
	Speed    string   `json:"speed,omitempty"`
	ETA      float64  `json:"eta"`
	Elapsed  float64  `json:"elapsed"`
	Step     string   `json:"step,omitempty"`
	Error    string   `json:"error,omitempty"`
	Outputs  []string `json:"outputs,omitempty"`

	cancel   context.CancelFunc
	started  time.Time
	lastSent time.Time
}

type JobManager struct {
	mu     sync.Mutex
	jobs   map[string]*Job
	hub    *Hub
	wg     sync.WaitGroup
	onDone func()
}

func newJobManager(hub *Hub, onDone func()) *JobManager {
	return &JobManager{jobs: map[string]*Job{}, hub: hub, onDone: onDone}
}

// JobCtl is handed to a running job so it can report progress.
type JobCtl struct {
	m   *JobManager
	job *Job
}

func (c *JobCtl) Step(s string) {
	c.m.mu.Lock()
	c.job.Step = s
	c.m.mu.Unlock()
	c.m.publish(c.job, true)
}

func (c *JobCtl) Progress(p float64, speed string) {
	if p < 0 {
		p = 0
	}
	if p > 1 {
		p = 1
	}
	c.m.mu.Lock()
	j := c.job
	j.Progress = p
	j.Speed = speed
	j.Elapsed = time.Since(j.started).Seconds()
	if p > 0.01 {
		j.ETA = j.Elapsed / p * (1 - p)
	}
	c.m.mu.Unlock()
	c.m.publish(j, false)
}

// AddOutput records a finished output file.
func (c *JobCtl) AddOutput(path string) {
	c.m.mu.Lock()
	c.job.Outputs = append(c.job.Outputs, path)
	c.m.mu.Unlock()
}

func (m *JobManager) Start(kind, title, sourceID string, fn func(ctx context.Context, c *JobCtl) error) *Job {
	ctx, cancel := context.WithCancel(context.Background())
	j := &Job{
		ID:       randomHex(6),
		Kind:     kind,
		Title:    title,
		SourceID: sourceID,
		State:    "running",
		cancel:   cancel,
		started:  time.Now(),
	}
	m.mu.Lock()
	m.jobs[j.ID] = j
	m.pruneLocked()
	m.mu.Unlock()
	m.publish(j, true)

	m.wg.Add(1)
	go func() {
		defer m.wg.Done()
		defer cancel()
		err := fn(ctx, &JobCtl{m: m, job: j})
		m.mu.Lock()
		j.Elapsed = time.Since(j.started).Seconds()
		j.ETA = 0
		switch {
		case ctx.Err() != nil:
			j.State = "canceled"
		case err != nil:
			j.State = "error"
			j.Error = err.Error()
		default:
			j.State = "done"
			j.Progress = 1
		}
		state, elapsed := j.State, j.Elapsed
		m.mu.Unlock()
		if kind == "export" {
			if err != nil && state == "error" {
				log.Printf("export failed after %.1fs: %v", elapsed, err)
			} else {
				log.Printf("export %s after %.1fs", state, elapsed)
			}
		}
		m.publish(j, true)
		if m.onDone != nil {
			m.onDone()
		}
	}()
	return j
}

func (m *JobManager) publish(j *Job, force bool) {
	m.mu.Lock()
	if !force && time.Since(j.lastSent) < 200*time.Millisecond {
		m.mu.Unlock()
		return
	}
	j.lastSent = time.Now()
	snap := *j
	m.mu.Unlock()
	m.hub.Publish("job", &snap)
}

// pruneLocked forgets old finished jobs.
func (m *JobManager) pruneLocked() {
	var finished []*Job
	for _, j := range m.jobs {
		if j.State != "running" {
			finished = append(finished, j)
		}
	}
	if len(finished) <= 20 {
		return
	}
	sort.Slice(finished, func(a, b int) bool { return finished[a].started.Before(finished[b].started) })
	for _, j := range finished[:len(finished)-20] {
		delete(m.jobs, j.ID)
	}
}

func (m *JobManager) Snapshot() []Job {
	m.mu.Lock()
	defer m.mu.Unlock()
	out := make([]Job, 0, len(m.jobs))
	for _, j := range m.jobs {
		out = append(out, *j)
	}
	sort.Slice(out, func(a, b int) bool { return out[a].started.Before(out[b].started) })
	return out
}

func (m *JobManager) Cancel(id string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	if j, ok := m.jobs[id]; ok && j.State == "running" {
		j.cancel()
		return true
	}
	return false
}

func (m *JobManager) CancelAll() {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, j := range m.jobs {
		if j.State == "running" {
			j.cancel()
		}
	}
}

// Find returns the running job of the given kind for a source, if any.
func (m *JobManager) Find(kind, sourceID string) *Job {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, j := range m.jobs {
		if j.Kind == kind && j.SourceID == sourceID && j.State == "running" {
			return j
		}
	}
	return nil
}

func (m *JobManager) AnyRunning(kind string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, j := range m.jobs {
		if j.Kind == kind && j.State == "running" {
			return true
		}
	}
	return false
}

func (m *JobManager) Wait(timeout time.Duration) {
	done := make(chan struct{})
	go func() { m.wg.Wait(); close(done) }()
	select {
	case <-done:
	case <-time.After(timeout):
	}
}