live/edge-agent/agent/mediamtx.go

package agent

import (
	"bufio"
	"embed"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
)

type MediaMTXManager struct {
	cfg        Config
	fs         embed.FS
	binPath    string
	configPath string
	cmd        *exec.Cmd
	mu         sync.RWMutex
	metrics    MediaMetrics
}

type MediaMetrics struct {
	ViewerCount int
	FrameDrops  int
	PacketLoss  int
	Available   bool
}

func NewMediaMTXManager(cfg Config, fs embed.FS) *MediaMTXManager {
	return &MediaMTXManager{
		cfg:        cfg,
		fs:         fs,
		binPath:    filepath.Join(cfg.DataDir, "mediamtx"),
		configPath: filepath.Join(cfg.DataDir, "mediamtx.yml"),
	}
}

func (m *MediaMTXManager) Extract() error {
	if _, err := os.Stat(m.binPath); err == nil {
		return nil
	}

	srcPath := "mediamtx/"
	if runtime.GOOS == "windows" {
		srcPath += "mediamtx.exe"
	} else {
		srcPath += "mediamtx"
	}

	entries, err := m.fs.ReadDir("mediamtx")
	if err != nil {
		return fmt.Errorf("read embedded dir: %w", err)
	}

	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()
		if runtime.GOOS == "windows" && !strings.Contains(name, ".exe") {
			continue
		}
		if runtime.GOOS != "windows" && strings.Contains(name, ".exe") {
			continue
		}

		src, err := m.fs.Open("mediamtx/" + name)
		if err != nil {
			return fmt.Errorf("open embedded %s: %w", name, err)
		}
		defer src.Close()

		destName := "mediamtx"
		if runtime.GOOS == "windows" {
			destName = "mediamtx.exe"
		}
		destPath := filepath.Join(m.cfg.DataDir, destName)
		dest, err := os.OpenFile(destPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
		if err != nil {
			src.Close()
			return fmt.Errorf("create %s: %w", destPath, err)
		}

		if _, err := io.Copy(dest, src); err != nil {
			dest.Close()
			src.Close()
			return fmt.Errorf("copy: %w", err)
		}
		dest.Close()
		src.Close()
		m.binPath = destPath
		break
	}

	return nil
}

func (m *MediaMTXManager) WriteConfig() error {
	config := fmt.Sprintf(`logLevel: info
logDestinations: [stdout]

rtmp: no
rtsp: no

hls: yes
hlsAddress: :%d

api: yes
apiAddress: :%d

paths:
  all:
    source: rtmp://%s:1935/live/$MTX_PATH
    sourceOnDemand: yes
`,
		m.cfg.HlsPort,
		m.cfg.ApiPort,
		strings.TrimPrefix(strings.TrimPrefix(m.cfg.OriginURL, "http://"), "https://"),
	)

	return os.WriteFile(m.configPath, []byte(config), 0644)
}

func (m *MediaMTXManager) Start() error {
	m.cmd = exec.Command(m.binPath, m.configPath)
	stdout, err := m.cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("capture mediamtx stdout: %w", err)
	}
	stderr, err := m.cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("capture mediamtx stderr: %w", err)
	}
	if err := m.cmd.Start(); err != nil {
		return fmt.Errorf("start mediamtx: %w", err)
	}
	go m.parseLogs(stdout)
	go m.parseLogs(stderr)
	return nil
}

func (m *MediaMTXManager) parseLogs(reader io.Reader) {
	scanner := bufio.NewScanner(reader)
	for scanner.Scan() {
		line := scanner.Text()
		fmt.Fprintln(os.Stdout, line)
		lower := strings.ToLower(line)
		m.mu.Lock()
		if strings.Contains(lower, "reader opened") || strings.Contains(lower, "hls client connected") {
			m.metrics.Available = true
			m.metrics.ViewerCount++
		}
		if strings.Contains(lower, "reader closed") || strings.Contains(lower, "hls client disconnected") {
			m.metrics.Available = true
			if m.metrics.ViewerCount > 0 {
				m.metrics.ViewerCount--
			}
		}
		if strings.Contains(lower, "frame drop") || strings.Contains(lower, "dropped frame") {
			m.metrics.Available = true
			m.metrics.FrameDrops++
		}
		if strings.Contains(lower, "packet loss") || strings.Contains(lower, "lost packet") {
			m.metrics.Available = true
			m.metrics.PacketLoss++
		}
		m.mu.Unlock()
	}
}

func (m *MediaMTXManager) Metrics() MediaMetrics {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.metrics
}

func (m *MediaMTXManager) Stop() {
	if m.cmd != nil && m.cmd.Process != nil {
		m.cmd.Process.Kill()
		m.cmd.Wait()
	}
}