hamlab/internal/tasks/tasks.go

// Package tasks runs long commands (apt upgrades, and later anything else the UI kicks
// off) in the background, keeps their output, and lets any number of clients follow it
// live. It is also the task history: what ran, when, and how it ended.
package tasks

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"os/exec"
	"sort"
	"sync"
	"syscall"
	"time"
)

const (
	maxOutput  = 1 << 20 // bytes kept per task
	maxHistory = 200
)

type Task struct {
	ID       string
	Kind     string // "updates", "cron", …
	Title    string
	Command  string
	Started  time.Time
	Finished time.Time
	Running  bool
	ExitCode int
	Err      string

	mu       sync.Mutex
	buf      bytes.Buffer
	subs     map[chan []byte]struct{}
	cancel   context.CancelFunc
	onCancel func()
}

type Snapshot struct {
	ID       string     `json:"id"`
	Kind     string     `json:"kind"`
	Title    string     `json:"title"`
	Command  string     `json:"command"`
	Started  time.Time  `json:"started"`
	Finished *time.Time `json:"finished,omitempty"`
	Running  bool       `json:"running"`
	ExitCode int        `json:"exit_code"`
	Err      string     `json:"error,omitempty"`
	Output   string     `json:"output,omitempty"`
}

type Manager struct {
	mu    sync.Mutex
	tasks map[string]*Task
}

func New() *Manager { return &Manager{tasks: map[string]*Task{}} }

func newID() string {
	b := make([]byte, 6)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}

// Start launches cmd with its output captured into a new task. cmd must not have
// Stdout/Stderr set. The task ends when the process exits or Cancel is called.
func (m *Manager) Start(kind, title string, cmd *exec.Cmd) (*Task, error) {
	ctx, cancel := context.WithCancel(context.Background())
	t := &Task{ID: newID(), Kind: kind, Title: title, Command: cmd.String(), Started: time.Now(), Running: true,
		subs: map[chan []byte]struct{}{}, cancel: cancel}
	w := &writer{t: t}
	cmd.Stdout, cmd.Stderr = w, w
	// Its own process group: a task is usually `sh -c "a && b"`, and cancelling must reach a and
	// b, not just the shell — killing only the shell left compose running and, because compose
	// still held the output pipe, left the task "running" too. WaitDelay is the backstop for
	// anything that escapes the group and keeps the pipe open.
	if cmd.SysProcAttr == nil {
		cmd.SysProcAttr = &syscall.SysProcAttr{}
	}
	cmd.SysProcAttr.Setpgid = true
	cmd.WaitDelay = 5 * time.Second
	if err := cmd.Start(); err != nil {
		cancel()
		return nil, err
	}
	m.mu.Lock()
	m.tasks[t.ID] = t
	m.prune()
	m.mu.Unlock()

	go func() {
		done := make(chan error, 1)
		go func() { done <- cmd.Wait() }()
		var err error
		select {
		case err = <-done:
		case <-ctx.Done():
			t.mu.Lock()
			hook := t.onCancel
			t.mu.Unlock()
			if hook != nil {
				go hook()
			}
			_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM) // asked first: compose and apt both stop cleanly
			select {
			case err = <-done:
			case <-time.After(4 * time.Second):
				_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
				_ = cmd.Process.Kill()
				err = <-done
			}
			t.mu.Lock()
			t.Err = "cancelled"
			t.mu.Unlock()
		}
		t.mu.Lock()
		t.Running, t.Finished = false, time.Now()
		if err != nil {
			var ee *exec.ExitError
			if errors.As(err, &ee) {
				t.ExitCode = ee.ExitCode()
			} else {
				t.ExitCode = -1
				if t.Err == "" {
					t.Err = err.Error()
				}
			}
		}
		for ch := range t.subs {
			close(ch)
		}
		t.subs = map[chan []byte]struct{}{}
		t.mu.Unlock()
		cancel()
	}()
	return t, nil
}

// prune keeps the newest maxHistory finished tasks. Caller holds m.mu.
func (m *Manager) prune() {
	if len(m.tasks) <= maxHistory {
		return
	}
	var finished []*Task
	for _, t := range m.tasks {
		if !t.running() {
			finished = append(finished, t)
		}
	}
	sort.Slice(finished, func(i, j int) bool { return finished[i].Started.Before(finished[j].Started) })
	for i := 0; i < len(finished) && len(m.tasks) > maxHistory; i++ {
		delete(m.tasks, finished[i].ID)
	}
}

func (m *Manager) Get(id string) (*Task, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	t, ok := m.tasks[id]
	return t, ok
}

// List returns every task newest first, without output.
func (m *Manager) List() []Snapshot {
	m.mu.Lock()
	all := make([]*Task, 0, len(m.tasks))
	for _, t := range m.tasks {
		all = append(all, t)
	}
	m.mu.Unlock()
	sort.Slice(all, func(i, j int) bool { return all[i].Started.After(all[j].Started) })
	out := make([]Snapshot, 0, len(all))
	for _, t := range all {
		out = append(out, t.Snapshot(false))
	}
	return out
}

// RunningOf reports whether a task of the given kind is in flight.
func (m *Manager) RunningOf(kind string) (*Task, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, t := range m.tasks {
		if t.Kind == kind && t.running() {
			return t, true
		}
	}
	return nil, false
}

// running reads the flag the waiter goroutine clears, under the task's own lock.
func (t *Task) running() bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.Running
}

func (t *Task) Snapshot(withOutput bool) Snapshot {
	t.mu.Lock()
	defer t.mu.Unlock()
	s := Snapshot{ID: t.ID, Kind: t.Kind, Title: t.Title, Command: t.Command, Started: t.Started, Running: t.Running, ExitCode: t.ExitCode, Err: t.Err}
	if !t.Running {
		f := t.Finished
		s.Finished = &f
	}
	if withOutput {
		s.Output = t.buf.String()
	}
	return s
}

// Subscribe returns the output so far plus a channel of further chunks, closed when the
// task ends. Call unsub when done following.
func (t *Task) Subscribe() (backlog []byte, ch chan []byte, unsub func()) {
	t.mu.Lock()
	defer t.mu.Unlock()
	backlog = append([]byte(nil), t.buf.Bytes()...)
	if !t.Running {
		c := make(chan []byte)
		close(c)
		return backlog, c, func() {}
	}
	ch = make(chan []byte, 256)
	t.subs[ch] = struct{}{}
	return backlog, ch, func() {
		t.mu.Lock()
		if _, ok := t.subs[ch]; ok {
			delete(t.subs, ch)
			close(ch)
		}
		t.mu.Unlock()
	}
}

func (t *Task) Cancel() { t.cancel() }

// OnCancel sets something to do, besides stopping the task's own processes, when the task is
// cancelled: whatever it started that isn't its child (a command inside a container).
func (t *Task) OnCancel(fn func()) {
	t.mu.Lock()
	t.onCancel = fn
	t.mu.Unlock()
}

type writer struct{ t *Task }

func (w *writer) Write(p []byte) (int, error) {
	t := w.t
	t.mu.Lock()
	defer t.mu.Unlock()
	t.buf.Write(p)
	if t.buf.Len() > maxOutput {
		b := t.buf.Bytes()
		t.buf = *bytes.NewBuffer(append([]byte("…\n"), b[len(b)-maxOutput*3/4:]...))
	}
	chunk := append([]byte(nil), p...)
	for ch := range t.subs {
		select {
		case ch <- chunk:
		default: // slow follower: drop the chunk rather than stall the process
		}
	}
	return len(p), nil
}