change type of memlimit from string to int64

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

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=

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()
 	}

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()

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"`

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")