From a37c8ca53194fd220b65799a8906a0bd19528b62 Mon Sep 17 00:00:00 2001
From: wangzhuowei <wangzhuowei@bytedance.com>
Date: Mon, 3 Jun 2024 19:54:55 +0800
Subject: [PATCH 1/4] fix: mcache skip manage > 1gb buffer

---
 lang/mcache/mcache.go      | 20 +++++++++++++++-----
 lang/mcache/mcache_test.go | 29 ++++++++++++++++++++++++++---
 2 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/lang/mcache/mcache.go b/lang/mcache/mcache.go
index 1e07472f..600427ac 100644
--- a/lang/mcache/mcache.go
+++ b/lang/mcache/mcache.go
@@ -20,13 +20,18 @@ import (
 	"github.com/bytedance/gopkg/lang/dirtmake"
 )
 
-const maxSize = 46
+const (
+	maxCachedIndex = 30 // len(caches) = maxCachedIndex+1
+
+	// mcache will never cache the buffer that out of maxCachedSize, which will cause high memory usage
+	maxCachedSize = 1 << maxCachedIndex // 1GB
+)
 
 // index contains []byte which cap is 1<<index
-var caches [maxSize]sync.Pool
+var caches [maxCachedIndex + 1]sync.Pool
 
 func init() {
-	for i := 0; i < maxSize; i++ {
+	for i := 0; i <= maxCachedIndex; i++ {
 		size := 1 << i
 		caches[i].New = func() interface{} {
 			buf := dirtmake.Bytes(0, size)
@@ -57,7 +62,12 @@ func Malloc(size int, capacity ...int) []byte {
 	if len(capacity) > 0 && capacity[0] > size {
 		c = capacity[0]
 	}
-	var ret = caches[calcIndex(c)].Get().([]byte)
+	if c > maxCachedSize {
+		return make([]byte, size, c)
+	}
+
+	sizeIdx := calcIndex(c)
+	var ret = caches[sizeIdx].Get().([]byte)
 	ret = ret[:size]
 	return ret
 }
@@ -65,7 +75,7 @@ func Malloc(size int, capacity ...int) []byte {
 // Free should be called when the buf is no longer used.
 func Free(buf []byte) {
 	size := cap(buf)
-	if !isPowerOfTwo(size) {
+	if size > maxCachedSize || !isPowerOfTwo(size) {
 		return
 	}
 	buf = buf[:0]
diff --git a/lang/mcache/mcache_test.go b/lang/mcache/mcache_test.go
index 1cfc9de4..964084a0 100644
--- a/lang/mcache/mcache_test.go
+++ b/lang/mcache/mcache_test.go
@@ -14,11 +14,34 @@
 
 package mcache
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
 
 func TestMalloc(t *testing.T) {
-	buf := Malloc(4096)
-	t.Log(cap(buf))
+	// cached by mcache
+	size := 4096 - 1
+	buf := Malloc(size)
+	assert.Equal(t, len(buf), size)
+	assert.Equal(t, cap(buf), size+1)
+	_ = buf[:size+1] // resize
+	Free(buf)
+
+	// cached by mcache
+	size = 1024 * 1024 * 1024
+	buf = Malloc(size)
+	assert.Equal(t, len(buf), size)
+	assert.Equal(t, cap(buf), size)
+	Free(buf)
+
+	// not cached by mcache
+	size = 1024*1024*1024 + 1
+	buf = Malloc(size)
+	assert.Equal(t, len(buf), size)
+	assert.Equal(t, cap(buf), size)
+	Free(buf)
 }
 
 func BenchmarkNormal4096(b *testing.B) {

From 21ff3be0138dfbf45aa08bd490049cb428f64041 Mon Sep 17 00:00:00 2001
From: wangzhuowei <wangzhuowei@bytedance.com>
Date: Tue, 4 Jun 2024 14:25:26 +0800
Subject: [PATCH 2/4] fix: use dirtymake for large malloc

---
 lang/mcache/mcache.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lang/mcache/mcache.go b/lang/mcache/mcache.go
index 600427ac..402bfb6f 100644
--- a/lang/mcache/mcache.go
+++ b/lang/mcache/mcache.go
@@ -63,7 +63,7 @@ func Malloc(size int, capacity ...int) []byte {
 		c = capacity[0]
 	}
 	if c > maxCachedSize {
-		return make([]byte, size, c)
+		return dirtmake.Bytes(size, c)
 	}
 
 	sizeIdx := calcIndex(c)

From 65caa4f9bb95db9488ce7356775c36ab86b6ad19 Mon Sep 17 00:00:00 2001
From: wangzhuowei <wangzhuowei@bytedance.com>
Date: Fri, 7 Jun 2024 11:17:00 +0800
Subject: [PATCH 3/4] fix: mcache Malloc should repect capacity

---
 lang/mcache/mcache.go      |  2 +-
 lang/mcache/mcache_test.go | 11 ++++++++---
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/lang/mcache/mcache.go b/lang/mcache/mcache.go
index 402bfb6f..ab32ea4f 100644
--- a/lang/mcache/mcache.go
+++ b/lang/mcache/mcache.go
@@ -68,7 +68,7 @@ func Malloc(size int, capacity ...int) []byte {
 
 	sizeIdx := calcIndex(c)
 	var ret = caches[sizeIdx].Get().([]byte)
-	ret = ret[:size]
+	ret = ret[:size:c]
 	return ret
 }
 
diff --git a/lang/mcache/mcache_test.go b/lang/mcache/mcache_test.go
index 964084a0..32e0f7a2 100644
--- a/lang/mcache/mcache_test.go
+++ b/lang/mcache/mcache_test.go
@@ -22,11 +22,16 @@ import (
 
 func TestMalloc(t *testing.T) {
 	// cached by mcache
-	size := 4096 - 1
-	buf := Malloc(size)
+	size := 128
+	buf := Malloc(size, size+1)
 	assert.Equal(t, len(buf), size)
 	assert.Equal(t, cap(buf), size+1)
-	_ = buf[:size+1] // resize
+
+	// cached by mcache
+	size = 4096 - 1
+	buf = Malloc(size)
+	assert.Equal(t, len(buf), size)
+	assert.Equal(t, cap(buf), size)
 	Free(buf)
 
 	// cached by mcache

From 44ee697eb646448f17c233626db7d54748e35f09 Mon Sep 17 00:00:00 2001
From: wangzhuowei <wangzhuowei@bytedance.com>
Date: Fri, 7 Jun 2024 11:39:30 +0800
Subject: [PATCH 4/4] Revert "fix: mcache Malloc should repect capacity"

This reverts commit 65caa4f9bb95db9488ce7356775c36ab86b6ad19.
---
 lang/mcache/mcache.go      |  2 +-
 lang/mcache/mcache_test.go | 11 +++--------
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/lang/mcache/mcache.go b/lang/mcache/mcache.go
index ab32ea4f..402bfb6f 100644
--- a/lang/mcache/mcache.go
+++ b/lang/mcache/mcache.go
@@ -68,7 +68,7 @@ func Malloc(size int, capacity ...int) []byte {
 
 	sizeIdx := calcIndex(c)
 	var ret = caches[sizeIdx].Get().([]byte)
-	ret = ret[:size:c]
+	ret = ret[:size]
 	return ret
 }
 
diff --git a/lang/mcache/mcache_test.go b/lang/mcache/mcache_test.go
index 32e0f7a2..964084a0 100644
--- a/lang/mcache/mcache_test.go
+++ b/lang/mcache/mcache_test.go
@@ -22,16 +22,11 @@ import (
 
 func TestMalloc(t *testing.T) {
 	// cached by mcache
-	size := 128
-	buf := Malloc(size, size+1)
+	size := 4096 - 1
+	buf := Malloc(size)
 	assert.Equal(t, len(buf), size)
 	assert.Equal(t, cap(buf), size+1)
-
-	// cached by mcache
-	size = 4096 - 1
-	buf = Malloc(size)
-	assert.Equal(t, len(buf), size)
-	assert.Equal(t, cap(buf), size)
+	_ = buf[:size+1] // resize
 	Free(buf)
 
 	// cached by mcache