From 48538702c6167710ce0fb5a2b046592c900bea0f Mon Sep 17 00:00:00 2001
From: wangzhuowei <wangzhuowei@bytedance.com>
Date: Mon, 10 Jun 2024 11:52:31 +0800
Subject: [PATCH] chore: add OOM test for cross nodes peek

---
 nocopy_linkbuffer.go      | 12 ++++++++++
 nocopy_linkbuffer_test.go | 48 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/nocopy_linkbuffer.go b/nocopy_linkbuffer.go
index bfa80d38..0dbe5831 100644
--- a/nocopy_linkbuffer.go
+++ b/nocopy_linkbuffer.go
@@ -761,6 +761,18 @@ func (b *UnsafeLinkBuffer) isSingleNode(readN int) (single bool) {
 	return l >= readN
 }
 
+// memorySize return the real memory size in bytes the LinkBuffer occupied
+func (b *LinkBuffer) memorySize() (bytes int) {
+	for node := b.head; node != nil; node = node.next {
+		bytes += cap(node.buf)
+	}
+	for _, c := range b.caches {
+		bytes += cap(c)
+	}
+	bytes += cap(b.cachePeek)
+	return bytes
+}
+
 // ------------------------------------------ implement link node ------------------------------------------
 
 // newLinkBufferNode create or reuse linkBufferNode.
diff --git a/nocopy_linkbuffer_test.go b/nocopy_linkbuffer_test.go
index 8abbebe0..bf3d2cc1 100644
--- a/nocopy_linkbuffer_test.go
+++ b/nocopy_linkbuffer_test.go
@@ -19,6 +19,7 @@ package netpoll
 
 import (
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"runtime"
 	"sync/atomic"
@@ -724,6 +725,53 @@ func TestLinkBufferIndexByte(t *testing.T) {
 	}
 }
 
+func TestLinkBufferPeekOutOfMemory(t *testing.T) {
+	bufCap := 1024 * 8
+	bufNodes := 100
+	magicN := uint64(2024)
+	buf := NewLinkBuffer(bufCap)
+	MustTrue(t, buf.IsEmpty())
+	Equal(t, cap(buf.write.buf), bufCap)
+	Equal(t, buf.memorySize(), bufCap)
+
+	var p []byte
+	var err error
+	// write data that cross multi nodes
+	for n := 0; n < bufNodes; n++ {
+		p, err = buf.Malloc(bufCap)
+		MustNil(t, err)
+		Equal(t, len(p), bufCap)
+		binary.BigEndian.PutUint64(p, magicN)
+	}
+	Equal(t, buf.MallocLen(), bufCap*bufNodes)
+	buf.Flush()
+	Equal(t, buf.MallocLen(), 0)
+
+	// peak data that in single node
+	for i := 0; i < 10; i++ {
+		p, err = buf.Peek(bufCap)
+		Equal(t, binary.BigEndian.Uint64(p), magicN)
+		MustNil(t, err)
+		Equal(t, len(p), bufCap)
+		Equal(t, buf.memorySize(), bufCap*bufNodes)
+	}
+
+	// peak data that cross nodes
+	memorySize := 0
+	for i := 0; i < 1024; i++ {
+		p, err = buf.Peek(bufCap + 1)
+		MustNil(t, err)
+		Equal(t, binary.BigEndian.Uint64(p), magicN)
+		Equal(t, len(p), bufCap+1)
+		if memorySize == 0 {
+			memorySize = buf.memorySize()
+			t.Logf("after Peek: memorySize=%d", memorySize)
+		} else {
+			Equal(t, buf.memorySize(), memorySize)
+		}
+	}
+}
+
 func BenchmarkStringToSliceByte(b *testing.B) {
 	b.StopTimer()
 	s := "hello world"