Fix vectorised writers

Author: nekohasekai

common/bufio/vectorised_test.go

@@ -3,6 +3,7 @@ package bufio
 import (
 	"crypto/rand"
 	"io"
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -26,6 +27,8 @@ func TestWriteVectorised(t *testing.T) {
 	finish := Timeout(t)
 	_, err = WriteVectorised(vectorisedWriter, [][]byte{bufA[:], bufB[:]})
 	require.NoError(t, err)
+	_, err = WriteVectorised(vectorisedWriter, [][]byte{})
+	require.Error(t, err, os.ErrInvalid)
 	output := make([]byte, 2048)
 	_, err = io.ReadFull(outputConn, output)
 	finish()

common/bufio/vectorised_unix.go

@@ -25,9 +25,14 @@ func (w *SyscallVectorisedWriter) WriteVectorised(buffers []*buf.Buffer) error {
 	defer buf.ReleaseMulti(buffers)
 	iovecList := w.iovecList
 	for _, buffer := range buffers {
+		if buffer.IsEmpty() {
+			continue
+		}
 		iovecList = append(iovecList, buffer.Iovec(buffer.Len()))
 	}
-	if cap(iovecList) > cap(w.iovecList) {
+	if len(iovecList) == 0 {
+		return os.ErrInvalid
+	} else if cap(iovecList) > cap(w.iovecList) {
 		w.iovecList = iovecList[:0]
 	}
 	var innerErr unix.Errno
@@ -37,16 +42,14 @@ func (w *SyscallVectorisedWriter) WriteVectorised(buffers []*buf.Buffer) error {
 			var r0 uintptr
 			//nolint:staticcheck
 			r0, _, innerErr = unix.RawSyscall(unix.SYS_WRITEV, fd, uintptr(unsafe.Pointer(&writeIovecList[0])), uintptr(len(writeIovecList)))
-			writeN := int(r0)
-			for writeN > 0 {
-				if buffers[0].Len() > writeN {
-					buffers[0].Advance(writeN)
-					writeIovecList[0] = buffers[0].Iovec(buffers[0].Len())
+			writeN := uint64(r0)
+			for writeN > 0 && len(writeIovecList) > 0 {
+				if uint64(writeIovecList[0].Len) > writeN {
+					writeIovecList[0].Base = (*byte)(unsafe.Add(unsafe.Pointer(writeIovecList[0].Base), writeN))
+					writeIovecList[0].SetLen(int(uint64(writeIovecList[0].Len) - writeN))
 					break
 				} else {
-					writeN -= buffers[0].Len()
-					buffers[0].Release()
-					buffers = buffers[1:]
+					writeN -= uint64(writeIovecList[0].Len)
 					writeIovecList = writeIovecList[1:]
 				}
 			}
@@ -69,6 +72,9 @@ func (w *SyscallVectorisedPacketWriter) WriteVectorisedPacket(buffers []*buf.Buf
 	defer buf.ReleaseMulti(buffers)
 	iovecList := w.iovecList
 	for _, buffer := range buffers {
+		if buffer.IsEmpty() {
+			continue
+		}
 		iovecList = append(iovecList, buffer.Iovec(buffer.Len()))
 	}
 	if cap(iovecList) > cap(w.iovecList) {

common/bufio/vectorised_windows.go

@@ -45,15 +45,22 @@ func (w *SyscallVectorisedPacketWriter) WriteVectorisedPacket(buffers []*buf.Buf
 	defer buf.ReleaseMulti(buffers)
 	iovecList := w.iovecList
 	for _, buffer := range buffers {
+		if buffer.IsEmpty() {
+			continue
+		}
 		iovecList = append(iovecList, buffer.Iovec(buffer.Len()))
 	}
 	var n uint32
 	var innerErr error
 	err := w.rawConn.Write(func(fd uintptr) (done bool) {
 		name, nameLen := ToSockaddr(destination.AddrPort())
+		var bufs *windows.WSABuf
+		if len(iovecList) > 0 {
+			bufs = &iovecList[0]
+		}
 		innerErr = windows.WSASendTo(
 			windows.Handle(fd),
-			&iovecList[0],
+			bufs,
 			uint32(len(iovecList)),
 			&n,
 			0,