Add support for reading total memory

internal/cgroups/cgroups.go

@@ -41,6 +41,8 @@ const (
 	// _cgroupCPUCFSPeriodUsParam is the file name for the CGroup CFS period
 	// parameter.
 	_cgroupCPUCFSPeriodUsParam = "cpu.cfs_period_us"
+	// _cgroupMemoryLimitBytes is the file name for CGroup total memory parameter.
+	_cgroupMemoryLimitBytes = "memory.limit_in_bytes"
 )
 
 const (
@@ -115,3 +117,17 @@ func (cg CGroups) CPUQuota() (float64, bool, error) {
 
 	return float64(cfsQuotaUs) / float64(cfsPeriodUs), true, nil
 }
+
+// TotalMemoryQuota returns total memory quota from `memory.limit_in_bytes`.
+func (cg CGroups) TotalMemoryQuota() (int64, bool, error) {
+	cpuCGroup, exists := cg[_cgroupSubsysMemory]
+	if !exists {
+		return -1, false, nil
+	}
+
+	memLimitBytes, err := cpuCGroup.readInt(_cgroupMemoryLimitBytes)
+	if defined := memLimitBytes > 0; err != nil || !defined {
+		return -1, defined, err
+	}
+	return int64(memLimitBytes), true, nil
+}

internal/cgroups/cgroups_test.go

@@ -27,6 +27,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestNewCGroups(t *testing.T) {
@@ -131,3 +132,12 @@ func TestCGroupsCPUQuota(t *testing.T) {
 		}
 	}
 }
+
+func TestTotalMemory(t *testing.T) {
+	cgroups, err := NewCGroupsForCurrentProcess()
+	require.NoError(t, err)
+	quota, defined, err := cgroups.TotalMemoryQuota()
+	require.NoError(t, err)
+	assert.Equal(t, true, defined)
+	assert.True(t, quota > 0)
+}

internal/runtime/cpu_quota_linux.go

@@ -47,3 +47,18 @@ func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
 	}
 	return maxProcs, CPUQuotaUsed, nil
 }
+
+// TotalMemory returns total available memory.
+func TotalMemory() (int64, TotalMemoryStatus, error) {
+	cgroups, err := cg.NewCGroupsForCurrentProcess()
+	if err != nil {
+		return -1, TotalMemoryUndefined, err
+	}
+
+	quota, defined, err := cgroups.TotalMemoryQuota()
+	if !defined || err != nil {
+		return -1, TotalMemoryUndefined, err
+	}
+
+	return quota, TotalMemoryUsed, nil
+}

internal/runtime/cpu_quota_unsupported.go

@@ -28,3 +28,8 @@ package runtime
 func CPUQuotaToGOMAXPROCS(_ int) (int, CPUQuotaStatus, error) {
 	return -1, CPUQuotaUndefined, nil
 }
+
+// TotalMemory returns total available memory.
+func TotalMemory() (int64, TotalMemoryStatus, error) {
+	return -1, TotalMemoryUndefined, nil
+}

internal/runtime/runtime.go

@@ -23,11 +23,19 @@ package runtime
 // CPUQuotaStatus presents the status of how CPU quota is used
 type CPUQuotaStatus int
 
+// TotalMemoryStatus presents the status of how total memory quota is used
+type TotalMemoryStatus int
+
 const (
 	// CPUQuotaUndefined is returned when CPU quota is undefined
 	CPUQuotaUndefined CPUQuotaStatus = iota
 	// CPUQuotaUsed is returned when a valid CPU quota can be used
 	CPUQuotaUsed
 	// CPUQuotaMinUsed is return when CPU quota is smaller than the min value
 	CPUQuotaMinUsed
+
+	// TotalMemoryUndefined is returned when total memory is undefined
+	TotalMemoryUndefined TotalMemoryStatus = iota
+	// TotalMemoryUsed is returned when a valid total memory quota can be used
+	TotalMemoryUsed
 )

totalmemory/totalmemory.go

@@ -0,0 +1,34 @@
+// Copyright (c) 2017 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package totalmemory
+
+import iruntime "go.uber.org/automaxprocs/internal/runtime"
+
+// TotalMemory returns total memory in bytes or -1 when no total memory cannot be read from the system.
+// The backing implementation reads the total memory from cgroup `memory.limit_in_bytes`.
+// When running on no linux systems the implementation returns -1.
+func TotalMemory() (int64, error) {
+	totalMemory, status, err := iruntime.TotalMemory()
+	if status == iruntime.TotalMemoryUndefined || err != nil {
+		return -1, err
+	}
+	return totalMemory, nil
+}

totalmemory/totalmemory_test.go

@@ -0,0 +1,34 @@
+// Copyright (c) 2017 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package totalmemory
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTotalMemory(t *testing.T) {
+	totalMemory, err := TotalMemory()
+	require.NoError(t, err)
+	assert.True(t, totalMemory > 0)
+}