test(thrift): add tests for WriteDirect

internal/testutils/netpoll/netpoll.go

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2024 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package netpoll
+
+// NetpollDirectWriter implements NocopyWriter for fastcodec.
+//
+// It's only used for testing purposes,
+// see the WriteDirect implementation of netpoll.Writer for details.
+type NetpollDirectWriter struct {
+	data []byte
+
+	// for WriteDirect
+	wbuf [][]byte
+	wend []int // remainCap
+}
+
+// Bytes returns the actual result of FastWriteNocopy
+func (p *NetpollDirectWriter) Bytes() []byte {
+	ret := make([]byte, 0, len(p.data))
+	start := 0
+	for i := 0; i < len(p.wend); i++ {
+		end := len(p.data) - p.wend[i]
+		ret = append(ret, p.data[start:end]...) // bytes from p.data
+		ret = append(ret, p.wbuf[i]...)         // bytes from WriteDirect
+		start = end
+	}
+	// copy left bytes
+	ret = append(ret, p.data[start:start+len(p.data)-len(ret)]...)
+	if len(ret) != len(p.data) {
+		panic("size not match")
+	}
+	return ret
+}
+
+// Malloc creates a new buffer for FastWriteNocopy
+func (p *NetpollDirectWriter) Malloc(n int) []byte {
+	p.wbuf = p.wbuf[:0]
+	p.wend = p.wend[:0]
+	p.data = make([]byte, n)
+	return p.data
+}
+
+// WriteDirect implements NocopyWriter for fastcodec
+func (p *NetpollDirectWriter) WriteDirect(b []byte, remainCap int) error {
+	if remainCap < len(b) {
+		panic("buffer too small to fit the input")
+	}
+	p.wbuf = append(p.wbuf, b)
+	p.wend = append(p.wend, remainCap)
+	return nil
+}
+
+func (p *NetpollDirectWriter) WriteDirectN() int {
+	return len(p.wbuf)
+}

internal/testutils/netpoll/netpoll_test.go

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2024 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package netpoll
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNetpollDirectWriter(t *testing.T) {
+	w := &NetpollDirectWriter{}
+
+	i := 0
+	b := w.Malloc(100)
+	vv := [][]byte{ // 10 x 10
+		[]byte("1234567890"),
+		[]byte("2234567890"),
+		[]byte("3234567890"),
+		[]byte("4234567890"),
+		[]byte("5234567890"),
+		[]byte("6234567890"),
+		[]byte("7234567890"),
+		[]byte("8234567890"),
+		[]byte("9234567890"),
+		[]byte("0234567890"),
+	}
+
+	for j, v := range vv {
+		if j != 0 && j%2 == 0 {
+			_ = w.WriteDirect(v, len(b)-i)
+		} else {
+			i += copy(b[i:], v)
+		}
+	}
+	expect := bytes.Join(vv, []byte{})
+	require.Equal(t, expect, w.Bytes())
+}

protocol/thrift/binary_test.go

@@ -17,8 +17,12 @@
 package thrift
 
 import (
+	"bytes"
+	"math/rand"
+	"strings"
 	"testing"
 
+	"github.com/cloudwego/gopkg/internal/testutils/netpoll"
 	"github.com/stretchr/testify/require"
 )
 
@@ -416,3 +420,34 @@ func TestBinarySkip(t *testing.T) {
 	_, err = Binary.Skip(b, TType(122))
 	require.Error(t, err)
 }
+
+func TestNocopyWrite(t *testing.T) {
+	n := 100
+	// generate strings with the size above or below nocopyWriteThreshold
+	ss := make([]string, n)
+	for i := 0; i < len(ss); i++ {
+		c := string(rune('A' + i%61)) // 'A'(65) - '}"(125)
+		// 50% possibility will larger than nocopyWriteThreshold
+		ss[i] = strings.Repeat(c, 3*nocopyWriteThreshold/4+rand.Intn(nocopyWriteThreshold/2))
+	}
+
+	// generate expected data
+	x := BinaryProtocol{}
+	expectb := make([]byte, 0, n*2*nocopyWriteThreshold)
+	for _, s := range ss {
+		expectb = x.AppendString(expectb, s)
+	}
+
+	// generate testing data
+	w := &netpoll.NetpollDirectWriter{}
+	i := 0
+	b := w.Malloc(len(expectb))
+	for _, s := range ss {
+		i += x.WriteStringNocopy(b[i:], w, s)
+	}
+	b = w.Bytes()
+	wn := w.WriteDirectN()
+	require.True(t, wn > 0 && wn != n, wn)
+	require.Equal(t, len(expectb), len(b))
+	require.True(t, bytes.Equal(expectb, b))
+}