From 9baca975ac5f0054956091aadbacfee12ab373a8 Mon Sep 17 00:00:00 2001
From: Szymon Matejczyk <szymon.matejczyk@snowflake.com>
Date: Tue, 28 Jan 2025 13:08:06 +0100
Subject: [PATCH] Expose oom_kill as part of SwapMemory

oom_kill from `/proc/vmstat` represents the number of times OOM Killer killed a process. It's useful to debug process memory issues.
---
 mem/mem.go                                        |  1 +
 mem/mem_linux.go                                  |  6 +++
 mem/mem_linux_test.go                             | 45 +++++++++++++++++++++++
 mem/mem_test.go                                   |  3 +-
 mem/testdata/linux/swapmemory/oomkill/proc/vmstat | 15 ++++++++
 5 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 mem/testdata/linux/swapmemory/oomkill/proc/vmstat

diff --git a/mem/mem.go b/mem/mem.go
index 0da71a9..77a3f12 100644
--- a/mem/mem.go
+++ b/mem/mem.go
@@ -97,6 +97,7 @@ type SwapMemoryStat struct {
 	// Linux specific numbers
 	// https://www.kernel.org/doc/Documentation/cgroup-v2.txt
 	PgMajFault uint64 `json:"pgMajFault"`
+	OomKill    uint64 `json:"oomKill"`
 }
 
 func (m VirtualMemoryStat) String() string {
diff --git a/mem/mem_linux.go b/mem/mem_linux.go
index 05bfdaf..e5dc1da 100644
--- a/mem/mem_linux.go
+++ b/mem/mem_linux.go
@@ -375,6 +375,12 @@ func SwapMemoryWithContext(ctx context.Context) (*SwapMemoryStat, error) {
 				continue
 			}
 			ret.PgMajFault = value * 4 * 1024
+		case "oom_kill":
+			value, err := strconv.ParseUint(fields[1], 10, 64)
+			if err != nil {
+				continue
+			}
+			ret.OomKill = value
 		}
 	}
 	return ret, nil
diff --git a/mem/mem_linux_test.go b/mem/mem_linux_test.go
index b71b71c..9676dd8 100644
--- a/mem/mem_linux_test.go
+++ b/mem/mem_linux_test.go
@@ -177,3 +177,48 @@ func TestParseSwapsFile_EmptyFile(t *testing.T) {
 	_, err := parseSwapsFile(context.Background(), strings.NewReader(""))
 	assert.Error(t, err)
 }
+
+var swapMemoryVmstatTests = []struct {
+	mockedRootFS string
+	swap         *SwapMemoryStat
+}{
+	{
+		"oomkill", &SwapMemoryStat{
+			// not checked
+			Total:       0,
+			Used:        0,
+			Free:        0,
+			UsedPercent: 0,
+
+			// checked
+			PgIn:       1 * 4 * 1024,
+			PgOut:      2 * 4 * 1024,
+			PgFault:    3 * 4 * 1024,
+			PgMajFault: 4 * 4 * 1024,
+			Sin:        3 * 4 * 1024,
+			Sout:       4 * 4 * 1024,
+			OomKill:    5,
+		},
+	},
+}
+
+func TestSwapVmstatMemoryLinux(t *testing.T) {
+	for _, tt := range swapMemoryVmstatTests {
+		t.Run(tt.mockedRootFS, func(t *testing.T) {
+			t.Setenv("HOST_PROC", filepath.Join("testdata/linux/swapmemory/", tt.mockedRootFS, "proc"))
+
+			stat, err := SwapMemory()
+			stat.Total = 0
+			stat.Used = 0
+			stat.Free = 0
+			stat.UsedPercent = 0
+			skipIfNotImplementedErr(t, err)
+			if err != nil {
+				t.Errorf("error %v", err)
+			}
+			if !reflect.DeepEqual(stat, tt.swap) {
+				t.Errorf("got: %+v\nwant: %+v", stat, tt.swap)
+			}
+		})
+	}
+}
diff --git a/mem/mem_test.go b/mem/mem_test.go
index 8d449da..8d42e15 100644
--- a/mem/mem_test.go
+++ b/mem/mem_test.go
@@ -108,8 +108,9 @@ func TestSwapMemoryStat_String(t *testing.T) {
 		PgOut:       4,
 		PgFault:     5,
 		PgMajFault:  6,
+		OomKill:     7,
 	}
-	e := `{"total":10,"used":30,"free":40,"usedPercent":30.1,"sin":1,"sout":2,"pgIn":3,"pgOut":4,"pgFault":5,"pgMajFault":6}`
+	e := `{"total":10,"used":30,"free":40,"usedPercent":30.1,"sin":1,"sout":2,"pgIn":3,"pgOut":4,"pgFault":5,"pgMajFault":6,"oomKill":7}`
 	if e != fmt.Sprintf("%v", v) {
 		t.Errorf("SwapMemoryStat string is invalid: %v", v)
 	}
diff --git a/mem/testdata/linux/swapmemory/oomkill/proc/vmstat b/mem/testdata/linux/swapmemory/oomkill/proc/vmstat
new file mode 100644
index 0000000..ad2bec9
--- /dev/null
+++ b/mem/testdata/linux/swapmemory/oomkill/proc/vmstat
@@ -0,0 +1,15 @@
+pgpgin 1
+pgpgout 2
+pswpin 3
+pswpout 4
+pgalloc_dma 5
+pgalloc_dma32 6
+pgalloc_normal 7
+pgalloc_movable 8
+pgfree 9
+pgactivate 10
+pgdeactivate 1
+pglazyfree 2
+pgfault 3
+pgmajfault 4
+oom_kill 5
\ No newline at end of file