hamlab/internal/media/hls.go

package media

// Playing what the browser cannot play as it is: an HLS stream, re-encoded on demand by
// ffmpeg on the machine's own GPU where it has one it can use, on its CPU otherwise. A
// stream is one file at one codec and size, in a folder of its own under
// <data dir>/transcode/.
//
// A re-encode is seekable anywhere from the start. Its playlist is written up front, for
// the whole film: a fixed grid of 4 s segments, so the player's timeline is the film's.
// ffmpeg runs ahead from the top as fast as it can; a request for a segment it hasn't
// reached and won't reach soon — a seek ahead, or back into a stretch never made — ends that
// run and starts another at the segment asked for, which stops on its own where segments
// already exist. When a run finishes, the next gap is filled, until the whole file is on
// disk. Every run is made the same way (segArgs), so segments from different runs join.
//
// "Repackaged" (copy) can't cut on a fixed grid — its keyframes are the source's — so it is
// an event playlist that grows as ffmpeg writes it, seekable within what's done. Repackaging
// runs at disk speed, so that is all of it within seconds.
//
// Nothing is kept: a stream nobody has touched for a while is ended and its folder removed,
// and the folder is emptied when the daemon starts.

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Stream is one re-encode, in progress or finished.
type Stream struct {
	ID       string    `json:"id"`
	Path     string    `json:"path"`
	Codec    string    `json:"codec"`
	Quality  string    `json:"quality"`
	Engine   string    `json:"engine"`
	Width    int       `json:"width"`
	Height   int       `json:"height"`
	Kbps     int       `json:"kbps"`
	Source   Info      `json:"source"`
	Started  time.Time `json:"started"`
	Root     bool      `json:"root"`
	Seekable bool      `json:"seekable"` // seeking anywhere restarts the re-encode there

	key      string
	dir      string
	mu       sync.Mutex
	used     time.Time
	proc     *ffproc // the run going now (or the only one, for copy)
	finished bool
	failure  string
	gone     bool // forgotten by the Streamer

	// seekable streams only
	caps    Caps
	self    string
	recipe  Recipe
	n       int // segments in the playlist
	end     int // segments past this are not in the file after all (n until a run finds out)
	run     *segRun
	runMu   sync.Mutex // one run change at a time; held across ending one run and starting the next
	asked   time.Time  // when the request that last moved the run arrived: an older one never moves it back
	initOK  bool       // a later run's init segment was checked against the first's
	initBad bool
}

// segRun is one ffmpeg over segments from..to-1.
type segRun struct {
	from, to int
	proc     *ffproc
	stopped  bool // ended by us: what it wrote after that is not to be trusted
}

// StreamStatus is what the player asks for while the file is still going through.
type StreamStatus struct {
	Running  bool    `json:"running"`
	Done     bool    `json:"done"`
	Error    string  `json:"error,omitempty"`
	Progress float64 `json:"progress"` // seconds of output written (copy), or on disk (seekable)
	Speed    float64 `json:"speed"`    // × realtime
	Segments int     `json:"segments"`
	From     float64 `json:"from"` // seekable: where the run going now started, seconds
}

// Streamer keeps one machine's streams.
type Streamer struct {
	dir     string
	self    string
	limit   int
	mu      sync.Mutex
	streams map[string]*Stream
}

const streamIdle = 10 * time.Minute
const firstSegmentWait = 40 * time.Second

// segmentWait is how long a request for a segment waits for a run to write it.
const segmentWait = 45 * time.Second // under the hub proxy's 60 s wait for a node to answer

// lookahead: a segment this many past what the run has written is waited for, not
// restarted at. A few seconds of film, which any engine here writes faster than that.
const lookahead = 3

var streamFile = regexp.MustCompile(`^(index\.m3u8|init\.mp4|seg\d{5}\.m4s)$`)
var segName = regexp.MustCompile(`^seg(\d{5})\.m4s$`)

// NewStreamer empties the folder — a daemon restart orphans every stream — and starts the
// sweeper.
func NewStreamer(dataDir, self string, limit int) *Streamer {
	m := &Streamer{dir: filepath.Join(dataDir, "transcode"), self: self, limit: limit, streams: map[string]*Stream{}}
	_ = os.MkdirAll(m.dir, 0o755)
	_ = clearDir(m.dir)
	go func() {
		for range time.Tick(30 * time.Second) {
			m.Sweep(streamIdle)
		}
	}()
	return m
}

// Open starts a stream, or hands back the one already serving the same request. It
// returns once the first segment is on disk, so the player's first fetch is answered.
func (m *Streamer) Open(ctx context.Context, c Caps, root bool, path, codec, quality string, src Info) (*Stream, error) {
	if c.Error != "" {
		return nil, errors.New(c.Error)
	}
	if err := ValidChoice(codec, quality); err != nil {
		return nil, err
	}
	q, _ := qualityByID(quality)
	if src.Video == nil {
		return nil, &DecodeError{Msg: "no video stream"}
	}
	if src.Video.Width <= 0 || src.Video.Height <= 0 {
		return nil, &DecodeError{Msg: "the picture's size is unknown"}
	}
	key := fmt.Sprintf("%s|%d|%.0f|%s|%s|%v", path, src.Size, src.Duration, codec, quality, root)
	m.mu.Lock()
	for _, s := range m.streams {
		if s.key == key {
			if s.failed() { // not the stream to hand back: let it go and make a fresh one
				delete(m.streams, s.ID)
				_ = os.RemoveAll(s.dir)
				break
			}
			m.mu.Unlock()
			s.touch()
			return s, nil
		}
	}
	running := 0
	for _, s := range m.streams {
		if s.running() {
			running++
		}
	}
	if running >= m.limit {
		m.mu.Unlock()
		return nil, fmt.Errorf("%d videos are already being re-encoded on this machine; one must finish or be closed first", running)
	}
	s := &Stream{ID: newID(), Path: path, Codec: codec, Quality: quality, Source: src, Started: time.Now(), Root: root, key: key, used: time.Now(),
		caps: c, self: m.self}
	s.dir = filepath.Join(m.dir, s.ID)
	// A fixed grid needs the film's length; without one it can only grow as it goes.
	s.Seekable = codec != "copy" && src.Duration > 0
	if s.Seekable {
		s.n = segCount(src.Duration)
		s.end = s.n
	}
	m.streams[s.ID] = s
	m.mu.Unlock()

	s.Width, s.Height, s.Kbps = fit(q, src, codec)
	if codec == "copy" {
		s.Width, s.Height, s.Kbps = src.Video.Width, src.Video.Height, int(src.BitRate/1000)
	}
	rs := recipes(c, codec, src, s.Width, s.Height, s.Kbps)
	if len(rs) == 0 {
		m.forget(s)
		return nil, errors.New("nothing on this machine encodes " + codec)
	}
	var last string
	for _, r := range rs {
		if err := os.MkdirAll(s.dir, 0o755); err != nil { // the root helper's ffmpeg writes here too
			m.forget(s)
			return nil, err
		}
		_ = clearDir(s.dir)
		var args []string
		if s.Seekable {
			args = segArgs(s.dir, path, src, r, 0, s.n, s.n, "init.mp4")
		} else {
			args = hlsArgs(s.dir, path, src, r, codec == "copy")
		}
		p, err := startFF(c, m.self, root, args)
		if err != nil {
			last = r.Engine + ": " + err.Error()
			continue
		}
		s.mu.Lock()
		s.proc, s.Engine, s.recipe = p, r.Engine, r
		if s.Seekable {
			s.run = &segRun{from: 0, to: s.n, proc: p}
		}
		s.mu.Unlock()
		msg, ok := s.waitFirstSegment(ctx)
		if ok {
			if s.Seekable {
				go s.watchRun(s.run)
			} else {
				go s.watch()
			}
			log.Printf("media: stream %s: %s → %s %s %dx%d %dk on %s", s.ID, filepath.Base(path), codec, quality, s.Width, s.Height, s.Kbps, r.Engine)
			return s, nil
		}
		p.end()
		if ctx.Err() != nil || s.closed() { // the caller left, or somebody closed it meanwhile: not the recipe's fault
			m.forget(s)
			return nil, errors.New("cancelled")
		}
		last = r.Engine + ": " + msg
		log.Printf("media: stream %s on %s: %s", filepath.Base(path), r.Engine, msg)
	}
	m.forget(s)
	return nil, &DecodeError{Msg: last}
}

// waitFirstSegment watches for the first segment, or ffmpeg giving up.
func (s *Stream) waitFirstSegment(ctx context.Context) (string, bool) {
	deadline := time.Now().Add(firstSegmentWait)
	playlist := filepath.Join(s.dir, "index.m3u8")
	for {
		if s.Seekable && exists(s.segPath(0)) || !s.Seekable && hasSegments(playlist) {
			return "", true
		}
		if done, err := s.proc.exited(); done {
			msg := s.proc.lastWords()
			if msg == "" && err != nil {
				msg = err.Error()
			}
			if msg == "" {
				msg = "ffmpeg ended before writing a segment"
			}
			return msg, false
		}
		if time.Now().After(deadline) {
			return "no segment after " + firstSegmentWait.String(), false
		}
		select {
		case <-ctx.Done():
			return ctx.Err().Error(), false
		case <-time.After(250 * time.Millisecond):
		}
	}
}

// watch records how the encode ended (copy streams: one run, start to end).
func (s *Stream) watch() {
	<-s.proc.done
	s.mu.Lock()
	s.finished = true
	if s.proc.err != nil {
		s.failure = s.proc.lastWords()
		if s.failure == "" {
			s.failure = s.proc.err.Error()
		}
	}
	s.mu.Unlock()
	if s.failure != "" {
		log.Printf("media: stream %s ended early: %s", s.ID, s.failure)
	}
}

// ---- seekable streams ----

func (s *Stream) segPath(i int) string { return filepath.Join(s.dir, fmt.Sprintf("seg%05d.m4s", i)) }

// written is the last segment of a run on disk, counting on from where it started.
func (s *Stream) written(r *segRun) int {
	i := r.from
	for i < r.to && exists(s.segPath(i)) {
		i++
	}
	return i - 1
}

// watchRun follows one run to its end. A run that ends by itself — the end of the film, or
// the segments it was to stop short of — is followed by the next gap. One that failed fails
// the stream; one we ended is the business of whoever ended it.
func (s *Stream) watchRun(r *segRun) {
	<-r.proc.done
	s.mu.Lock()
	if s.gone || r.stopped || s.run != r {
		s.mu.Unlock()
		return
	}
	if r.proc.err != nil {
		s.failure = r.proc.lastWords()
		if s.failure == "" {
			s.failure = r.proc.err.Error()
		}
		s.finished = true
		s.mu.Unlock()
		log.Printf("media: stream %s: run from segment %d failed: %s", s.ID, r.from, s.failure)
		return
	}
	if w := s.written(r); w+1 < r.to && r.to == s.n {
		// ran to the end of the file before the end of the grid: the picture is shorter than
		// the container says, and those last segments will never exist
		s.end = w + 1
		log.Printf("media: stream %s: the film ends at segment %d of %d", s.ID, s.end, s.n)
	}
	s.checkInit(r)
	s.mu.Unlock()
	s.fillNext(r.to)
}

// checkInit compares a later run's init segment with the first run's, once. They come out
// identical (same encoder, same settings); if one day they don't, the log says so.
func (s *Stream) checkInit(r *segRun) {
	if r.from == 0 || s.initOK {
		return
	}
	a, err1 := os.ReadFile(filepath.Join(s.dir, "init.mp4"))
	b, err2 := os.ReadFile(filepath.Join(s.dir, "init-run.mp4"))
	if err1 != nil || err2 != nil {
		return
	}
	s.initOK = true
	if !bytes.Equal(a, b) && !s.initBad {
		s.initBad = true
		log.Printf("media: stream %s: a restarted run's init segment differs from the first (%s); its segments may not join", s.ID, s.Engine)
	}
}

// missingFrom is the first segment not on disk at or after i (wrapping round to the start),
// or -1 when the whole film is there.
func (s *Stream) missingFrom(i int) int {
	end := s.end
	for k := 0; k < end; k++ {
		j := (i + k) % end
		if !exists(s.segPath(j)) {
			return j
		}
	}
	return -1
}

// fillNext starts a run at the next gap after `after`; when there is none, the stream is done.
func (s *Stream) fillNext(after int) {
	s.runMu.Lock()
	defer s.runMu.Unlock()
	s.mu.Lock()
	if s.gone || s.failure != "" || (s.run != nil && !s.run.done()) {
		s.mu.Unlock()
		return
	}
	end := s.end
	s.mu.Unlock()
	if after >= end {
		after = 0
	}
	next := s.missingFrom(after)
	if next < 0 {
		s.mu.Lock()
		s.finished, s.proc = true, nil
		s.mu.Unlock()
		log.Printf("media: stream %s: all %d segments on disk", s.ID, end)
		return
	}
	s.startRun(next)
}

func (r *segRun) done() bool { d, _ := r.proc.exited(); return d }

// startRun ends the run going now and starts one at segment from, stopping short of the next
// segment already on disk. The run it ends is asked to stop and waited for, and whatever it
// finished writing after being asked is removed: ffmpeg closes the segment it was in the
// middle of on its way out, short. Called with runMu held, mu not.
func (s *Stream) startRun(from int) {
	s.mu.Lock()
	old := s.run
	end := s.end
	s.mu.Unlock()
	if old != nil && !old.done() {
		had := map[int]bool{}
		for i := old.from; i < old.to; i++ {
			had[i] = exists(s.segPath(i))
		}
		s.mu.Lock()
		old.stopped = true
		s.mu.Unlock()
		old.proc.end()
		for i := old.from; i < old.to; i++ {
			if !had[i] && exists(s.segPath(i)) {
				_ = os.Remove(s.segPath(i))
			}
		}
	}
	to := from + 1
	for to < end && !exists(s.segPath(to)) {
		to++
	}
	s.mu.Lock()
	if s.gone {
		s.mu.Unlock()
		return
	}
	s.mu.Unlock()
	init := "init-run.mp4"
	if !exists(filepath.Join(s.dir, "init.mp4")) {
		init = "init.mp4"
	}
	p, err := startFF(s.caps, s.self, s.Root, segArgs(s.dir, s.Path, s.Source, s.recipe, from, to, s.n, init))
	s.mu.Lock()
	defer s.mu.Unlock()
	if err != nil {
		s.failure, s.finished = "restarting at "+strconv.Itoa(from*hlsSeconds)+" s: "+err.Error(), true
		return
	}
	r := &segRun{from: from, to: to, proc: p}
	s.run, s.proc, s.finished = r, p, false
	go s.watchRun(r)
	log.Printf("media: stream %s: run from %d s to %d s", s.ID, from*hlsSeconds, to*hlsSeconds)
}

// Playlist is the whole film's, written before most of it exists.
func (s *Stream) Playlist() []byte {
	var b strings.Builder
	b.WriteString("#EXTM3U\n#EXT-X-VERSION:7\n")
	fmt.Fprintf(&b, "#EXT-X-TARGETDURATION:%d\n", hlsSeconds)
	b.WriteString("#EXT-X-MEDIA-SEQUENCE:0\n#EXT-X-PLAYLIST-TYPE:VOD\n#EXT-X-INDEPENDENT-SEGMENTS\n#EXT-X-MAP:URI=\"init.mp4\"\n")
	for i := 0; i < s.n; i++ {
		d := float64(hlsSeconds)
		if i == s.n-1 {
			d = s.Source.Duration - float64(i*hlsSeconds)
		}
		fmt.Fprintf(&b, "#EXTINF:%.6f,\nseg%05d.m4s\n", d, i)
	}
	b.WriteString("#EXT-X-ENDLIST\n")
	return []byte(b.String())
}

// ErrNotInFilm is a segment the playlist lists but the film turned out not to reach.
var ErrNotInFilm = errors.New("past the end of the film")

// Segment is the path of segment i once it is on disk. When the run going now will write it
// soon, it waits; otherwise this is a seek, and the run moves there — unless a newer request
// has moved it since this one arrived.
func (s *Stream) Segment(ctx context.Context, i int) (string, error) {
	asked := time.Now()
	deadline := asked.Add(segmentWait)
	p := s.segPath(i)
	for {
		s.touch()
		if exists(p) {
			return p, nil
		}
		s.mu.Lock()
		switch {
		case s.gone:
			s.mu.Unlock()
			return "", errors.New("the stream was closed")
		case i >= s.end:
			s.mu.Unlock()
			return "", ErrNotInFilm
		case s.failure != "":
			f := s.failure
			s.mu.Unlock()
			return "", errors.New(f)
		}
		r := s.run
		coming := r != nil && !r.done() && r.from <= i && i < r.to && i <= s.written(r)+lookahead
		move := !coming && !asked.Before(s.asked)
		if move {
			s.asked = asked
		}
		s.mu.Unlock()
		if move {
			s.runMu.Lock()
			// still the newest, and still not on its way? (a run may have moved for a
			// newer request while this one waited for the lock)
			s.mu.Lock()
			r = s.run
			coming = r != nil && !r.done() && r.from <= i && i < r.to && i <= s.written(r)+lookahead
			newest := !asked.Before(s.asked)
			s.mu.Unlock()
			if !coming && newest && !exists(p) {
				log.Printf("media: stream %s: seek to %d s", s.ID, i*hlsSeconds)
				s.startRun(i)
			}
			s.runMu.Unlock()
			continue
		}
		if time.Now().After(deadline) {
			return "", fmt.Errorf("segment %d didn't arrive in %s", i, segmentWait)
		}
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-time.After(150 * time.Millisecond):
		}
	}
}

// SegmentIndex reads the number out of a segment's name.
func SegmentIndex(name string) (int, bool) {
	m := segName.FindStringSubmatch(name)
	if m == nil {
		return 0, false
	}
	i, err := strconv.Atoi(m[1])
	return i, err == nil
}

// ---- both kinds ----

func (m *Streamer) Find(id string) (*Stream, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	s, ok := m.streams[id]
	return s, ok
}

// Close ends a stream and removes its folder.
func (m *Streamer) Close(id string) {
	m.mu.Lock()
	s, ok := m.streams[id]
	m.mu.Unlock()
	if !ok {
		return
	}
	s.mu.Lock()
	s.gone = true // a run ending now is ours, not a gap to fill
	p := s.proc
	if s.run != nil {
		s.run.stopped = true
	}
	s.mu.Unlock()
	if p != nil {
		p.end()
	}
	m.forget(s)
}

func (m *Streamer) forget(s *Stream) {
	m.mu.Lock()
	delete(m.streams, s.ID)
	m.mu.Unlock()
	s.mu.Lock()
	s.gone = true
	s.mu.Unlock()
	_ = os.RemoveAll(s.dir)
}

func (s *Stream) closed() bool { s.mu.Lock(); defer s.mu.Unlock(); return s.gone }

func (s *Stream) failed() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.finished && s.failure != ""
}

// Sweep closes what nobody has touched for a while.
func (m *Streamer) Sweep(idle time.Duration) {
	m.mu.Lock()
	var old []string
	for id, s := range m.streams {
		s.mu.Lock()
		if time.Since(s.used) > idle {
			old = append(old, id)
		}
		s.mu.Unlock()
	}
	m.mu.Unlock()
	for _, id := range old {
		log.Printf("media: stream %s idle, closed", id)
		m.Close(id)
	}
}

// File is the on-disk path of one of the stream's files — the playlist (copy streams; a
// seekable one's is Playlist), the init segment or a media segment — and nothing else.
func (s *Stream) File(name string) (string, bool) {
	if !streamFile.MatchString(name) {
		return "", false
	}
	s.touch()
	return filepath.Join(s.dir, name), true
}

func (s *Stream) touch() { s.mu.Lock(); s.used = time.Now(); s.mu.Unlock() }

func (s *Stream) running() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.proc != nil && !s.finished
}

// Status is read from ffmpeg's own progress file and what is on disk.
func (s *Stream) Status() StreamStatus {
	s.mu.Lock()
	st := StreamStatus{Running: s.proc != nil && !s.finished, Done: s.finished && s.failure == "", Error: s.failure}
	seekable, end, run := s.Seekable, s.end, s.run
	s.mu.Unlock()
	b, err := os.ReadFile(filepath.Join(s.dir, "progress"))
	if err == nil {
		st.Progress, st.Speed = parseProgress(b)
	}
	if !seekable {
		if b, err := os.ReadFile(filepath.Join(s.dir, "index.m3u8")); err == nil {
			st.Segments = strings.Count(string(b), "#EXTINF")
		}
		return st
	}
	for i := 0; i < end; i++ {
		if exists(s.segPath(i)) {
			st.Segments++
		}
	}
	st.Progress = float64(st.Segments * hlsSeconds)
	if st.Progress > s.Source.Duration {
		st.Progress = s.Source.Duration
	}
	if run != nil && st.Running {
		st.From = float64(run.from * hlsSeconds)
	}
	return st
}

// parseProgress takes the last out_time_us and speed from ffmpeg's -progress output.
func parseProgress(b []byte) (seconds, speed float64) {
	sc := bufio.NewScanner(strings.NewReader(string(b)))
	for sc.Scan() {
		k, v, ok := strings.Cut(sc.Text(), "=")
		if !ok {
			continue
		}
		v = strings.TrimSpace(v)
		switch k {
		case "out_time_us":
			if us, err := strconv.ParseInt(v, 10, 64); err == nil && us >= 0 {
				seconds = float64(us) / 1e6
			}
		case "speed":
			if x, err := strconv.ParseFloat(strings.TrimSuffix(v, "x"), 64); err == nil {
				speed = x
			}
		}
	}
	return seconds, speed
}

func newID() string {
	b := make([]byte, 8)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}