change type of memlimit from string to int64

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