YCMediaKit/internal/pkg/server/ffmpegServer/toHlsCmd.go

package ffmpegServer

import (
	"errors"
	"fmt"
	"os/exec"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"ycmediakit/internal/pkg/unit/keepalive"
	"ycmediakit/internal/pkg/util"

	ffmpegCmd "github.com/u2takey/ffmpeg-go"
	"go.uber.org/zap"
)

var (
	log *zap.Logger
)

func init() {
	log = zap.L()
}

type ToHlsServer struct {
	Cfg        *ToHlsConfig
	m          sync.Mutex
	cmdMap     map[string]*ToHlsInfo
	Supervisor *keepalive.Supervisor
}

type ToHlsInfo struct {
	ToHlsConfig
	url        string
	streamPath string
	cmd        *exec.Cmd
}

func NewToHlsServer(cfg *ToHlsConfig) *ToHlsServer {
	s := &ToHlsServer{
		Cfg:    cfg,
		cmdMap: make(map[string]*ToHlsInfo),
	}
	if s.Cfg.Keepalive.Enable {
		s.Supervisor = keepalive.NewSupervisor(&s.Cfg.Keepalive)
		handlePath = "/" + cfg.HlsRoot
		s.Supervisor.Register(handlePath, s)
	}
	return s
}

func (s *ToHlsServer) Add(url string, streamPath string, timeoutMs int, curHlsCfg *ToHlsConfig) bool {
	s.m.Lock()
	defer s.m.Unlock()
	_, ok := s.cmdMap[streamPath]
	if ok {
		return true
	}
	// merge curHlsCfg, Cfg
	cmd := s.BuildToHlsCmd(url, streamPath, timeoutMs, curHlsCfg)
	info := &ToHlsInfo{
		ToHlsConfig: *curHlsCfg,
		url:         url,
		streamPath:  streamPath,
		cmd:         cmd,
	}
	s.cmdMap[streamPath] = info
	return true
}

func (s *ToHlsServer) Exists(streamPath string) bool {
	s.m.Lock()
	defer s.m.Unlock()
	_, ok := s.cmdMap[streamPath]
	return ok
}

func (s *ToHlsServer) Delete(streamPath string) bool {
	s.m.Lock()
	defer s.m.Unlock()
	delete(s.cmdMap, streamPath)
	dir := filepath.Join(util.GetWorkPath(), s.Cfg.HlsRoot, BuildStreamFilePath(streamPath), "..")
	err := util.RemoveDir(dir)
	if err != nil {
		log.Error(err.Error())
	}
	return true
}

func (s *ToHlsServer) GetList() []string {
	s.m.Lock()
	defer s.m.Unlock()
	var slice []string
	for _, elm := range s.cmdMap {
		slice = append(slice, elm.streamPath)
	}
	sort.Strings(slice)
	return slice
}

func (s *ToHlsServer) DeleteAndStop(streamPath string) error {
	s.m.Lock()
	defer s.m.Unlock()
	delete(s.cmdMap, streamPath)
	_, err := s.StopToHlsCmd(streamPath)
	if err != nil {
		return err
	}
	dir := filepath.Join(util.GetWorkPath(), s.Cfg.HlsRoot, BuildStreamFilePath(streamPath))
	err = util.RemoveDir(dir)
	if err != nil {
		log.Error(err.Error())
	}
	return nil
}

func (s *ToHlsServer) BuildToHlsCmd(url string, streamPath string, timeoutMs int, hlsCfg *ToHlsConfig) *exec.Cmd {
	dir := filepath.Join(util.GetWorkPath(), s.Cfg.HlsRoot, BuildStreamFilePath(streamPath))
	_, err := util.CheckDir(dir)
	if err != nil {
		log.Error(err.Error())
	}
	p := filepath.Join(dir, hlsCfg.HlsName+".m3u8")
	buffer := new(strings.Builder)
	cmd := ffmpegCmd.
		Input(url, ffmpegCmd.KwArgs{
			"stimeout": timeoutMs * 1000,
			"loglevel": "error",
		}).
		Output(p, ffmpegCmd.KwArgs{
			"c:v":              hlsCfg.EncodeCodec,
			"profile:v":        hlsCfg.Profile,
			"vf":               util.Merge("scale=", hlsCfg.VfScale),
			"force_key_frames": util.Merge("expr:gte(t,n_forced*", hlsCfg.GopSize, ")"),
			"hls_time":         hlsCfg.HlsTime,
			"hls_list_size":    hlsCfg.HlsListSize,
			"hls_wrap":         hlsCfg.HlsWarp,
			"loglevel":         "error",
		}).
		OverWriteOutput().WithErrorOutput(buffer).Compile()
	return cmd
}

func (s *ToHlsServer) RunToHlsCmd(streamPath string) (bool, error) {
	info, ok := s.cmdMap[streamPath]
	if !ok {
		return false, nil
	}
	s.Supervisor.Activate(streamPath)
	defer s.Supervisor.DeActivate(streamPath)
	err := info.cmd.Run()
	if err != nil {
		errStr := strings.TrimSpace(fmt.Sprint(info.cmd.Stderr))
		return false, errors.New(errStr)
	}
	return true, nil
}

func (s *ToHlsServer) StopToHlsCmd(streamPath string) (bool, error) {
	info, ok := s.cmdMap[streamPath]
	if !ok {
		return true, nil
	}
	err := info.cmd.Process.Kill()
	if err != nil {
		return false, err
	}
	return true, nil
}

func (s *ToHlsServer) BuildHlsPath(streamPath string) string {
	return path.Join("/"+s.Cfg.HlsRoot, streamPath, s.Cfg.HlsName+".m3u8")
}

var (
	handlePath string
)

func (s *ToHlsServer) CreateKey(oriUrl string) (string, bool) {
	url := oriUrl
	paramArr := strings.Split(url, "/")
	if len(paramArr) != 5 || paramArr[1] != s.Cfg.HlsRoot {
		return "", false
	}
	streamPath := paramArr[2] + "/" + paramArr[3]
	return streamPath, true
}

func (s *ToHlsServer) Expire(key string) bool {
	_, ok := s.cmdMap[key]
	if !ok {
		return false
	}
	_ = s.DeleteAndStop(key)
	return true
}