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/: ffmpeg writes an event playlist and fragmented-MP4 segments there
// as fast as it can, and the player reads them as they land, so the range it can seek in
// grows until the whole file is through — on a GPU, several times faster than watching.
// 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"
	"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"`

	key      string
	dir      string
	mu       sync.Mutex
	used     time.Time
	proc     *ffproc
	finished bool
	failure  string
	gone     bool // forgotten by the Streamer
}

// 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
	Speed    float64 `json:"speed"`    // × realtime
	Segments int     `json:"segments"`
}

// 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

var streamFile = regexp.MustCompile(`^(index\.m3u8|init\.mp4|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()}
	s.dir = filepath.Join(m.dir, s.ID)
	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)
		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 = p, r.Engine
		s.mu.Unlock()
		msg, ok := s.waitFirstSegment(ctx)
		if ok {
			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 playlist's 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 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.
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)
	}
}

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()
	p := s.proc
	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, 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 the playlist.
func (s *Stream) Status() StreamStatus {
	s.mu.Lock()
	st := StreamStatus{Running: s.proc != nil && !s.finished, Done: s.finished && s.failure == "", Error: s.failure}
	s.mu.Unlock()
	if b, err := os.ReadFile(filepath.Join(s.dir, "progress")); err == nil {
		st.Progress, st.Speed = parseProgress(b)
	}
	if b, err := os.ReadFile(filepath.Join(s.dir, "index.m3u8")); err == nil {
		st.Segments = strings.Count(string(b), "#EXTINF")
	}
	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)
}