tapir/internal/uploads/worker.go

package uploads

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"sync"
	"time"

	"git.d-ma.be/mathias/tapir/internal/domain"
	"git.d-ma.be/mathias/tapir/internal/errkind"
	"git.d-ma.be/mathias/tapir/internal/progress"
)

// ErrStopped is returned by Submit once the worker has stopped.
var ErrStopped = errors.New("uploads: worker stopped")

// ErrQueueFull is returned by Submit when the queue, or one user's share of it, is at
// its cap. Every queued job holds up to the upload cap on disk until its turn, so an
// unbounded queue is a disk-filling lever (ADR-035 review).
var ErrQueueFull = errors.New("uploads: queue full")

// Converter turns the uploaded file into the 16 kHz mono PCM WAV whisper wants.
type Converter interface {
	ToWAV(ctx context.Context, src, dst string) error
}

// Transcriber turns a WAV into a timed transcription. The gateway chain
// (iguana primary, berget fallback) implements it.
type Transcriber interface {
	Transcribe(ctx context.Context, wavPath string) (domain.Transcription, error)
}

// Planner is an optional extra on a Transcriber: it declares its own phases (for
// example one per engine) so the page can show them. A Transcriber that does not
// implement it gets a single "transcribing" phase from the worker. Whatever a Planner
// declares it reports itself, through progress.From(ctx) (ADR-035).
type Planner interface {
	Plan() []progress.Planned
}

// WorkerStore is the slice of the private transcript store the worker drives.
type WorkerStore interface {
	StartTranscribing(ctx context.Context, userID, id string) error
	CompleteUpload(ctx context.Context, userID, id string, t domain.Transcription) error
	FailUpload(ctx context.Context, userID, id string, kind domain.UploadErrorKind) error
	SaveProgress(ctx context.Context, userID, id string, steps []domain.ProgressStep, at time.Time) error
	Heartbeat(ctx context.Context, userID, id string, at time.Time) error
}

// WorkerOption tunes a Worker.
type WorkerOption func(*Worker)

// WithHeartbeat sets how often a running job stamps its row alive and how often it logs
// which phase is still running. The defaults are 10 s and 5 min (ADR-035).
func WithHeartbeat(beat, still time.Duration) WorkerOption {
	return func(w *Worker) { w.beatEvery, w.stillEvery = beat, still }
}

// WithQueueLimits bounds how many jobs may WAIT: in total, and per user. A job being
// processed is not waiting. The defaults are 20 and 3 (one instance, one job at a time).
func WithQueueLimits(total, perUser int) WorkerOption {
	return func(w *Worker) { w.maxQueued, w.maxPerUser = total, perUser }
}

// WithMaxDuration makes the worker measure the converted audio and fail a job whose real
// length exceeds the limit, or falls far short of what the file claimed. Zero (the
// default) measures nothing. The limit is first checked against the file's header, which
// can lie (ADR-035 review).
func WithMaxDuration(d time.Duration) WorkerOption { return func(w *Worker) { w.maxDuration = d } }

// WithLogger replaces slog.Default() as the worker's logger.
func WithLogger(l *slog.Logger) WorkerOption { return func(w *Worker) { w.log = l } }

const (
	defaultMaxQueued  = 20
	defaultMaxPerUser = 3
	defaultBeatEvery  = 10 * time.Second
	defaultStillEvery = 5 * time.Minute
)

// jobSink binds the store to one job so a progress.Tracker can persist through it.
type jobSink struct {
	st         WorkerStore
	userID, id string
}

func (s jobSink) SaveProgress(ctx context.Context, steps []domain.ProgressStep, at time.Time) error {
	return s.st.SaveProgress(ctx, s.userID, s.id, steps, at)
}

func (s jobSink) Beat(ctx context.Context, at time.Time) error {
	return s.st.Heartbeat(ctx, s.userID, s.id, at)
}

// Worker transcribes accepted uploads in-process, ONE AT A TIME (ADR-033). It is
// not a second deployable: it shares the web process's lifetime, which is why it
// inherits ADR-018's single-replica assumption — the queue is process-local. The
// concurrency of one is deliberate: two simultaneous 90-minute jobs on iguana were
// never measured, and a Metal OOM there latches the whisper backend (infra#273).
type Worker struct {
	store   WorkerStore
	conv    Converter
	tr      Transcriber
	tmpDir  string
	timeout time.Duration

	log                   *slog.Logger
	beatEvery, stillEvery time.Duration
	maxDuration           time.Duration
	maxQueued, maxPerUser int

	mu      sync.Mutex
	queue   []Job
	stopped bool
	wake    chan struct{}
}

// NewWorker wires a Worker. timeout bounds each job (conversion + transcription);
// a job that outlives it is failed rather than left spinning.
func NewWorker(s WorkerStore, c Converter, t Transcriber, tmpDir string, timeout time.Duration, opts ...WorkerOption) *Worker {
	w := &Worker{
		store: s, conv: c, tr: t, tmpDir: tmpDir, timeout: timeout, wake: make(chan struct{}, 1),
		beatEvery: defaultBeatEvery, stillEvery: defaultStillEvery,
		maxQueued: defaultMaxQueued, maxPerUser: defaultMaxPerUser,
	}
	for _, o := range opts {
		o(w)
	}
	return w
}

// HasRoom reports whether a job from this user would be accepted right now. It lets the
// intake refuse BEFORE spooling up to the upload cap and running ffprobe; Submit still
// enforces the same limits atomically, so a race costs at worst one refused submit.
func (w *Worker) HasRoom(userID string) bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.stopped || len(w.queue) >= w.maxQueued {
		return false
	}
	mine := 0
	for _, q := range w.queue {
		if q.UserID == userID {
			mine++
		}
	}
	return mine < w.maxPerUser
}

// Submit queues a job. It never blocks. It satisfies Submitter.
func (w *Worker) Submit(j Job) error {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.stopped {
		return ErrStopped
	}
	if len(w.queue) >= w.maxQueued {
		return ErrQueueFull
	}
	mine := 0
	for _, q := range w.queue {
		if q.UserID == j.UserID {
			mine++
		}
	}
	if mine >= w.maxPerUser {
		return ErrQueueFull
	}
	w.queue = append(w.queue, j)
	select {
	case w.wake <- struct{}{}:
	default:
	}
	return nil
}

// Run processes queued jobs until ctx ends, then deletes the audio of anything
// still queued and refuses new jobs. The job running at shutdown is recorded interrupted
// by Process; the rows of jobs still queued stay queued and are failed as interrupted by
// the start-up sweep (store.FailInterruptedUploads): the audio is not retained, so they
// cannot resume.
func (w *Worker) Run(ctx context.Context) {
	defer w.stop()
	for {
		job, ok := w.next(ctx)
		if !ok {
			return
		}
		w.Process(ctx, job)
	}
}

// next blocks for the next job. Cancellation wins over a non-empty queue: a job
// still queued at shutdown must not be STARTED (it would begin a long conversion
// with an already-cancelled context); stop() deletes its audio instead.
func (w *Worker) next(ctx context.Context) (Job, bool) {
	for {
		if ctx.Err() != nil {
			return Job{}, false
		}
		w.mu.Lock()
		if len(w.queue) > 0 {
			j := w.queue[0]
			w.queue = w.queue[1:]
			w.mu.Unlock()
			return j, true
		}
		w.mu.Unlock()
		select {
		case <-w.wake:
		case <-ctx.Done():
			return Job{}, false
		}
	}
}

func (w *Worker) stop() {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.stopped = true
	for _, j := range w.queue {
		_ = os.Remove(j.Path)
	}
	w.queue = nil
}

// Process runs one job: row -> transcribing, convert, transcribe, store. The
// uploaded file and the converted WAV are deleted on every path — the audio is
// never retained (ADR-033). A failure is recorded with a TYPED kind; a shutdown
// mid-job is recorded as interrupted at once (see fail).
//
// Every phase is reported through one progress.Tracker, which saves it to the row for
// the page and logs it, and a heartbeat stamps the row on its own while the job runs
// (ADR-035). Nothing logged or saved here carries transcript text or the file's name.
func (w *Worker) Process(ctx context.Context, job Job) {
	defer func() { _ = os.Remove(job.Path) }()
	base := w.log
	if base == nil {
		base = slog.Default()
	}
	log := base.With("upload", job.UploadID)
	start := time.Now()

	if err := w.store.StartTranscribing(ctx, job.UserID, job.UploadID); err != nil {
		log.Error("upload: start transcribing", "kind", errkind.Of(err))
		w.fail(ctx, job, log, domain.UploadErrTranscriptionFailed)
		return
	}

	plan := []progress.Planned{{Key: "received", Label: "Received"}, {Key: "audio", Label: "Preparing the audio"}}
	planner, planned := w.tr.(Planner)
	if planned {
		plan = append(plan, planner.Plan()...)
	} else {
		plan = append(plan, progress.Planned{Key: "transcribe", Label: "Transcribing"})
	}
	plan = append(plan, progress.Planned{Key: "save", Label: "Saving the transcript"})
	tr := progress.NewTracker(ctx, jobSink{st: w.store, userID: job.UserID, id: job.UploadID}, log, plan)
	tr.Start()
	tr.Begin("received", 0)
	tr.Done("received")

	hbCtx, stopBeat := context.WithCancel(ctx)
	hbDone := make(chan struct{})
	go func() {
		defer close(hbDone)
		tr.Heartbeat(hbCtx, w.beatEvery, w.stillEvery)
	}()
	defer func() { stopBeat(); <-hbDone }() // no heartbeat outlives its job

	// abort records where the job stopped. A shutdown (the parent context ended) is not
	// a failure of the step: it stays as it was, and fail records the row as interrupted.
	abort := func(kind domain.UploadErrorKind) {
		if ctx.Err() == nil {
			tr.FailRunning()
		}
		w.fail(ctx, job, log, kind)
	}

	wav, err := w.tempWAV()
	if err != nil {
		log.Error("upload: temp wav", "kind", errkind.Of(err))
		abort(domain.UploadErrTranscriptionFailed)
		return
	}
	defer func() { _ = os.Remove(wav) }()

	jobCtx, cancel := context.WithTimeout(ctx, w.timeout)
	defer cancel()
	tr.Begin("audio", 0)
	if err := w.conv.ToWAV(jobCtx, job.Path, wav); err != nil {
		log.Error("upload: convert", "kind", errkind.Of(err))
		abort(domain.UploadErrTranscriptionFailed)
		return
	}
	got, tooLong := w.audioLengthProblem(wav)
	if tooLong {
		abort(domain.UploadErrTooLong)
		log.Error("upload: converted audio is over the limit", "audio_s", int(got.Seconds()), "limit_s", int(w.maxDuration.Seconds()))
		return
	}
	if w.maxDuration > 0 && job.Duration > 0 {
		if note := shortfallNote(got, job.Duration); note != "" {
			log.Warn("upload: audio is much shorter than the file claimed", "audio_s", int(got.Seconds()), "claimed_s", int(job.Duration.Seconds()))
			tr.Annotate("audio", note)
		}
	}
	tr.Done("audio")

	if !planned {
		tr.Begin("transcribe", 0)
	}
	t, err := w.tr.Transcribe(domain.WithLanguage(progress.With(jobCtx, tr), job.Language), wav)
	if err != nil {
		log.Error("upload: transcribe", "kind", errkind.Of(err), "elapsed", time.Since(start).Round(time.Second))
		abort(domain.UploadErrTranscriptionFailed)
		return
	}
	if !planned {
		tr.Done("transcribe")
	}

	tr.Begin("save", 0)
	if err := w.store.CompleteUpload(ctx, job.UserID, job.UploadID, t); err != nil {
		log.Error("upload: store transcript", "kind", errkind.Of(err))
		abort(domain.UploadErrTranscriptionFailed)
		return
	}
	tr.Done("save")
	tr.Finish()
	log.Info("upload transcribed", "endpoint", t.Endpoint, "audio_s", t.AudioSeconds,
		"segments", len(t.Segments), "elapsed", time.Since(start).Round(time.Second))
}

// interruptedWriteTimeout bounds the one write a dying process makes to record its job.
const interruptedWriteTimeout = 3 * time.Second

// fail records a transcription failure. When the PARENT context ended the process is
// shutting down: the job is recorded interrupted at once, on a short context of its own,
// because the process may exit before any other pod's start-up sweep would do it (#50).
// The sweep remains the backstop for a process that is killed outright.
func (w *Worker) fail(ctx context.Context, job Job, log *slog.Logger, kind domain.UploadErrorKind) {
	if ctx.Err() != nil {
		kind = domain.UploadErrInterrupted
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(context.WithoutCancel(ctx), interruptedWriteTimeout)
		defer cancel()
	}
	if err := w.store.FailUpload(ctx, job.UserID, job.UploadID, kind); err != nil {
		log.Error("upload: record failure", "kind", errkind.Of(err))
	}
}

func (w *Worker) tempWAV() (string, error) {
	f, err := os.CreateTemp(w.tmpDir, "wav-*")
	if err != nil {
		return "", fmt.Errorf("uploads: temp wav: %w", err)
	}
	name := f.Name()
	if err := f.Close(); err != nil {
		_ = os.Remove(name)
		return "", fmt.Errorf("uploads: temp wav: %w", err)
	}
	return name, nil
}

const (
	// tooLongSlack is how far over the limit the converted audio may be before the claim
	// in the file's header is called a lie: ffmpeg rounds.
	tooLongSlack = 30 * time.Second
	// A file counts as truncated when it converts to under 80% of the length it claimed
	// AND more than two minutes short. Probe estimates for variable-bitrate files are
	// imprecise; refusing real recordings would be worse than the problem.
	incompleteBelow     = 0.8
	incompleteShortfall = 2 * time.Minute
)

// audioLengthProblem measures the converted WAV (16 kHz mono 16-bit PCM is 32,000 bytes a
// second after a 44-byte header) against the limit. It returns the real length and whether
// that is over the limit by more than the tolerance.
func (w *Worker) audioLengthProblem(wav string) (got time.Duration, tooLong bool) {
	if w.maxDuration <= 0 {
		return 0, false
	}
	info, err := os.Stat(wav)
	if err != nil {
		return 0, false
	}
	got = time.Duration(float64(max(info.Size()-44, 0)) / 32000 * float64(time.Second))
	return got, got > w.maxDuration+tooLongSlack
}

// shortfallNote is a remark for the audio step when the converted audio is much shorter
// than the file claimed, or "" when it is not. It is a NOTE and never a failure: a
// container header can over-estimate a complete file (a variable-bitrate MP3 without a Xing
// header claimed 79 minutes for 20 real ones), so refusing would turn away real
// recordings, while saying nothing about a truncated one would hide a partial transcript.
func shortfallNote(got, claimed time.Duration) string {
	if float64(got) >= incompleteBelow*float64(claimed) || claimed-got <= incompleteShortfall {
		return ""
	}
	return fmt.Sprintf("The audio is %s long; the file claimed %s. If the recording was cut off, upload the complete one.",
		FormatDuration(got), FormatDuration(claimed))
}