Allows ignoring memory modules

memory.go

@@ -32,14 +32,43 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-func NewMemory(root string) *memoryController {
-	return &memoryController{
-		root: filepath.Join(root, string(Memory)),
+// NewMemory returns a Memory controller given the root folder of cgroups.
+// It may optionally accept other configuration options, such as IgnoreModules(...)
+func NewMemory(root string, options ...func(*memoryController)) *memoryController {
+	mc := &memoryController{
+		root:    filepath.Join(root, string(Memory)),
+		ignored: map[string]struct{}{},
+	}
+	for _, opt := range options {
+		opt(mc)
+	}
+	return mc
+}
+
+// IgnoreModules configure the memory controller to not read memory metrics for some
+// module names (e.g. passing "memsw" would avoid all the memory.memsw.* entries)
+func IgnoreModules(names ...string) func(*memoryController) {
+	return func(mc *memoryController) {
+		for _, name := range names {
+			mc.ignored[name] = struct{}{}
+		}
+	}
+}
+
+// OptionalSwap allows the memory controller to not fail if cgroups is not accounting
+// Swap memory (there are no memory.memsw.* entries)
+func OptionalSwap() func(*memoryController) {
+	return func(mc *memoryController) {
+		_, err := os.Stat(filepath.Join(mc.root, "memory.memsw.usage_in_bytes"))
+		if os.IsNotExist(err) {
+			mc.ignored["memsw"] = struct{}{}
+		}
 	}
 }
 
 type memoryController struct {
-	root string
+	root    string
+	ignored map[string]struct{}
 }
 
 func (m *memoryController) Name() Name {
@@ -133,6 +162,9 @@ func (m *memoryController) Stat(path string, stats *v1.Metrics) error {
 			entry:  stats.Memory.KernelTCP,
 		},
 	} {
+		if _, ok := m.ignored[t.module]; ok {
+			continue
+		}
 		for _, tt := range []struct {
 			name  string
 			value *uint64

memory_test.go

@@ -17,6 +17,10 @@
 package cgroups
 
 import (
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path"
 	"strings"
 	"testing"
 
@@ -106,3 +110,154 @@ func TestParseMemoryStats(t *testing.T) {
 		}
 	}
 }
+
+func TestMemoryController_Stat(t *testing.T) {
+	// GIVEN a cgroups folder with all the memory metrics
+	modules := []string{"", "memsw", "kmem", "kmem.tcp"}
+	metrics := []string{"usage_in_bytes", "max_usage_in_bytes", "failcnt", "limit_in_bytes"}
+	tmpRoot := buildMemoryMetrics(t, modules, metrics)
+
+	// WHEN the memory controller reads the metrics stats
+	mc := NewMemory(tmpRoot)
+	stats := v1.Metrics{}
+	if err := mc.Stat("", &stats); err != nil {
+		t.Errorf("can't get stats: %v", err)
+	}
+
+	// THEN all the memory stats have been completely loaded in memory
+	checkMemoryStatIsComplete(t, stats.Memory)
+}
+
+func TestMemoryController_Stat_IgnoreModules(t *testing.T) {
+	// GIVEN a cgroups folder that accounts for all the metrics BUT swap memory
+	modules := []string{"", "kmem", "kmem.tcp"}
+	metrics := []string{"usage_in_bytes", "max_usage_in_bytes", "failcnt", "limit_in_bytes"}
+	tmpRoot := buildMemoryMetrics(t, modules, metrics)
+
+	// WHEN the memory controller explicitly ignores memsw module and reads the data
+	mc := NewMemory(tmpRoot, IgnoreModules("memsw"))
+	stats := v1.Metrics{}
+	if err := mc.Stat("", &stats); err != nil {
+		t.Errorf("can't get stats: %v", err)
+	}
+
+	// THEN the swap memory stats are not loaded but all the other memory metrics are
+	checkMemoryStatHasNoSwap(t, stats.Memory)
+}
+
+func TestMemoryController_Stat_OptionalSwap_HasSwap(t *testing.T) {
+	// GIVEN a cgroups folder with all the memory metrics
+	modules := []string{"", "memsw", "kmem", "kmem.tcp"}
+	metrics := []string{"usage_in_bytes", "max_usage_in_bytes", "failcnt", "limit_in_bytes"}
+	tmpRoot := buildMemoryMetrics(t, modules, metrics)
+
+	// WHEN a memory controller that ignores swap only if it is missing reads stats
+	mc := NewMemory(tmpRoot, OptionalSwap())
+	stats := v1.Metrics{}
+	if err := mc.Stat("", &stats); err != nil {
+		t.Errorf("can't get stats: %v", err)
+	}
+
+	// THEN all the memory stats have been completely loaded in memory
+	checkMemoryStatIsComplete(t, stats.Memory)
+}
+
+func TestMemoryController_Stat_OptionalSwap_NoSwap(t *testing.T) {
+	// GIVEN a cgroups folder that accounts for all the metrics BUT swap memory
+	modules := []string{"", "kmem", "kmem.tcp"}
+	metrics := []string{"usage_in_bytes", "max_usage_in_bytes", "failcnt", "limit_in_bytes"}
+	tmpRoot := buildMemoryMetrics(t, modules, metrics)
+
+	// WHEN a memory controller that ignores swap only if it is missing reads stats
+	mc := NewMemory(tmpRoot, OptionalSwap())
+	stats := v1.Metrics{}
+	if err := mc.Stat("", &stats); err != nil {
+		t.Errorf("can't get stats: %v", err)
+	}
+
+	// THEN the swap memory stats are not loaded but all the other memory metrics are
+	checkMemoryStatHasNoSwap(t, stats.Memory)
+}
+
+func checkMemoryStatIsComplete(t *testing.T, mem *v1.MemoryStat) {
+	index := []uint64{
+		mem.Usage.Usage,
+		mem.Usage.Max,
+		mem.Usage.Failcnt,
+		mem.Usage.Limit,
+		mem.Swap.Usage,
+		mem.Swap.Max,
+		mem.Swap.Failcnt,
+		mem.Swap.Limit,
+		mem.Kernel.Usage,
+		mem.Kernel.Max,
+		mem.Kernel.Failcnt,
+		mem.Kernel.Limit,
+		mem.KernelTCP.Usage,
+		mem.KernelTCP.Max,
+		mem.KernelTCP.Failcnt,
+		mem.KernelTCP.Limit,
+	}
+	for i, v := range index {
+		if v != uint64(i) {
+			t.Errorf("expected value at index %d to be %d but received %d", i, i, v)
+		}
+	}
+}
+
+func checkMemoryStatHasNoSwap(t *testing.T, mem *v1.MemoryStat) {
+	if mem.Swap.Usage != 0 || mem.Swap.Limit != 0 ||
+		mem.Swap.Max != 0 || mem.Swap.Failcnt != 0 {
+		t.Errorf("swap memory should have been ignored. Got: %+v", mem.Swap)
+	}
+	index := []uint64{
+		mem.Usage.Usage,
+		mem.Usage.Max,
+		mem.Usage.Failcnt,
+		mem.Usage.Limit,
+		mem.Kernel.Usage,
+		mem.Kernel.Max,
+		mem.Kernel.Failcnt,
+		mem.Kernel.Limit,
+		mem.KernelTCP.Usage,
+		mem.KernelTCP.Max,
+		mem.KernelTCP.Failcnt,
+		mem.KernelTCP.Limit,
+	}
+	for i, v := range index {
+		if v != uint64(i) {
+			t.Errorf("expected value at index %d to be %d but received %d", i, i, v)
+		}
+	}
+}
+
+// buildMemoryMetrics creates fake cgroups memory entries in a temporary dir. Returns the fake cgroups root
+func buildMemoryMetrics(t *testing.T, modules []string, metrics []string) string {
+	tmpRoot, err := ioutil.TempDir("", "memtests")
+	if err != nil {
+		t.Fatal(err)
+	}
+	tmpDir := path.Join(tmpRoot, string(Memory))
+	if err := os.MkdirAll(tmpDir, defaultDirPerm); err != nil {
+		t.Fatal(err)
+	}
+	if err := ioutil.WriteFile(path.Join(tmpDir, "memory.stat"), []byte(memoryData), defaultFilePerm); err != nil {
+		t.Fatal(err)
+	}
+	cnt := 0
+	for _, mod := range modules {
+		for _, metric := range metrics {
+			var fileName string
+			if mod == "" {
+				fileName = path.Join(tmpDir, strings.Join([]string{"memory", metric}, "."))
+			} else {
+				fileName = path.Join(tmpDir, strings.Join([]string{"memory", mod, metric}, "."))
+			}
+			if err := ioutil.WriteFile(fileName, []byte(fmt.Sprintln(cnt)), defaultFilePerm); err != nil {
+				t.Fatal(err)
+			}
+			cnt++
+		}
+	}
+	return tmpRoot
+}