Vellum/internal/transcode/session_test.go

package transcode

import (
	"context"
	"errors"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// testManager builds a Manager that never touches real hardware or ffmpeg: the
// capability and the command are injected, the same way detect injects its
// probe. The stub command blocks until its context is cancelled, standing in
// for ffmpeg.
func testManager(t *testing.T) (*Manager, func() int) {
	t.Helper()
	var spawns atomic.Int32
	m := New()
	m.hw = func() Hardware {
		return Hardware{Encoder: EncoderVAAPI, Args: []string{"-vaapi_device", "/dev/dri/renderD128", "-c:v", "h264_vaapi"}}
	}
	m.newCmd = func(ctx context.Context, args []string) *exec.Cmd {
		spawns.Add(1)
		return exec.CommandContext(ctx, "sleep", "3600")
	}
	t.Cleanup(m.Shutdown)
	return m, func() int { return int(spawns.Load()) }
}

// TestStartAndReuse is the reason Start exists: two requests for the same item
// must not start two encoders.
func TestStartAndReuse(t *testing.T) {
	m, spawns := testManager(t)
	cacheDir := t.TempDir()

	a, err := m.Start(7, "/media/movie.mkv", cacheDir)
	if err != nil {
		t.Fatalf("Start: %v", err)
	}
	b, err := m.Start(7, "/media/movie.mkv", cacheDir)
	if err != nil {
		t.Fatalf("Start (reuse): %v", err)
	}
	if a != b {
		t.Errorf("second Start returned a different session: %v then %v", a.Dir(), b.Dir())
	}
	if got := spawns(); got != 1 {
		t.Errorf("spawned %d ffmpegs, want 1", got)
	}
	if want := filepath.Join(a.Dir(), "playlist.m3u8"); a.PlaylistPath() != want {
		t.Errorf("PlaylistPath = %q, want %q", a.PlaylistPath(), want)
	}
	if !strings.HasPrefix(a.Dir(), filepath.Join(cacheDir, "transcode")) {
		t.Errorf("Dir = %q, want it under %q", a.Dir(), filepath.Join(cacheDir, "transcode"))
	}
}

// TestStartArgs pins the ffmpeg command line the task requires.
func TestStartArgs(t *testing.T) {
	m := New()
	t.Cleanup(m.Shutdown)
	m.hw = func() Hardware {
		return Hardware{Encoder: EncoderNVENC, Args: []string{"-c:v", "h264_nvenc"}}
	}
	var got []string
	m.newCmd = func(ctx context.Context, args []string) *exec.Cmd {
		got = append([]string(nil), args...)
		return exec.CommandContext(ctx, "sleep", "3600")
	}

	s, err := m.Start(9, "/m/in.mkv", t.TempDir())
	if err != nil {
		t.Fatalf("Start: %v", err)
	}
	// The hardware args must come AFTER -i. Before it, ffmpeg treats them as
	// INPUT options and "-c:v h264_nvenc" is looked up as a decoder: on a real
	// NVENC host that fails with "Unknown decoder 'h264_nvenc'" and produces no
	// segments at all. This ordering is load-bearing, not cosmetic.
	want := []string{
		"-i", "/m/in.mkv",
		"-c:v", "h264_nvenc",
		"-c:a", "aac", "-b:a", "128k",
		"-f", "hls",
		"-hls_time", "6",
		"-hls_playlist_type", "event",
		"-hls_segment_filename", filepath.Join(s.Dir(), "seg%05d.ts"),
		filepath.Join(s.Dir(), "playlist.m3u8"),
	}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("args = %q, want %q", got, want)
	}

	// State the invariant directly as well, so a future reordering fails with a
	// message that explains why rather than only showing two arg lists.
	in, enc := indexOf(got, "-i"), indexOf(got, "-c:v")
	if in < 0 || enc < 0 || enc < in {
		t.Errorf("encoder at %d, -i at %d: the hardware args must follow the input", enc, in)
	}
}

func indexOf(args []string, want string) int {
	for i, a := range args {
		if a == want {
			return i
		}
	}
	return -1
}

// TestIdleReap proves a session nobody touches is stopped, so an abandoned
// stream does not burn the GPU forever.
func TestIdleReap(t *testing.T) {
	m, spawns := testManager(t)
	m.idle = time.Minute
	cacheDir := t.TempDir()

	a, err := m.Start(1, "/m/in.mkv", cacheDir)
	if err != nil {
		t.Fatalf("Start: %v", err)
	}

	// Not yet idle: the session survives and is still reused.
	m.reap(m.now())
	if b, _ := m.Start(1, "/m/in.mkv", cacheDir); b != a {
		t.Fatal("session reaped before it was idle")
	}

	// Idle: reaped, so the next request starts a fresh encoder.
	m.reap(m.now().Add(time.Hour))
	c, err := m.Start(1, "/m/in.mkv", cacheDir)
	if err != nil {
		t.Fatalf("Start after reap: %v", err)
	}
	if c == a {
		t.Error("idle session was reused after being reaped")
	}
	if got := spawns(); got != 2 {
		t.Errorf("spawned %d ffmpegs, want 2", got)
	}
}

// TestAtCapacity proves a request past the cap is refused, not queued.
func TestAtCapacity(t *testing.T) {
	m, _ := testManager(t)
	m.max = 1
	cacheDir := t.TempDir()

	if _, err := m.Start(1, "/m/a.mkv", cacheDir); err != nil {
		t.Fatalf("Start: %v", err)
	}
	_, err := m.Start(2, "/m/b.mkv", cacheDir)
	if !errors.Is(err, ErrAtCapacity) {
		t.Errorf("Start past cap: got %v, want ErrAtCapacity", err)
	}
}

// TestStartNoCapability proves a GPU-less host gets a clear refusal.
func TestStartNoCapability(t *testing.T) {
	m := New()
	t.Cleanup(m.Shutdown)
	m.hw = func() Hardware { return Hardware{} }

	_, err := m.Start(1, "/m/in.mkv", t.TempDir())
	if !errors.Is(err, ErrUnavailable) {
		t.Errorf("Start with no capability: got %v, want ErrUnavailable", err)
	}
}

// TestShutdownStopsSessions proves shutdown tears down running encoders and
// refuses new work afterwards.
func TestShutdownStopsSessions(t *testing.T) {
	m := New()
	m.hw = func() Hardware { return Hardware{Encoder: EncoderNVENC, Args: []string{"-c:v", "h264_nvenc"}} }
	m.newCmd = func(ctx context.Context, args []string) *exec.Cmd {
		return exec.CommandContext(ctx, "sleep", "3600")
	}
	if _, err := m.Start(1, "/m/a.mkv", t.TempDir()); err != nil {
		t.Fatalf("Start: %v", err)
	}

	m.Shutdown()

	if _, err := m.Start(2, "/m/b.mkv", t.TempDir()); !errors.Is(err, ErrShutdown) {
		t.Errorf("Start after Shutdown: got %v, want ErrShutdown", err)
	}
}