From a669946357b93540d95178375cd5326ff40924ff Mon Sep 17 00:00:00 2001
From: Anis Eleuch <vadmeste@users.noreply.github.com>
Date: Tue, 30 Jan 2024 20:13:27 +0100
Subject: [PATCH] Add cgroup v2 support for memory limit (#18905)

---
 cmd/erasure-server-pool.go   |  7 +++++-
 cmd/handler-api.go           | 41 +++++++++++++++++++++++-------------
 internal/bpool/bpool.go      |  2 +-
 internal/bpool/bpool_test.go |  6 +++---
 4 files changed, 36 insertions(+), 20 deletions(-)

diff --git a/cmd/erasure-server-pool.go b/cmd/erasure-server-pool.go
index 3dea1b628..44d0c5ca9 100644
--- a/cmd/erasure-server-pool.go
+++ b/cmd/erasure-server-pool.go
@@ -85,7 +85,7 @@ func newErasureServerPools(ctx context.Context, endpointServerPools EndpointServ
 	)
 
 	// Maximum number of reusable buffers per node at any given point in time.
-	n := 1024 // single node single/multiple drives set this to 1024 entries
+	n := uint64(1024) // single node single/multiple drives set this to 1024 entries
 
 	if globalIsDistErasure {
 		n = 2048
@@ -95,6 +95,11 @@ func newErasureServerPools(ctx context.Context, endpointServerPools EndpointServ
 		n = 256 // 256MiB for CI/CD environments is sufficient
 	}
 
+	// Avoid allocating more than half of the available memory
+	if maxN := availableMemory() / (blockSizeV2 * 2); n > maxN {
+		n = maxN
+	}
+
 	// Initialize byte pool once for all sets, bpool size is set to
 	// setCount * setDriveCount with each memory upto blockSizeV2.
 	globalBytePoolCap = bpool.NewBytePoolCap(n, blockSizeV2, blockSizeV2*2)
diff --git a/cmd/handler-api.go b/cmd/handler-api.go
index ae7781116..e344095da 100644
--- a/cmd/handler-api.go
+++ b/cmd/handler-api.go
@@ -56,16 +56,31 @@ type apiConfig struct {
 	syncEvents                  bool
 }
 
-const cgroupLimitFile = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
+const (
+	cgroupV1MemLimitFile = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
+	cgroupV2MemLimitFile = "/sys/fs/cgroup/memory.max"
+	cgroupMemNoLimit     = 9223372036854771712
+)
 
-func cgroupLimit(limitFile string) (limit uint64) {
-	buf, err := os.ReadFile(limitFile)
+func cgroupMemLimit() (limit uint64) {
+	buf, err := os.ReadFile(cgroupV2MemLimitFile)
 	if err != nil {
-		return 9223372036854771712
+		buf, err = os.ReadFile(cgroupV1MemLimitFile)
+	}
+	if err != nil {
+		return 0
 	}
 	limit, err = strconv.ParseUint(string(buf), 10, 64)
 	if err != nil {
-		return 9223372036854771712
+		// The kernel can return valid but non integer values
+		// but still, no need to interpret more
+		return 0
+	}
+	if limit == cgroupMemNoLimit {
+		// No limit set, It's the highest positive signed 64-bit
+		// integer (2^63-1), rounded down to multiples of 4096 (2^12),
+		// the most common page size on x86 systems - for cgroup_limits.
+		return 0
 	}
 	return limit
 }
@@ -74,23 +89,19 @@ func availableMemory() (available uint64) {
 	available = 8 << 30 // Default to 8 GiB when we can't find the limits.
 
 	if runtime.GOOS == "linux" {
-		available = cgroupLimit(cgroupLimitFile)
-
-		// No limit set, It's the highest positive signed 64-bit
-		// integer (2^63-1), rounded down to multiples of 4096 (2^12),
-		// the most common page size on x86 systems - for cgroup_limits.
-		if available != 9223372036854771712 {
-			// This means cgroup memory limit is configured.
+		// Useful in container mode
+		limit := cgroupMemLimit()
+		if limit > 0 {
+			// A valid value is found
+			available = limit
 			return
-		} // no-limit set proceed to set the limits based on virtual memory.
-
+		}
 	} // for all other platforms limits are based on virtual memory.
 
 	memStats, err := mem.VirtualMemory()
 	if err != nil {
 		return
 	}
-
 	available = memStats.Available / 2
 	return
 }
diff --git a/internal/bpool/bpool.go b/internal/bpool/bpool.go
index 4f69f7f7a..d73367346 100644
--- a/internal/bpool/bpool.go
+++ b/internal/bpool/bpool.go
@@ -28,7 +28,7 @@ type BytePoolCap struct {
 
 // NewBytePoolCap creates a new BytePool bounded to the given maxSize, with new
 // byte arrays sized based on width.
-func NewBytePoolCap(maxSize int, width int, capwidth int) (bp *BytePoolCap) {
+func NewBytePoolCap(maxSize uint64, width int, capwidth int) (bp *BytePoolCap) {
 	if capwidth > 0 && capwidth < 64 {
 		panic("buffer capped with smaller than 64 bytes is not supported")
 	}
diff --git a/internal/bpool/bpool_test.go b/internal/bpool/bpool_test.go
index 28fee5993..1b122284e 100644
--- a/internal/bpool/bpool_test.go
+++ b/internal/bpool/bpool_test.go
@@ -21,7 +21,7 @@ import "testing"
 
 // Tests - bytePool functionality.
 func TestBytePool(t *testing.T) {
-	size := 4
+	size := uint64(4)
 	width := 1024
 	capWidth := 2048
 
@@ -49,7 +49,7 @@ func TestBytePool(t *testing.T) {
 	bufPool.Put(b)
 
 	// Fill the pool beyond the capped pool size.
-	for i := 0; i < size*2; i++ {
+	for i := uint64(0); i < size*2; i++ {
 		bufPool.Put(make([]byte, bufPool.w))
 	}
 
@@ -67,7 +67,7 @@ func TestBytePool(t *testing.T) {
 	close(bufPool.c)
 
 	// Check the size of the pool.
-	if len(bufPool.c) != size {
+	if uint64(len(bufPool.c)) != size {
 		t.Fatalf("bytepool size invalid: got %v want %v", len(bufPool.c), size)
 	}