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

package ffmpegServer

import (
	"context"
	"errors"
	"fmt"
	ffmpegCmd "github.com/u2takey/ffmpeg-go"
	"go.uber.org/zap"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"ycmediakit/internal/pkg/util"
)

type FfmpegCmdServer struct {
	m      sync.Mutex
	cmdMap map[string]*CmdBaseInfo
}

type CmdBaseInfo struct {
	CmdBaseConfig
}

type ToHlsServer struct {
	hlsCfg *ToHlsConfig
	lock   sync.Mutex
	cmdMap map[string]*ToHlsInfo
}

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

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

func (server *ToHlsServer) Add(url string, streamPath string, timeoutMs int, curHlsCfg *ToHlsConfig) bool {
	server.lock.Lock()
	defer server.lock.Unlock()
	_, ok := server.cmdMap[streamPath]
	if ok {
		return true
	}
	// merge curHlsCfg, hlsCfg
	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.lock.Lock()
	defer server.lock.Unlock()
	_, ok := server.cmdMap[streamPath]
	return ok
}

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

func (server *ToHlsServer) GetList() []string {
	server.lock.Lock()
	defer server.lock.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.lock.Lock()
	defer server.lock.Unlock()
	delete(server.cmdMap, streamPath)
	_, err := server.StopToHlsCmd(streamPath)
	if err != nil {
		return err
	}
	path := filepath.Join(util.GetWorkPath(), server.hlsCfg.HlsRoot, streamPath)
	err = util.RemoveDir(path)
	if err != nil {
		zap.L().Error(err.Error())
	}
	return nil
}

func (server *ToHlsServer) PrepareUrl(url string) (string, bool) {
	if url == "" {
		return "", false
	}
	urlArr := strings.Split(url, "?")
	if len(urlArr) > 2 {
		return "", false
	}
	url = urlArr[0]
	return url, true
}

func (server *ToHlsServer) PrepareStreamPath(streamPath string) (string, bool) {
	streamArr := strings.Split(streamPath, "/")
	if len(streamArr) != 2 {
		return "", false
	}
	streamPath = filepath.Join(streamArr...)
	return streamPath, true
}

func (server *ToHlsServer) ProbeWithTimeout(url string, timeoutMs int) error {
	args := ffmpegCmd.ConvertKwargsToCmdLineArgs(ffmpegCmd.KwArgs{
		"stimeout": timeoutMs * 1000,
		"loglevel": "error",
	})
	args = append(args, url)
	ctx := context.Background()
	cmd := exec.CommandContext(ctx, "ffprobe", args...)
	errBuf := new(strings.Builder)
	cmd.Stderr = errBuf
	err := cmd.Run()
	if err != nil {
		errStr := strings.TrimSpace(errBuf.String())
		return errors.New(errStr)
	}
	return nil
}

func (server *ToHlsServer) BuildToHlsCmd(url string, streamPath string, timeoutMs int, hlsCfg *ToHlsConfig) *exec.Cmd {
	path := filepath.Join(util.GetWorkPath(), server.hlsCfg.HlsRoot, streamPath)
	_, err := util.CheckDir(path)
	if err != nil {
		zap.L().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.hlsCfg.HlsRoot, streamPath, server.hlsCfg.HlsName+".m3u8")
}