record-asr/internal/store/store.go

// Package store persists transcription tasks to disk.
//
// State lives in <DataDir>/tasks.json plus <DataDir>/uploads/<id><ext> audio
// files. The store is the single source of truth for the task queue and is
// safe for concurrent use from the HTTP handlers and the worker.
package store

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"

	"git.misaka.ren/m1saka/record-asr/internal/transcribe"
)

const (
	StatusPending   = "pending"
	StatusProcessing = "processing"
	StatusCompleted  = "completed"
	StatusFailed    = "failed"
	StatusSkipped   = "skipped" // user chose to abandon a failed task ("continue")
)

// Task is the persisted representation of one transcription request.
type Task struct {
	ID           string                `json:"id"`
	Filename     string                `json:"filename"`       // path relative to DataDir, e.g. uploads/<id>.wav
	OriginalName string                `json:"original_name"`
	Status       string                `json:"status"`
	CreatedAt    time.Time             `json:"created_at"`
	UpdatedAt    time.Time             `json:"updated_at"`
	Attempts     int                   `json:"attempts"`
	MaxRetries   int                   `json:"max_retries"`
	NextRetryAt  *time.Time            `json:"next_retry_at,omitempty"`
	Error        string                `json:"error,omitempty"`
	Duration     float64               `json:"duration,omitempty"`
	Result       *transcribe.Result    `json:"result,omitempty"`
}

// IsTerminal reports whether the task is in a state that needs no more work.
func (t *Task) IsTerminal() bool {
	return t.Status == StatusCompleted || t.Status == StatusSkipped
}

// Store is the on-disk task repository.
type Store struct {
	dir        string
	tasksFile  string
	uploadsDir string

	mu       sync.Mutex
	tasks    map[string]*Task
	sessions map[string]*UploadSession // in-progress chunked uploads
}

// Open initializes the store at dir, creating directories and loading any
// persisted tasks. Tasks left "processing" from a crash are reset to
// "pending" so they re-run on the next worker tick (the crash-recovery half
// of the "continue on error" contract).
func Open(dir string) (*Store, error) {
	uploads := filepath.Join(dir, "uploads")
	for _, d := range []string{dir, uploads} {
		if err := os.MkdirAll(d, 0o755); err != nil {
			return nil, fmt.Errorf("mkdir %s: %w", d, err)
		}
	}
	s := &Store{
		dir:        dir,
		tasksFile:  filepath.Join(dir, "tasks.json"),
		uploadsDir: uploads,
		tasks:      map[string]*Task{},
		sessions:   map[string]*UploadSession{},
	}
	if err := s.load(); err != nil {
		return nil, err
	}
	// In-memory upload sessions do not survive a restart; sweep any .part
	// files left behind by an interrupted chunked upload so they don't leak.
	s.sweepParts()
	return s, nil
}

func (s *Store) load() error {
	data, err := os.ReadFile(s.tasksFile)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return fmt.Errorf("read tasks file: %w", err)
	}
	var list []*Task
	if err := json.Unmarshal(data, &list); err != nil {
		return fmt.Errorf("parse tasks file: %w", err)
	}
	for _, t := range list {
		if t.ID == "" {
			continue
		}
		// Crash recovery: a task that was mid-processing when we died cannot
		// be trusted; requeue it.
		if t.Status == StatusProcessing {
			t.Status = StatusPending
			t.UpdatedAt = time.Now().UTC()
		}
		s.tasks[t.ID] = t
	}
	return nil
}

// save persists all tasks atomically. Caller must hold s.mu.
func (s *Store) save() error {
	list := s.snapshotLocked()
	data, err := json.MarshalIndent(list, "", "  ")
	if err != nil {
		return err
	}
	tmp := s.tasksFile + ".tmp"
	if err := os.WriteFile(tmp, data, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, s.tasksFile)
}

// Create reads an uploaded audio blob, stores it, and enqueues a pending task.
func (s *Store) Create(originalName string, r io.Reader, maxBytes int64, maxRetries int) (*Task, error) {
	ext := filepath.Ext(originalName)
	id := newID()
	rel := filepath.Join("uploads", id+ext)
	dst := filepath.Join(s.dir, rel)

	// Copy with a size cap so a malicious upload can't exhaust disk. The cap
	// is inclusive: a body of exactly maxBytes is accepted, and anything past
	// it is rejected.
	tmp := dst + ".tmp"
	f, err := os.Create(tmp)
	if err != nil {
		return nil, fmt.Errorf("create upload: %w", err)
	}
	n, err := io.Copy(f, &limitedReader{r: r, max: maxBytes})
	f.Close()
	if err != nil {
		os.Remove(tmp)
		return nil, fmt.Errorf("read upload: %w", err)
	}
	_ = n // bytes copied (<= maxBytes); the cap above is the real guard
	if err := os.Rename(tmp, dst); err != nil {
		os.Remove(tmp)
		return nil, fmt.Errorf("save upload: %w", err)
	}

	now := time.Now().UTC()
	t := &Task{
		ID:           id,
		Filename:     rel,
		OriginalName: originalName,
		Status:       StatusPending,
		CreatedAt:    now,
		UpdatedAt:    now,
		MaxRetries:   maxRetries,
	}
	s.mu.Lock()
	s.tasks[id] = t
	serr := s.save()
	s.mu.Unlock()
	if serr != nil {
		return nil, fmt.Errorf("persist task: %w", serr)
	}
	return t, nil
}

// Get returns a copy of a task by id, or nil if absent.
func (s *Store) Get(id string) *Task {
	s.mu.Lock()
	defer s.mu.Unlock()
	if t, ok := s.tasks[id]; ok {
		cp := *t
		return &cp
	}
	return nil
}

// List returns all tasks, newest-first.
func (s *Store) List() []*Task {
	s.mu.Lock()
	defer s.mu.Unlock()
	list := s.snapshotLocked()
	sort.Slice(list, func(i, j int) bool {
		return list[i].CreatedAt.After(list[j].CreatedAt)
	})
	return list
}

func (s *Store) snapshotLocked() []*Task {
	list := make([]*Task, 0, len(s.tasks))
	for _, t := range s.tasks {
		cp := *t
		list = append(list, &cp)
	}
	return list
}

// Delete removes a task and its uploaded audio file.
func (s *Store) Delete(id string) error {
	s.mu.Lock()
	t, ok := s.tasks[id]
	if !ok {
		s.mu.Unlock()
		return ErrNotFound
	}
	delete(s.tasks, id)
	serr := s.save()
	s.mu.Unlock()
	if serr != nil {
		return fmt.Errorf("persist deletion: %w", serr)
	}
	// Best-effort file cleanup; a missing file is not an error.
	if t != nil && t.Filename != "" {
		_ = os.Remove(filepath.Join(s.dir, t.Filename))
	}
	return nil
}

// Claim atomically selects the oldest runnable task (pending and whose retry
// timer has elapsed) and marks it processing. Returns nil if none is ready.
// This is how multiple workers avoid double-processing.
func (s *Store) Claim(now time.Time) *Task {
	s.mu.Lock()
	defer s.mu.Unlock()
	var pick *Task
	var pickT time.Time
	for _, t := range s.tasks {
		if t.Status != StatusPending {
			continue
		}
		if t.NextRetryAt != nil && t.NextRetryAt.After(now) {
			continue
		}
		if pick == nil || t.CreatedAt.Before(pickT) {
			pick = t
			pickT = t.CreatedAt
		}
	}
	if pick == nil {
		return nil
	}
	pick.Status = StatusProcessing
	pick.Attempts++
	pick.UpdatedAt = now.UTC()
	cp := *pick
	_ = s.save()
	return &cp
}

// Complete marks a task done with its transcript.
func (s *Store) Complete(id string, res *transcribe.Result) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	t, ok := s.tasks[id]
	if !ok {
		return ErrNotFound
	}
	t.Status = StatusCompleted
	t.Result = res
	if res != nil {
		t.Duration = res.Duration
	}
	t.Error = ""
	t.NextRetryAt = nil
	t.UpdatedAt = time.Now().UTC()
	return s.save()
}

// Fail records a processing failure. If attempts remain within MaxRetries the
// task is requeued (pending) with an exponential backoff; otherwise it is
// marked failed and left for manual retry or skip.
func (s *Store) Fail(id, errMsg string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	t, ok := s.tasks[id]
	if !ok {
		return ErrNotFound
	}
	t.Error = errMsg
	t.UpdatedAt = time.Now().UTC()
	if t.Attempts <= t.MaxRetries {
		// Retry: requeue with backoff.
		t.Status = StatusPending
		at := time.Now().UTC().Add(backoff(t.Attempts))
		t.NextRetryAt = &at
	} else {
		t.Status = StatusFailed
		t.NextRetryAt = nil
	}
	return s.save()
}

// Retry manually requeues a (typically failed) task: resets the attempt count
// and clears any error so the worker picks it up immediately.
func (s *Store) Retry(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	t, ok := s.tasks[id]
	if !ok {
		return ErrNotFound
	}
	t.Status = StatusPending
	t.Attempts = 0
	t.Error = ""
	t.NextRetryAt = nil
	t.UpdatedAt = time.Now().UTC()
	return s.save()
}

// Skip abandons a task: it becomes terminal ("skipped") and the queue moves
// on. This is the "continue" half of error handling — give up on this one.
func (s *Store) Skip(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	t, ok := s.tasks[id]
	if !ok {
		return ErrNotFound
	}
	t.Status = StatusSkipped
	t.NextRetryAt = nil
	t.UpdatedAt = time.Now().UTC()
	return s.save()
}

// AudioPath returns the absolute path to a task's uploaded audio file.
func (s *Store) AudioPath(t *Task) string {
	return filepath.Join(s.dir, t.Filename)
}

// Has reports whether a task id exists.
func (s *Store) Has(id string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	_, ok := s.tasks[id]
	return ok
}

// backoff is an exponential schedule capped at 60s.
func backoff(attempts int) time.Duration {
	d := time.Duration(5<<uint(attempts-1)) * time.Second // 5s, 10s, 20s, 40s...
	if d > 60*time.Second {
		d = 60 * time.Second
	}
	if d < 1*time.Second {
		d = 1 * time.Second
	}
	return d
}

// newID returns a short, collision-resistant id (8 random bytes hex).
func newID() string {
	b := make([]byte, 8)
	if _, err := rand.Read(b); err != nil {
		// Fallback to time-based if crypto/rand ever fails.
		return fmt.Sprintf("%x", time.Now().UnixNano())
	}
	return hex.EncodeToString(b)
}

// limitedReader wraps io.Reader with an inclusive byte cap: a body that ends
// at exactly max is accepted (max is the allowed maximum), while any bytes
// past it are rejected so an oversized body can't fill the disk.
type limitedReader struct {
	r   io.Reader
	max int64
	n   int64
}

func (l *limitedReader) Read(p []byte) (int, error) {
	if l.n < l.max {
		remain := l.max - l.n
		if int64(len(p)) > remain {
			p = p[:remain]
		}
		n, err := l.r.Read(p)
		l.n += int64(n)
		return n, err
	}
	// Already at the cap. Tell an exactly-at-cap body (reader now exhausted)
	// from one that still has bytes (over the cap): read once — any non-zero
	// return means the body exceeded the cap.
	n, err := l.r.Read(p)
	l.n += int64(n)
	if n > 0 {
		return n, fmt.Errorf("upload exceeds %d bytes", l.max)
	}
	return 0, err
}