add memory abstraction

2 files changed, 126 insertions, 0 deletions

diff --git a/client/go/jvm/memory.go b/client/go/jvm/memory.go
new file mode 100644
index 00000000000..dcbf2d3d8b1
--- /dev/null
+++ b/client/go/jvm/memory.go
@@ -0,0 +1,79 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Author: arnej
+
+package jvm
+
+import (
+	"fmt"
+)
+
+const (
+	_ = 1 << (10 * iota)
+	PowerOfTwo10
+	PowerOfTwo20
+	PowerOfTwo30
+)
+
+type AmountOfMemory struct {
+	numBytes int64
+}
+
+func BytesOfMemory(v int64) AmountOfMemory {
+	return AmountOfMemory{numBytes: v}
+}
+func KiloBytesOfMemory(v int64) AmountOfMemory {
+	return BytesOfMemory(v * PowerOfTwo10)
+}
+func MegaBytesOfMemory(v int) AmountOfMemory {
+	return BytesOfMemory(int64(v) * PowerOfTwo20)
+}
+func GigaBytesOfMemory(v int) AmountOfMemory {
+	return BytesOfMemory(int64(v) * PowerOfTwo30)
+}
+
+func (v AmountOfMemory) ToBytes() int64 {
+	return v.numBytes
+}
+func (v AmountOfMemory) ToKB() int64 {
+	return v.numBytes / PowerOfTwo10
+}
+func (v AmountOfMemory) ToMB() int {
+	return int(v.numBytes / PowerOfTwo20)
+}
+func (v AmountOfMemory) ToGB() int {
+	return int(v.numBytes / PowerOfTwo30)
+}
+func (v AmountOfMemory) AsJvmSpec() string {
+	val := v.ToKB()
+	suffix := "k"
+	if val%PowerOfTwo10 == 0 {
+		val = val / PowerOfTwo10
+		suffix = "m"
+		if val%PowerOfTwo10 == 0 {
+			val = val / PowerOfTwo10
+			suffix = "g"
+		}
+	}
+	return fmt.Sprintf("%d%s", val, suffix)
+}
+
+func ParseJvmMemorySpec(spec string) (result AmountOfMemory, err error) {
+	result = BytesOfMemory(0)
+	var n int
+	var val int64
+	var suffix rune
+	n, err = fmt.Sscanf(spec, "%d%c", &val, &suffix)
+	if n == 2 && err == nil {
+		switch suffix {
+		case 'k':
+			result = KiloBytesOfMemory(val)
+		case 'm':
+			result = MegaBytesOfMemory(int(val))
+		case 'g':
+			result = GigaBytesOfMemory(int(val))
+		default:
+			err = fmt.Errorf("Unknown suffix in JVM memory spec '%s'", spec)
+		}
+	}
+	return
+}
diff --git a/client/go/jvm/memory_test.go b/client/go/jvm/memory_test.go
new file mode 100644
index 00000000000..c2956513811
--- /dev/null
+++ b/client/go/jvm/memory_test.go
@@ -0,0 +1,47 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+package jvm
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestConversion(t *testing.T) {
+	v1 := GigaBytesOfMemory(17)
+	v2 := MegaBytesOfMemory(17 * 1024)
+	v3 := KiloBytesOfMemory(17 * 1024 * 1024)
+	var numBytes int64 = 17 * 1024 * 1024 * 1024
+	assert.Equal(t, v1, v2)
+	assert.Equal(t, v1, v3)
+	assert.Equal(t, v2, v1)
+	assert.Equal(t, v2, v3)
+	assert.Equal(t, v3, v1)
+	assert.Equal(t, v3, v2)
+	assert.Equal(t, numBytes, v1.ToBytes())
+	assert.Equal(t, numBytes, v2.ToBytes())
+	assert.Equal(t, numBytes, v3.ToBytes())
+	assert.Equal(t, "17g", v1.AsJvmSpec())
+	assert.Equal(t, "17g", v2.AsJvmSpec())
+	assert.Equal(t, "17g", v3.AsJvmSpec())
+
+	v1 = GigaBytesOfMemory(17)
+	v2 = MegaBytesOfMemory(17 * 1000)
+	v3 = KiloBytesOfMemory(17 * 1000 * 1000)
+	assert.Equal(t, "17g", v1.AsJvmSpec())
+	assert.Equal(t, "17000m", v2.AsJvmSpec())
+	assert.Equal(t, "17000000k", v3.AsJvmSpec())
+
+	var result AmountOfMemory
+	var err error
+	result, err = ParseJvmMemorySpec("17g")
+	assert.Nil(t, err)
+	assert.Equal(t, v1, result)
+	result, err = ParseJvmMemorySpec("17000m")
+	assert.Nil(t, err)
+	assert.Equal(t, v2, result)
+	result, err = ParseJvmMemorySpec("17000000k")
+	assert.Nil(t, err)
+	assert.Equal(t, v3, result)
+}