support cgroup for limit resource

run_linux.go

@@ -0,0 +1,138 @@
+package ffmpeg_go
+
+import (
+	"context"
+	"errors"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strconv"
+
+	"github.com/u2takey/go-utils/rand"
+)
+
+const (
+	cgroupConfigKey = "cgroupConfig"
+	cpuRoot         = "/sys/fs/cgroup/cpu,cpuacct"
+	cpuSetRoot      = "/sys/fs/cgroup/cpuset"
+	procsFile       = "cgroup.procs"
+	cpuSharesFile   = "cpu.shares"
+	cfsPeriodUsFile = "cpu.cfs_period_us"
+	cfsQuotaUsFile  = "cpu.cfs_quota_us"
+	cpuSetCpusFile  = "cpuset.cpus"
+	cpuSetMemsFile  = "cpuset.mems"
+)
+
+type cgroupConfig struct {
+	cpuRequest float32
+	cpuLimit   float32
+	cpuset     string
+	memset     string
+}
+
+func (s *Stream) setCGroupConfig(f func(config *cgroupConfig)) *Stream {
+	a := s.Context.Value(cgroupConfigKey)
+	if a == nil {
+		a = &cgroupConfig{}
+	}
+	f(a.(*cgroupConfig))
+	s.Context = context.WithValue(s.Context, cgroupConfigKey, a)
+	return s
+}
+
+func (s *Stream) WithCpuCoreRequest(n float32) *Stream {
+	return s.setCGroupConfig(func(config *cgroupConfig) {
+		config.cpuRequest = n
+	})
+}
+
+func (s *Stream) WithCpuCoreLimit(n float32) *Stream {
+	return s.setCGroupConfig(func(config *cgroupConfig) {
+		config.cpuLimit = n
+	})
+}
+
+func (s *Stream) WithCpuSet(n string) *Stream {
+	return s.setCGroupConfig(func(config *cgroupConfig) {
+		config.cpuset = n
+	})
+}
+
+func (s *Stream) WithMemSet(n string) *Stream {
+	return s.setCGroupConfig(func(config *cgroupConfig) {
+		config.memset = n
+	})
+}
+
+func writeCGroupFile(rootPath, file string, value string) error {
+	return ioutil.WriteFile(filepath.Join(rootPath, file), []byte(value), 0755)
+}
+
+func (s *Stream) RunLinux() error {
+	a := s.Context.Value(cgroupConfigKey).(*cgroupConfig)
+	if a.cpuRequest > a.cpuLimit {
+		return errors.New("cpuCoreLimit should greater or equal to cpuCoreRequest")
+	}
+	name := "ffmpeg_go_" + rand.String(6)
+	rootCpuPath, rootCpuSetPath := filepath.Join(cpuRoot, name), filepath.Join(cpuSetRoot, name)
+	err := os.MkdirAll(rootCpuPath, 0777)
+	if err != nil {
+		return err
+	}
+	err = os.MkdirAll(rootCpuSetPath, 0777)
+	if err != nil {
+		return err
+	}
+	defer func() { _ = os.Remove(rootCpuPath); _ = os.Remove(rootCpuSetPath) }()
+
+	share := int(1024 * a.cpuRequest)
+	period := 100000
+	quota := int(a.cpuLimit * 100000)
+
+	if share > 0 {
+		err = writeCGroupFile(rootCpuPath, cpuSharesFile, strconv.Itoa(share))
+		if err != nil {
+			return err
+		}
+	}
+	err = writeCGroupFile(rootCpuPath, cfsPeriodUsFile, strconv.Itoa(period))
+	if err != nil {
+		return err
+	}
+	if quota > 0 {
+		err = writeCGroupFile(rootCpuPath, cfsQuotaUsFile, strconv.Itoa(quota))
+		if err != nil {
+			return err
+		}
+	}
+	if a.cpuset != "" && a.memset != "" {
+		err = writeCGroupFile(rootCpuSetPath, cpuSetCpusFile, a.cpuset)
+		if err != nil {
+			return err
+		}
+		err = writeCGroupFile(rootCpuSetPath, cpuSetMemsFile, a.memset)
+		if err != nil {
+			return err
+		}
+	}
+
+	cmd := s.Compile()
+	err = cmd.Start()
+	if err != nil {
+		return err
+	}
+	if share > 0 || quota > 0 {
+		err = writeCGroupFile(rootCpuPath, procsFile, strconv.Itoa(cmd.Process.Pid))
+		if err != nil {
+			return err
+		}
+	}
+	if a.cpuset != "" && a.memset != "" {
+		err = writeCGroupFile(rootCpuSetPath, procsFile, strconv.Itoa(cmd.Process.Pid))
+		if err != nil {
+			return err
+		}
+	}
+
+	return cmd.Wait()
+}