record-asr/internal/store/upload.go

// upload.go extends the store with resumable chunked uploads: a way to receive
// a large file as many small requests (each under a size-capping proxy's body
// limit, e.g. Cloudflare's 100MB) and reassemble it into one task.
//
// Lifecycle:
//  1. CreateUpload  — allocate an id + an empty "<id>.part" file, record the
//     expected total size and chunk layout.
//  2. AppendChunk    — write chunk #index at offset index*chunkSize into the
//     .part file. Chunks may arrive in any order and re-sending one is
//     idempotent (it overwrites the same offset).
//  3. CompleteUpload — once every chunk has arrived, truncate the .part file
//     to the exact size, rename it to "<id><ext>", and enqueue a pending task
//     (reusing the upload id as the task id).
//  4. AbortUpload    — drop the session and delete the .part file.
//
// Sessions live in memory (the .part file is the on-disk durability). A
// restart therefore loses in-flight sessions; sweepParts cleans up orphaned
// .part files on Open, and gcUploads lazily reaps sessions idle too long.
//
// Locking: Store.mu guards the sessions map; each UploadSession has its own
// mutex guarding its received[] slice and the part file it owns. The rule that
// prevents deadlock is: never acquire Store.mu while holding a session mutex.
// Append/Complete/GC all phase the two locks separately to honour this.

package store

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// sessionTTL is how long an upload session may stay idle before it (and its
// .part file) is garbage-collected.
const sessionTTL = 2 * time.Hour

// UploadSession tracks a resumable chunked upload whose slices are written
// into a single .part file and finalized into a task on completion.
type UploadSession struct {
	UploadID    string
	Filename    string
	Size        int64
	ChunkSize   int64
	TotalChunks int

	// Unexported state — callers reach these only through Store methods.
	received  []bool
	partPath  string
	ext       string
	createdAt time.Time
	lastSeen  time.Time
	mu        sync.Mutex
}

// CreateUpload starts a resumable chunked upload session. It allocates an
// upload id, creates an empty .part file, and records the expected chunk
// layout. maxBytes caps the total file size; chunkSize is the per-chunk size
// (must be > 0). The returned session's id is later reused as the task id.
func (s *Store) CreateUpload(filename string, size, chunkSize, maxBytes int64) (*UploadSession, error) {
	if chunkSize <= 0 {
		return nil, fmt.Errorf("chunk size must be > 0")
	}
	if size <= 0 {
		return nil, fmt.Errorf("size must be > 0")
	}
	if size > maxBytes {
		return nil, fmt.Errorf("upload size %d exceeds cap %d", size, maxBytes)
	}
	ext := filepath.Ext(filename)
	id := newID()
	partPath := filepath.Join(s.uploadsDir, id+".part")
	f, err := os.OpenFile(partPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
	if err != nil {
		return nil, fmt.Errorf("create part file: %w", err)
	}
	f.Close()
	total := int((size + chunkSize - 1) / chunkSize)
	now := time.Now().UTC()
	sess := &UploadSession{
		UploadID:    id,
		Filename:    filename,
		Size:        size,
		ChunkSize:   chunkSize,
		TotalChunks: total,
		received:    make([]bool, total),
		partPath:    partPath,
		ext:         ext,
		createdAt:   now,
		lastSeen:    now,
	}
	s.gcUploads() // reap idle sessions opportunistically before adding a new one
	s.mu.Lock()
	s.sessions[id] = sess
	s.mu.Unlock()
	return sess, nil
}

// AppendChunk writes chunk #index (0-based) of the session's file at offset
// index*chunkSize and marks it received. Chunks may arrive in any order and
// are idempotent (re-sending a chunk overwrites the same offset). The reader
// is capped at maxChunkBytes so a malicious or corrupt body can't overrun the
// slot; the written length must match the expected chunk length.
func (s *Store) AppendChunk(uploadID string, index int, r io.Reader, maxChunkBytes int64) error {
	s.mu.Lock()
	sess, ok := s.sessions[uploadID]
	s.mu.Unlock()
	if !ok {
		return ErrUploadNotFound
	}
	sess.mu.Lock()
	defer sess.mu.Unlock()
	sess.lastSeen = time.Now().UTC()
	if index < 0 || index >= sess.TotalChunks {
		return fmt.Errorf("chunk index %d out of range [0,%d)", index, sess.TotalChunks)
	}
	offset := int64(index) * sess.ChunkSize
	f, err := os.OpenFile(sess.partPath, os.O_WRONLY, 0o644)
	if err != nil {
		return fmt.Errorf("open part: %w", err)
	}
	defer f.Close()
	if _, err := f.Seek(offset, io.SeekStart); err != nil {
		return fmt.Errorf("seek: %w", err)
	}
	n, err := io.Copy(f, &limitedReader{r: r, max: maxChunkBytes})
	if err != nil {
		return fmt.Errorf("write chunk: %w", err)
	}
	// Validate length: every full chunk == chunkSize; the last chunk is the
	// remainder. A short or oversized body is rejected (not marked received).
	expected := sess.ChunkSize
	if index == sess.TotalChunks-1 {
		expected = sess.Size - offset
	}
	if n != expected {
		return fmt.Errorf("chunk %d length mismatch: got %d want %d", index, n, expected)
	}
	sess.received[index] = true
	return nil
}

// CompleteUpload finalizes a chunked upload: verifies every chunk arrived,
// truncates the part file to the exact size (repairing any overwrite from a
// rejected chunk), renames it to the task's audio file, and enqueues a
// pending transcription task. The upload id becomes the task id.
func (s *Store) CompleteUpload(uploadID string, maxRetries int) (*Task, error) {
	s.mu.Lock()
	sess, ok := s.sessions[uploadID]
	s.mu.Unlock()
	if !ok {
		return nil, ErrUploadNotFound
	}
	// Session-scoped work (file ops) under the session lock only.
	sess.mu.Lock()
	for i, got := range sess.received {
		if !got {
			sess.mu.Unlock()
			return nil, fmt.Errorf("missing chunk %d of %d", i, sess.TotalChunks)
		}
	}
	if err := os.Truncate(sess.partPath, sess.Size); err != nil {
		sess.mu.Unlock()
		return nil, fmt.Errorf("truncate part: %w", err)
	}
	rel := filepath.Join("uploads", sess.UploadID+sess.ext)
	dst := filepath.Join(s.dir, rel)
	if err := os.Rename(sess.partPath, dst); err != nil {
		sess.mu.Unlock()
		return nil, fmt.Errorf("finalize upload: %w", err)
	}
	sess.mu.Unlock()

	// Register the task under the store lock (no session lock held).
	now := time.Now().UTC()
	t := &Task{
		ID:           sess.UploadID,
		Filename:     rel,
		OriginalName: sess.Filename,
		Status:       StatusPending,
		CreatedAt:    now,
		UpdatedAt:    now,
		MaxRetries:   maxRetries,
	}
	s.mu.Lock()
	delete(s.sessions, uploadID)
	s.tasks[uploadID] = t
	serr := s.save()
	s.mu.Unlock()
	if serr != nil {
		return nil, fmt.Errorf("persist task: %w", serr)
	}
	return t, nil
}

// AbortUpload abandons a chunked upload session and removes its part file.
// It is idempotent (a missing session is not an error).
func (s *Store) AbortUpload(uploadID string) error {
	s.mu.Lock()
	sess, ok := s.sessions[uploadID]
	if ok {
		delete(s.sessions, uploadID)
	}
	s.mu.Unlock()
	if !ok {
		return nil
	}
	_ = os.Remove(sess.partPath)
	return nil
}

// UploadStatus reports how many chunks of a session have arrived so far. It
// lets a client show progress and decide whether a complete is safe.
func (s *Store) UploadStatus(uploadID string) (received, total int, ok bool) {
	s.mu.Lock()
	sess, ok2 := s.sessions[uploadID]
	s.mu.Unlock()
	if !ok2 {
		return 0, 0, false
	}
	sess.mu.Lock()
	defer sess.mu.Unlock()
	n := 0
	for _, b := range sess.received {
		if b {
			n++
		}
	}
	return n, sess.TotalChunks, true
}

// gcUploads reaps sessions idle for longer than sessionTTL and deletes their
// .part files. It is called lazily from CreateUpload so no extra goroutine is
// needed. The two locks are phased (map under Store.mu, idle-check under the
// session mutex; file removal after Store.mu is released) to honour the
// "never hold Store.mu while taking a session lock for long" rule.
func (s *Store) gcUploads() {
	cutoff := time.Now().Add(-sessionTTL)
	var dead []*UploadSession
	s.mu.Lock()
	for id, sess := range s.sessions {
		sess.mu.Lock()
		idle := sess.lastSeen.Before(cutoff)
		sess.mu.Unlock()
		if idle {
			dead = append(dead, sess)
			delete(s.sessions, id)
		}
	}
	s.mu.Unlock()
	for _, sess := range dead {
		_ = os.Remove(sess.partPath)
	}
}

// sweepParts removes leftover .part files from interrupted chunked uploads
// after a restart (in-memory sessions don't survive). Called once from Open.
func (s *Store) sweepParts() {
	entries, err := os.ReadDir(s.uploadsDir)
	if err != nil {
		return
	}
	for _, e := range entries {
		if !e.Type().IsRegular() {
			continue
		}
		if strings.HasSuffix(e.Name(), ".part") {
			_ = os.Remove(filepath.Join(s.uploadsDir, e.Name()))
		}
	}
}