optimize: WriteDirect implementation to avoid panic and duplicate cre…

…ation of redundant LinkBufferNode when remainLen is 0 (cloudwego#250)
firedtoad · Jul 6, 2023 · 1feecad · 1feecad
1 parent e9b72f3
commit 1feecad
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 27 deletions.
diff --git a/nocopy_linkbuffer.go b/nocopy_linkbuffer.go
@@ -475,7 +475,7 @@ func (b *LinkBuffer) WriteDirect(p []byte, remainLen int) error {
 	// find origin
 	origin := b.flush
 	malloc := b.mallocSize - remainLen // calculate the remaining malloc length
-	for t := origin.malloc - len(origin.buf); t <= malloc; t = origin.malloc - len(origin.buf) {
+	for t := origin.malloc - len(origin.buf); t < malloc; t = origin.malloc - len(origin.buf) {
 		malloc -= t
 		origin = origin.next
 	}
@@ -486,18 +486,24 @@ func (b *LinkBuffer) WriteDirect(p []byte, remainLen int) error {
 	dataNode := newLinkBufferNode(0)
 	dataNode.buf, dataNode.malloc = p[:0], n
 
-	newNode := newLinkBufferNode(0)
-	newNode.off = malloc
-	newNode.buf = origin.buf[:malloc]
-	newNode.malloc = origin.malloc
-	newNode.readonly = false
-	origin.malloc = malloc
-	origin.readonly = true
-
-	// link nodes
-	dataNode.next = newNode
-	newNode.next = origin.next
-	origin.next = dataNode
+	if remainLen > 0 {
+		newNode := newLinkBufferNode(0)
+		newNode.off = malloc
+		newNode.buf = origin.buf[:malloc]
+		newNode.malloc = origin.malloc
+		newNode.readonly = false
+		origin.malloc = malloc
+		origin.readonly = true
+
+		// link nodes
+		dataNode.next = newNode
+		newNode.next = origin.next
+		origin.next = dataNode
+	} else {
+		// link nodes
+		dataNode.next = origin.next
+		origin.next = dataNode
+	}
 
 	// adjust b.write
 	for b.write.next != nil {

diff --git a/nocopy_linkbuffer_race.go b/nocopy_linkbuffer_race.go
@@ -513,7 +513,7 @@ func (b *LinkBuffer) WriteDirect(p []byte, remainLen int) error {
 	// find origin
 	origin := b.flush
 	malloc := b.mallocSize - remainLen // calculate the remaining malloc length
-	for t := origin.malloc - len(origin.buf); t <= malloc; t = origin.malloc - len(origin.buf) {
+	for t := origin.malloc - len(origin.buf); t < malloc; t = origin.malloc - len(origin.buf) {
 		malloc -= t
 		origin = origin.next
 	}
@@ -524,18 +524,24 @@ func (b *LinkBuffer) WriteDirect(p []byte, remainLen int) error {
 	dataNode := newLinkBufferNode(0)
 	dataNode.buf, dataNode.malloc = p[:0], n
 
-	newNode := newLinkBufferNode(0)
-	newNode.off = malloc
-	newNode.buf = origin.buf[:malloc]
-	newNode.malloc = origin.malloc
-	newNode.readonly = false
-	origin.malloc = malloc
-	origin.readonly = true
-
-	// link nodes
-	dataNode.next = newNode
-	newNode.next = origin.next
-	origin.next = dataNode
+	if remainLen > 0 {
+		newNode := newLinkBufferNode(0)
+		newNode.off = malloc
+		newNode.buf = origin.buf[:malloc]
+		newNode.malloc = origin.malloc
+		newNode.readonly = false
+		origin.malloc = malloc
+		origin.readonly = true
+
+		// link nodes
+		dataNode.next = newNode
+		newNode.next = origin.next
+		origin.next = dataNode
+	} else {
+		// link nodes
+		dataNode.next = origin.next
+		origin.next = dataNode
+	}
 
 	// adjust b.write
 	for b.write.next != nil {

diff --git a/nocopy_linkbuffer_test.go b/nocopy_linkbuffer_test.go
@@ -454,9 +454,11 @@ func TestWriteDirect(t *testing.T) {
 	buf.WriteDirect([]byte("nopqrst"), 28)
 	bt[4] = 'u'
 	buf.WriteDirect([]byte("vwxyz"), 27)
+	copy(bt[5:], "abcdefghijklmnopqrstuvwxyza")
+	buf.WriteDirect([]byte("abcdefghijklmnopqrstuvwxyz"), 0)
 	buf.Flush()
 	bs := buf.Bytes()
-	str := "abcdefghijklmnopqrstuvwxyz"
+	str := "abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzaabcdefghijklmnopqrstuvwxyz"
 	for i := 0; i < len(str); i++ {
 		if bs[i] != str[i] {
 			t.Error("not equal!")