From 9341eb8d3fd3f0670b239952fda5f66ae9fbf567 Mon Sep 17 00:00:00 2001
From: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com>
Date: Mon, 11 Nov 2024 14:21:45 +0200
Subject: [PATCH] introduce osutil.LoadAvg function, works on linux and darwin

Signed-off-by: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com>
---
 go/osutil/os.go      | 63 ++++++++++++++++++++++++++++++++++++
 go/osutil/os_test.go | 77 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 140 insertions(+)
 create mode 100644 go/osutil/os.go
 create mode 100644 go/osutil/os_test.go

diff --git a/go/osutil/os.go b/go/osutil/os.go
new file mode 100644
index 00000000000..a3fdc3bce6f
--- /dev/null
+++ b/go/osutil/os.go
@@ -0,0 +1,63 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package osutil
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"runtime"
+	"strconv"
+	"strings"
+)
+
+// LoadAvg returns the past 1 minute system load average. This works on linux and darwin systems.
+// On other systems, it returns 0 with no error.
+func LoadAvg() (float64, error) {
+	switch runtime.GOOS {
+	case "linux":
+		content, err := os.ReadFile("/proc/loadavg")
+		if err != nil {
+			return 0, err
+		}
+		return parseLoadAvg(string(content))
+	case "darwin":
+		cmd := exec.Command("sysctl", "-n", "vm.loadavg")
+		// Sample output: `{ 2.83 3.01 3.36 }`
+		output, err := cmd.CombinedOutput()
+		if err != nil {
+			return 0, err
+		}
+		if len(output) < 1 {
+			return 0, fmt.Errorf("unexpected sysctl output: %q", output)
+		}
+		output = output[1:] // Remove the leading `{ `
+		return parseLoadAvg(string(output))
+	default:
+		return 0, nil
+	}
+}
+
+// parseLoadAvg parses the load average from the content of /proc/loadavg or sysctl output.
+// Input such as "1.00 0.99 0.98 1/1 1", "2.83 3.01 3.36"
+func parseLoadAvg(content string) (float64, error) {
+	fields := strings.Fields(content)
+	if len(fields) == 0 {
+		return 0, fmt.Errorf("unexpected loadavg content: %s", content)
+	}
+	return strconv.ParseFloat(fields[0], 64)
+}
diff --git a/go/osutil/os_test.go b/go/osutil/os_test.go
new file mode 100644
index 00000000000..5f3831648c6
--- /dev/null
+++ b/go/osutil/os_test.go
@@ -0,0 +1,77 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package osutil
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLoadAvgValue(t *testing.T) {
+	tcases := []struct {
+		input   string
+		loadavg float64
+		isError bool
+	}{
+		{
+			input:   "",
+			isError: true,
+		},
+		{
+			input:   "{}",
+			isError: true,
+		},
+		{
+			input:   "{ x y z }",
+			isError: true,
+		},
+		{
+			input:   "1",
+			loadavg: 1.0,
+		},
+		{
+			input:   "0.00 0.00 0.00 1/1 1",
+			loadavg: 0.0,
+		},
+		{
+			input:   "2.72 2.89 3.17",
+			loadavg: 2.72,
+		},
+		{
+			input:   "   2.72 2.89 3.17",
+			loadavg: 2.72,
+		},
+	}
+	for _, tcase := range tcases {
+		t.Run(tcase.input, func(t *testing.T) {
+			loadavg, err := parseLoadAvg(tcase.input)
+			if tcase.isError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tcase.loadavg, loadavg)
+			}
+		})
+	}
+}
+
+func TestLoadAvg(t *testing.T) {
+	loadavg, err := LoadAvg()
+	assert.NoError(t, err)
+	assert.GreaterOrEqual(t, loadavg, 0.0)
+}