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

package ffmpegServer

import (
	"errors"
	"fmt"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"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
}

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

func NewToHlsServer(hlsCfg *ToHlsConfig) *ToHlsServer {
	return &ToHlsServer{
		cfg:    hlsCfg,
		cmdMap: make(map[string]*ToHlsInfo),
	}
}

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

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

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

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

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

func (server *ToHlsServer) BuildToHlsCmd(url string, streamPath string, timeoutMs int, hlsCfg *ToHlsConfig) *exec.Cmd {
	path := filepath.Join(util.GetWorkPath(), server.cfg.HlsRoot, streamPath)
	_, err := util.CheckDir(path)
	if err != nil {
		log.Error(err.Error())
	}
	path = filepath.Join(path, hlsCfg.HlsName+".m3u8")
	buffer := new(strings.Builder)
	cmd := ffmpegCmd.
		Input(url, ffmpegCmd.KwArgs{
			"stimeout": timeoutMs * 1000,
			"loglevel": "error",
		}).
		Output(path, 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 (server *ToHlsServer) RunToHlsCmd(streamPath string) (bool, error) {
	info, ok := server.cmdMap[streamPath]
	if !ok {
		return false, nil
	}
	err := info.cmd.Run()
	if err != nil {
		errStr := strings.TrimSpace(fmt.Sprint(info.cmd.Stderr))
		return false, errors.New(errStr)
	}
	return true, nil
}

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

func (server *ToHlsServer) BuildHlsPath(streamPath string) string {
	return filepath.Join(server.cfg.HlsRoot, streamPath, server.cfg.HlsName+".m3u8")
}