From 89d652b583cb04e9cf8a8c6090ed7f20d68daa3e Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Tue, 22 Sep 2020 23:26:44 +0300
Subject: [PATCH] lib/cgroup: attempt to obtain available CPU cores via
 /sys/devices/system/cpu/online

See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/685#issuecomment-674423728
---
 lib/cgroup/cpu.go      | 51 ++++++++++++++++++++++++++++++++++++++++++
 lib/cgroup/cpu_test.go | 24 ++++++++++++++++++++
 2 files changed, 75 insertions(+)
 create mode 100644 lib/cgroup/cpu_test.go

diff --git a/lib/cgroup/cpu.go b/lib/cgroup/cpu.go
index 54f4fcdd61..6d112979d2 100644
--- a/lib/cgroup/cpu.go
+++ b/lib/cgroup/cpu.go
@@ -1,8 +1,11 @@
 package cgroup
 
 import (
+	"io/ioutil"
 	"os"
 	"runtime"
+	"strconv"
+	"strings"
 
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/logger"
 )
@@ -40,9 +43,57 @@ func getCPUQuota() float64 {
 	if err != nil {
 		return 0
 	}
+	if quotaUS <= 0 {
+		// The quota isn't set. This may be the case in multilevel containers.
+		// See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/685#issuecomment-674423728
+		return getOnlineCPUCount()
+	}
 	periodUS, err := readInt64("/sys/fs/cgroup/cpu/cpu.cfs_period_us", "cat /sys/fs/cgroup/cpu$(cat /proc/self/cgroup | grep cpu, | cut -d: -f3)/cpu.cfs_period_us")
 	if err != nil {
 		return 0
 	}
 	return float64(quotaUS) / float64(periodUS)
 }
+
+func getOnlineCPUCount() float64 {
+	// See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/685#issuecomment-674423728
+	data, err := ioutil.ReadFile("/sys/devices/system/cpu/online")
+	if err != nil {
+		return -1
+	}
+	n := float64(countCPUs(string(data)))
+	if n <= 0 {
+		return -1
+	}
+	// Add a half of CPU core, since it looks like actual cores is usually bigger than online cores.
+	// See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/685#issuecomment-674423728
+	return n + 0.5
+}
+
+func countCPUs(data string) int {
+	data = strings.TrimSpace(data)
+	n := 0
+	for _, s := range strings.Split(data, ",") {
+		n++
+		if !strings.Contains(s, "-") {
+			if _, err := strconv.Atoi(s); err != nil {
+				return -1
+			}
+			continue
+		}
+		bounds := strings.Split(s, "-")
+		if len(bounds) != 2 {
+			return -1
+		}
+		start, err := strconv.Atoi(bounds[0])
+		if err != nil {
+			return -1
+		}
+		end, err := strconv.Atoi(bounds[1])
+		if err != nil {
+			return -1
+		}
+		n += end - start
+	}
+	return n
+}
diff --git a/lib/cgroup/cpu_test.go b/lib/cgroup/cpu_test.go
new file mode 100644
index 0000000000..37aaceb6c2
--- /dev/null
+++ b/lib/cgroup/cpu_test.go
@@ -0,0 +1,24 @@
+package cgroup
+
+import (
+	"testing"
+)
+
+func TestCountCPUs(t *testing.T) {
+	f := func(s string, nExpected int) {
+		t.Helper()
+		n := countCPUs(s)
+		if n != nExpected {
+			t.Fatalf("unexpected result from countCPUs(%q); got %d; want %d", s, n, nExpected)
+		}
+	}
+	f("", -1)
+	f("1", 1)
+	f("234", 1)
+	f("1,2", 2)
+	f("0-1", 2)
+	f("0-0", 1)
+	f("1-2,3,5-9,200-210", 19)
+	f("0-3", 4)
+	f("0-6", 7)
+}