Vellum/internal/transcode/session.go

package transcode

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"sync"
	"time"
)

// Errors Start returns. They are sentinel values so callers can tell a missing
// capability from a full server from a preference to retry later.
var (
	// ErrUnavailable means no hardware encoder passed detection's probe.
	ErrUnavailable = errors.New("transcode: no hardware encoder available")
	// ErrAtCapacity means the concurrent-session cap is reached. The request
	// is refused rather than queued: a queued transcode looks like a hang.
	ErrAtCapacity = errors.New("transcode: concurrent session cap reached")
	// ErrShutdown means the manager has been shut down.
	ErrShutdown = errors.New("transcode: manager shut down")
)

const (
	// maxSessions caps concurrent transcodes. A GPU encodes one or two
	// streams, not ten.
	maxSessions = 2
	// idleTimeout is how long a session may go untouched before its encoder
	// is stopped, so an abandoned stream does not pin the GPU forever.
	idleTimeout = 3 * time.Minute
	// reapEvery is how often the idle reaper runs.
	reapEvery = 30 * time.Second
)

// Session is one running transcode: ffmpeg writing an HLS rendition of a single
// item into its own directory. It is reused across requests for the same item,
// so a returned Session may have other holders. Close releases the caller's
// hold; the encoder is stopped once no holder remains, or when the session goes
// idle.
type Session interface {
	// PlaylistPath is the HLS playlist ffmpeg writes and a player fetches.
	PlaylistPath() string
	// Dir is the directory holding the playlist and its segments.
	Dir() string
	// Close releases this caller's hold on the session. Returns nil if it was
	// already released.
	Close() error
}

// Manager starts, reuses, reaps and caps transcode sessions.
//
// The same item gets one session while it is alive, so the two range requests a
// player makes for one file do not start two encoders. A session nobody has
// touched for idleTimeout is stopped. At most max sessions run at once; a
// request past the cap is an error, never a queue.
type Manager struct {
	// hw, newCmd and now are the impure edges, injected so a test can run a
	// session with no hardware and no ffmpeg, the same way detect injects its
	// probe.
	hw     func() Hardware
	newCmd func(ctx context.Context, args []string) *exec.Cmd
	now    func() time.Time

	// max caps concurrent sessions; idle is how long an untouched session is
	// kept before it is reaped.
	max  int
	idle time.Duration

	mu       sync.Mutex
	sessions map[int64]*session
	seq      int64
	closed   bool

	stop     chan struct{}
	stopOnce sync.Once
	wg       sync.WaitGroup
}

// New builds a Manager that uses the host's detected capability and the real
// ffmpeg binary, and starts its idle reaper.
func New() *Manager {
	m := &Manager{
		hw:       Detect,
		newCmd:   newFFmpeg,
		now:      time.Now,
		max:      maxSessions,
		idle:     idleTimeout,
		sessions: map[int64]*session{},
		stop:     make(chan struct{}),
	}
	m.wg.Add(1)
	go m.reaper()
	return m
}

// newFFmpeg is the production command builder.
func newFFmpeg(ctx context.Context, args []string) *exec.Cmd {
	return exec.CommandContext(ctx, "ffmpeg", args...)
}

// Start spawns ffmpeg to package inputPath as HLS under cacheDir/transcode/ and
// returns the session. Asking for the same item again while its session is alive
// reuses it. It returns ErrUnavailable, ErrAtCapacity or ErrShutdown instead of
// queueing.
func (m *Manager) Start(itemID int64, inputPath, cacheDir string) (Session, error) {
	hw := m.hw()
	if !hw.Available() {
		return nil, fmt.Errorf("%w: no VAAPI, NVENC or QSV encoder passed its probe", ErrUnavailable)
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	if m.closed {
		return nil, ErrShutdown
	}
	if s, ok := m.sessions[itemID]; ok {
		s.refs++
		s.last = m.now()
		return s, nil
	}
	if len(m.sessions) >= m.max {
		return nil, fmt.Errorf("%w (%d running)", ErrAtCapacity, m.max)
	}

	m.seq++
	dir := filepath.Join(cacheDir, "transcode", strconv.FormatInt(itemID, 10)+"-"+strconv.FormatInt(m.seq, 10))
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, fmt.Errorf("transcode: create session dir: %w", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	cmd := m.newCmd(ctx, sessionArgs(hw, inputPath, dir))
	if err := cmd.Start(); err != nil {
		cancel()
		_ = os.RemoveAll(dir)
		return nil, fmt.Errorf("transcode: start ffmpeg: %w", err)
	}
	// Reap the child when it exits; without a Wait it would linger as a zombie.
	go func() { _ = cmd.Wait() }()

	s := &session{m: m, itemID: itemID, dir: dir, cancel: cancel, refs: 1, last: m.now()}
	m.sessions[itemID] = s
	return s, nil
}

// sessionArgs is the ffmpeg command line: the encoder's hardware arguments,
// then the input, then AAC audio and event HLS packaging into dir.
func sessionArgs(hw Hardware, inputPath, dir string) []string {
	// The hardware arguments go AFTER -i. Placed before it, ffmpeg reads them as
	// INPUT options and "-c:v h264_nvenc" is looked up as a decoder, failing
	// with "Unknown decoder 'h264_nvenc'". Verified against a real NVENC host:
	// before -i produces no segments at all, after -i produces them.
	args := make([]string, 0, len(hw.Args)+14)
	args = append(args, "-i", inputPath)
	args = append(args, hw.Args...)
	return append(args,
		"-c:a", "aac", "-b:a", "128k",
		"-f", "hls",
		"-hls_time", "6",
		"-hls_playlist_type", "event",
		"-hls_segment_filename", filepath.Join(dir, "seg%05d.ts"),
		filepath.Join(dir, "playlist.m3u8"),
	)
}

// Shutdown stops every session and the idle reaper. Call it once at process
// exit; it is safe to call more than once.
func (m *Manager) Shutdown() {
	m.stopOnce.Do(func() { close(m.stop) })
	m.wg.Wait()

	m.mu.Lock()
	defer m.mu.Unlock()
	m.closed = true
	for _, s := range m.sessions {
		m.stopLocked(s)
	}
}

// reap stops every session untouched for at least idle. It is separate from the
// ticker so a test can drive it directly.
func (m *Manager) reap(now time.Time) {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, s := range m.sessions {
		if now.Sub(s.last) >= m.idle {
			m.stopLocked(s)
		}
	}
}

func (m *Manager) reaper() {
	defer m.wg.Done()
	t := time.NewTicker(reapEvery)
	defer t.Stop()
	for {
		select {
		case <-m.stop:
			return
		case <-t.C:
			m.reap(m.now())
		}
	}
}

// stopLocked kills a session's encoder and drops it from the manager. Callers
// must hold m.mu.
func (m *Manager) stopLocked(s *session) {
	if s.stopped {
		return
	}
	s.stopped = true
	s.cancel()
	if cur, ok := m.sessions[s.itemID]; ok && cur == s {
		delete(m.sessions, s.itemID)
	}
}

// session is Manager's concrete Session.
//
// shortcut: a stopped session leaves its segment directory on disk; the cache
// directory is expected to be pruned out of band. Add eviction here if that
// cache is not otherwise bounded.
type session struct {
	m      *Manager
	itemID int64
	dir    string
	cancel context.CancelFunc
	refs   int
	last   time.Time
	// stopped guards against double Close and against Close racing the reaper.
	stopped bool
}

func (s *session) Dir() string          { return s.dir }
func (s *session) PlaylistPath() string { return filepath.Join(s.dir, "playlist.m3u8") }

// Close releases one hold. The encoder is stopped once the last hold is gone.
func (s *session) Close() error {
	s.m.mu.Lock()
	defer s.m.mu.Unlock()
	if s.stopped {
		return nil
	}
	s.refs--
	if s.refs > 0 {
		return nil
	}
	s.m.stopLocked(s)
	return nil
}