From 3ce5c2ede39e74f6f82baf63df9f436927264f7a Mon Sep 17 00:00:00 2001 From: Miao Wang Date: Mon, 12 Jul 2021 19:09:25 +0800 Subject: [PATCH] change type of memlimit from string to int64 --- go.mod | 1 + go.sum | 2 ++ worker/cgroup.go | 8 ++++---- worker/cgroup_test.go | 5 +++-- worker/config.go | 29 ++++++++++++++++++++++++++++- worker/config_test.go | 7 +++++++ 6 files changed, 45 insertions(+), 7 deletions(-) diff --git a/go.mod b/go.mod index 85ebd77..058d01b 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/codeskyblue/go-sh v0.0.0-20190412065543-76bd3d59ff27 github.com/dennwc/btrfs v0.0.0-20190517175702-d917b30ff035 github.com/dgraph-io/badger/v2 v2.2007.2 + github.com/docker/go-units v0.4.0 // indirect github.com/flynn/go-shlex v0.0.0-20150515145356-3f9db97f8568 // indirect github.com/gin-gonic/gin v1.5.0 github.com/go-redis/redis/v8 v8.3.0 diff --git a/go.sum b/go.sum index 682d33c..cc10984 100644 --- a/go.sum +++ b/go.sum @@ -46,6 +46,8 @@ github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 h1:tdlZCpZ/P9DhczC github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw= +github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo= github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= github.com/flynn/go-shlex v0.0.0-20150515145356-3f9db97f8568 h1:BHsljHzVlRcyQhjrss6TZTdY2VfCqZPbv5k3iBFa2ZQ= diff --git a/worker/cgroup.go b/worker/cgroup.go index cf2b572..17d33ce 100644 --- a/worker/cgroup.go +++ b/worker/cgroup.go @@ -21,10 +21,10 @@ type cgroupHook struct { baseGroup string created bool subsystem string - memLimit string + memLimit MemBytes } -func newCgroupHook(p mirrorProvider, basePath, baseGroup, subsystem, memLimit string) *cgroupHook { +func newCgroupHook(p mirrorProvider, basePath, baseGroup, subsystem string, memLimit MemBytes) *cgroupHook { if basePath == "" { basePath = "/sys/fs/cgroup" } @@ -52,11 +52,11 @@ func (c *cgroupHook) preExec() error { if c.subsystem != "memory" { return nil } - if c.memLimit != "" { + if c.memLimit != 0 { gname := fmt.Sprintf("%s/%s", c.baseGroup, c.provider.Name()) return sh.Command( "cgset", "-r", - fmt.Sprintf("memory.limit_in_bytes=%s", c.memLimit), + fmt.Sprintf("memory.limit_in_bytes=%d", c.memLimit.Value()), gname, ).Run() } diff --git a/worker/cgroup_test.go b/worker/cgroup_test.go index 835ddd6..f56e59b 100644 --- a/worker/cgroup_test.go +++ b/worker/cgroup_test.go @@ -8,6 +8,7 @@ import ( "strings" "testing" "time" + units "github.com/docker/go-units" . "github.com/smartystreets/goconvey/convey" ) @@ -72,7 +73,7 @@ sleep 30 provider, err := newCmdProvider(c) So(err, ShouldBeNil) - cg := newCgroupHook(provider, "/sys/fs/cgroup", "tunasync", "cpu", "") + cg := newCgroupHook(provider, "/sys/fs/cgroup", "tunasync", "cpu", 0) provider.AddHook(cg) err = cg.preExec() @@ -132,7 +133,7 @@ sleep 30 provider, err := newRsyncProvider(c) So(err, ShouldBeNil) - cg := newCgroupHook(provider, "/sys/fs/cgroup", "tunasync", "cpu", "512M") + cg := newCgroupHook(provider, "/sys/fs/cgroup", "tunasync", "cpu", 512 * units.MiB) provider.AddHook(cg) err = cg.preExec() diff --git a/worker/config.go b/worker/config.go index b9e13f4..95a8c4e 100644 --- a/worker/config.go +++ b/worker/config.go @@ -7,6 +7,7 @@ import ( "github.com/BurntSushi/toml" "github.com/imdario/mergo" + units "github.com/docker/go-units" ) type providerEnum uint8 @@ -113,6 +114,32 @@ type includedMirrorConfig struct { Mirrors []mirrorConfig `toml:"mirrors"` } +type MemBytes int64 + +// Set sets the value of the MemBytes by passing a string +func (m *MemBytes) Set(value string) error { + val, err := units.RAMInBytes(value) + *m = MemBytes(val) + return err +} + +// Type returns the type +func (m *MemBytes) Type() string { + return "bytes" +} + +// Value returns the value in int64 +func (m *MemBytes) Value() int64 { + return int64(*m) +} + +// UnmarshalJSON is the customized unmarshaler for MemBytes +func (m *MemBytes) UnmarshalText(s []byte) error { + val, err := units.RAMInBytes(string(s)) + *m = MemBytes(val) + return err +} + type mirrorConfig struct { Name string `toml:"name"` Provider providerEnum `toml:"provider"` @@ -148,7 +175,7 @@ type mirrorConfig struct { RsyncOverride []string `toml:"rsync_override"` Stage1Profile string `toml:"stage1_profile"` - MemoryLimit string `toml:"memory_limit"` + MemoryLimit MemBytes `toml:"memory_limit"` DockerImage string `toml:"docker_image"` DockerVolumes []string `toml:"docker_volumes"` diff --git a/worker/config_test.go b/worker/config_test.go index ed054fc..c8de85d 100644 --- a/worker/config_test.go +++ b/worker/config_test.go @@ -7,6 +7,7 @@ import ( "path/filepath" "testing" "time" + units "github.com/docker/go-units" . "github.com/smartystreets/goconvey/convey" ) @@ -53,12 +54,15 @@ provider = "two-stage-rsync" stage1_profile = "debian" upstream = "rsync://ftp.debian.org/debian/" use_ipv6 = true +memory_limit = "256MiB" [[mirrors]] name = "fedora" provider = "rsync" upstream = "rsync://ftp.fedoraproject.org/fedora/" use_ipv6 = true +memory_limit = "128M" + exclude_file = "/etc/tunasync.d/fedora-exclude.txt" exec_on_failure = [ "bash -c 'echo ${TUNASYNC_JOB_EXIT_STATUS} > ${TUNASYNC_WORKING_DIR}/exit_status'" @@ -141,17 +145,20 @@ use_ipv6 = true So(m.Name, ShouldEqual, "debian") So(m.MirrorDir, ShouldEqual, "") So(m.Provider, ShouldEqual, provTwoStageRsync) + So(m.MemoryLimit.Value(), ShouldEqual, 256 * units.MiB) m = cfg.Mirrors[2] So(m.Name, ShouldEqual, "fedora") So(m.MirrorDir, ShouldEqual, "") So(m.Provider, ShouldEqual, provRsync) So(m.ExcludeFile, ShouldEqual, "/etc/tunasync.d/fedora-exclude.txt") + So(m.MemoryLimit.Value(), ShouldEqual, 128 * units.MiB) m = cfg.Mirrors[3] So(m.Name, ShouldEqual, "debian-cd") So(m.MirrorDir, ShouldEqual, "") So(m.Provider, ShouldEqual, provTwoStageRsync) + So(m.MemoryLimit.Value(), ShouldEqual, 0) m = cfg.Mirrors[4] So(m.Name, ShouldEqual, "debian-security")