Internal change.

PiperOrigin-RevId: 381375705

3 files changed, 198 insertions, 71 deletions

diff --git a/pkg/tcpip/link/fdbased/BUILD b/pkg/tcpip/link/fdbased/BUILD
index d971194e6..1d0163823 100644
--- a/pkg/tcpip/link/fdbased/BUILD
+++ b/pkg/tcpip/link/fdbased/BUILD
@@ -14,7 +14,6 @@ go_library(
     ],
     visibility = ["//visibility:public"],
     deps = [
-        "//pkg/iovec",
         "//pkg/sync",
         "//pkg/tcpip",
         "//pkg/tcpip/buffer",
diff --git a/pkg/tcpip/link/fdbased/endpoint.go b/pkg/tcpip/link/fdbased/endpoint.go
index 735c28da1..1b56d2b72 100644
--- a/pkg/tcpip/link/fdbased/endpoint.go
+++ b/pkg/tcpip/link/fdbased/endpoint.go
@@ -44,7 +44,6 @@ import (
 	"sync/atomic"
 
 	"golang.org/x/sys/unix"
-	"gvisor.dev/gvisor/pkg/iovec"
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/buffer"
@@ -138,6 +137,20 @@ type endpoint struct {
 
 	// gsoKind is the supported kind of GSO.
 	gsoKind stack.SupportedGSO
+
+	// maxSyscallHeaderBytes has the same meaning as
+	// Options.MaxSyscallHeaderBytes.
+	maxSyscallHeaderBytes uintptr
+
+	// writevMaxIovs is the maximum number of iovecs that may be passed to
+	// rawfile.NonBlockingWriteIovec, as possibly limited by
+	// maxSyscallHeaderBytes. (No analogous limit is defined for
+	// rawfile.NonBlockingSendMMsg, since in that case the maximum number of
+	// iovecs also depends on the number of mmsghdrs. Instead, if sendBatch
+	// encounters a packet whose iovec count is limited by
+	// maxSyscallHeaderBytes, it falls back to writing the packet using writev
+	// via WritePacket.)
+	writevMaxIovs int
 }
 
 // Options specify the details about the fd-based endpoint to be created.
@@ -186,6 +199,11 @@ type Options struct {
 	// RXChecksumOffload if true, indicates that this endpoints capability
 	// set should include CapabilityRXChecksumOffload.
 	RXChecksumOffload bool
+
+	// If MaxSyscallHeaderBytes is non-zero, it is the maximum number of bytes
+	// of struct iovec, msghdr, and mmsghdr that may be passed by each host
+	// system call.
+	MaxSyscallHeaderBytes int
 }
 
 // fanoutID is used for AF_PACKET based endpoints to enable PACKET_FANOUT
@@ -235,14 +253,25 @@ func New(opts *Options) (stack.LinkEndpoint, error) {
 		return nil, fmt.Errorf("opts.FD is empty, at least one FD must be specified")
 	}
 
+	if opts.MaxSyscallHeaderBytes < 0 {
+		return nil, fmt.Errorf("opts.MaxSyscallHeaderBytes is negative")
+	}
+
 	e := &endpoint{
-		fds:                opts.FDs,
-		mtu:                opts.MTU,
-		caps:               caps,
-		closed:             opts.ClosedFunc,
-		addr:               opts.Address,
-		hdrSize:            hdrSize,
-		packetDispatchMode: opts.PacketDispatchMode,
+		fds:                   opts.FDs,
+		mtu:                   opts.MTU,
+		caps:                  caps,
+		closed:                opts.ClosedFunc,
+		addr:                  opts.Address,
+		hdrSize:               hdrSize,
+		packetDispatchMode:    opts.PacketDispatchMode,
+		maxSyscallHeaderBytes: uintptr(opts.MaxSyscallHeaderBytes),
+		writevMaxIovs:         rawfile.MaxIovs,
+	}
+	if e.maxSyscallHeaderBytes != 0 {
+		if max := int(e.maxSyscallHeaderBytes / rawfile.SizeofIovec); max < e.writevMaxIovs {
+			e.writevMaxIovs = max
+		}
 	}
 
 	// Increment fanoutID to ensure that we don't re-use the same fanoutID for
@@ -470,9 +499,8 @@ func (e *endpoint) WritePacket(r stack.RouteInfo, protocol tcpip.NetworkProtocol
 		e.AddHeader(r.LocalLinkAddress, r.RemoteLinkAddress, protocol, pkt)
 	}
 
-	var builder iovec.Builder
-
 	fd := e.fds[pkt.Hash%uint32(len(e.fds))]
+	var vnetHdrBuf []byte
 	if e.gsoKind == stack.HWGSOSupported {
 		vnetHdr := virtioNetHdr{}
 		if pkt.GSOOptions.Type != stack.GSONone {
@@ -494,71 +522,123 @@ func (e *endpoint) WritePacket(r stack.RouteInfo, protocol tcpip.NetworkProtocol
 				vnetHdr.gsoSize = pkt.GSOOptions.MSS
 			}
 		}
+		vnetHdrBuf = vnetHdr.marshal()
+	}
 
-		vnetHdrBuf := vnetHdr.marshal()
-		builder.Add(vnetHdrBuf)
+	views := pkt.Views()
+	numIovecs := len(views)
+	if len(vnetHdrBuf) != 0 {
+		numIovecs++
+	}
+	if numIovecs > e.writevMaxIovs {
+		numIovecs = e.writevMaxIovs
 	}
 
-	for _, v := range pkt.Views() {
-		builder.Add(v)
+	// Allocate small iovec arrays on the stack.
+	var iovecsArr [8]unix.Iovec
+	iovecs := iovecsArr[:0]
+	if numIovecs > len(iovecsArr) {
+		iovecs = make([]unix.Iovec, 0, numIovecs)
+	}
+	iovecs = rawfile.AppendIovecFromBytes(iovecs, vnetHdrBuf, numIovecs)
+	for _, v := range views {
+		iovecs = rawfile.AppendIovecFromBytes(iovecs, v, numIovecs)
 	}
-	return rawfile.NonBlockingWriteIovec(fd, builder.Build())
+	return rawfile.NonBlockingWriteIovec(fd, iovecs)
 }
 
-func (e *endpoint) sendBatch(batchFD int, batch []*stack.PacketBuffer) (int, tcpip.Error) {
+func (e *endpoint) sendBatch(batchFD int, pkts []*stack.PacketBuffer) (int, tcpip.Error) {
 	// Send a batch of packets through batchFD.
-	mmsgHdrs := make([]rawfile.MMsgHdr, 0, len(batch))
-	for _, pkt := range batch {
-		if e.hdrSize > 0 {
-			e.AddHeader(pkt.EgressRoute.LocalLinkAddress, pkt.EgressRoute.RemoteLinkAddress, pkt.NetworkProtocolNumber, pkt)
-		}
+	mmsgHdrsStorage := make([]rawfile.MMsgHdr, 0, len(pkts))
+	packets := 0
+	for packets < len(pkts) {
+		mmsgHdrs := mmsgHdrsStorage
+		batch := pkts[packets:]
+		syscallHeaderBytes := uintptr(0)
+		for _, pkt := range batch {
+			if e.hdrSize > 0 {
+				e.AddHeader(pkt.EgressRoute.LocalLinkAddress, pkt.EgressRoute.RemoteLinkAddress, pkt.NetworkProtocolNumber, pkt)
+			}
 
-		var vnetHdrBuf []byte
-		if e.gsoKind == stack.HWGSOSupported {
-			vnetHdr := virtioNetHdr{}
-			if pkt.GSOOptions.Type != stack.GSONone {
-				vnetHdr.hdrLen = uint16(pkt.HeaderSize())
-				if pkt.GSOOptions.NeedsCsum {
-					vnetHdr.flags = _VIRTIO_NET_HDR_F_NEEDS_CSUM
-					vnetHdr.csumStart = header.EthernetMinimumSize + pkt.GSOOptions.L3HdrLen
-					vnetHdr.csumOffset = pkt.GSOOptions.CsumOffset
-				}
-				if pkt.GSOOptions.Type != stack.GSONone && uint16(pkt.Data().Size()) > pkt.GSOOptions.MSS {
-					switch pkt.GSOOptions.Type {
-					case stack.GSOTCPv4:
-						vnetHdr.gsoType = _VIRTIO_NET_HDR_GSO_TCPV4
-					case stack.GSOTCPv6:
-						vnetHdr.gsoType = _VIRTIO_NET_HDR_GSO_TCPV6
-					default:
-						panic(fmt.Sprintf("Unknown gso type: %v", pkt.GSOOptions.Type))
+			var vnetHdrBuf []byte
+			if e.gsoKind == stack.HWGSOSupported {
+				vnetHdr := virtioNetHdr{}
+				if pkt.GSOOptions.Type != stack.GSONone {
+					vnetHdr.hdrLen = uint16(pkt.HeaderSize())
+					if pkt.GSOOptions.NeedsCsum {
+						vnetHdr.flags = _VIRTIO_NET_HDR_F_NEEDS_CSUM
+						vnetHdr.csumStart = header.EthernetMinimumSize + pkt.GSOOptions.L3HdrLen
+						vnetHdr.csumOffset = pkt.GSOOptions.CsumOffset
+					}
+					if pkt.GSOOptions.Type != stack.GSONone && uint16(pkt.Data().Size()) > pkt.GSOOptions.MSS {
+						switch pkt.GSOOptions.Type {
+						case stack.GSOTCPv4:
+							vnetHdr.gsoType = _VIRTIO_NET_HDR_GSO_TCPV4
+						case stack.GSOTCPv6:
+							vnetHdr.gsoType = _VIRTIO_NET_HDR_GSO_TCPV6
+						default:
+							panic(fmt.Sprintf("Unknown gso type: %v", pkt.GSOOptions.Type))
+						}
+						vnetHdr.gsoSize = pkt.GSOOptions.MSS
 					}
-					vnetHdr.gsoSize = pkt.GSOOptions.MSS
 				}
+				vnetHdrBuf = vnetHdr.marshal()
 			}
-			vnetHdrBuf = vnetHdr.marshal()
-		}
 
-		var builder iovec.Builder
-		builder.Add(vnetHdrBuf)
-		for _, v := range pkt.Views() {
-			builder.Add(v)
-		}
-		iovecs := builder.Build()
+			views := pkt.Views()
+			numIovecs := len(views)
+			if len(vnetHdrBuf) != 0 {
+				numIovecs++
+			}
+			if numIovecs > rawfile.MaxIovs {
+				numIovecs = rawfile.MaxIovs
+			}
+			if e.maxSyscallHeaderBytes != 0 {
+				syscallHeaderBytes += rawfile.SizeofMMsgHdr + uintptr(numIovecs)*rawfile.SizeofIovec
+				if syscallHeaderBytes > e.maxSyscallHeaderBytes {
+					// We can't fit this packet into this call to sendmmsg().
+					// We could potentially do so if we reduced numIovecs
+					// further, but this might incur considerable extra
+					// copying. Leave it to the next batch instead.
+					break
+				}
+			}
 
-		var mmsgHdr rawfile.MMsgHdr
-		mmsgHdr.Msg.Iov = &iovecs[0]
-		mmsgHdr.Msg.SetIovlen((len(iovecs)))
-		mmsgHdrs = append(mmsgHdrs, mmsgHdr)
-	}
+			// We can't easily allocate iovec arrays on the stack here since
+			// they will escape this loop iteration via mmsgHdrs.
+			iovecs := make([]unix.Iovec, 0, numIovecs)
+			iovecs = rawfile.AppendIovecFromBytes(iovecs, vnetHdrBuf, numIovecs)
+			for _, v := range views {
+				iovecs = rawfile.AppendIovecFromBytes(iovecs, v, numIovecs)
+			}
 
-	packets := 0
-	for len(mmsgHdrs) > 0 {
-		sent, err := rawfile.NonBlockingSendMMsg(batchFD, mmsgHdrs)
-		if err != nil {
-			return packets, err
+			var mmsgHdr rawfile.MMsgHdr
+			mmsgHdr.Msg.Iov = &iovecs[0]
+			mmsgHdr.Msg.SetIovlen(len(iovecs))
+			mmsgHdrs = append(mmsgHdrs, mmsgHdr)
+		}
+
+		if len(mmsgHdrs) == 0 {
+			// We can't fit batch[0] into a mmsghdr while staying under
+			// e.maxSyscallHeaderBytes. Use WritePacket, which will avoid the
+			// mmsghdr (by using writev) and re-buffer iovecs more aggressively
+			// if necessary (by using e.writevMaxIovs instead of
+			// rawfile.MaxIovs).
+			pkt := batch[0]
+			if err := e.WritePacket(pkt.EgressRoute, pkt.NetworkProtocolNumber, pkt); err != nil {
+				return packets, err
+			}
+			packets++
+		} else {
+			for len(mmsgHdrs) > 0 {
+				sent, err := rawfile.NonBlockingSendMMsg(batchFD, mmsgHdrs)
+				if err != nil {
+					return packets, err
+				}
+				packets += sent
+				mmsgHdrs = mmsgHdrs[sent:]
+			}
 		}
-		packets += sent
-		mmsgHdrs = mmsgHdrs[sent:]
 	}
 
 	return packets, nil
@@ -676,8 +756,9 @@ func NewInjectable(fd int, mtu uint32, capabilities stack.LinkEndpointCapabiliti
 	unix.SetNonblock(fd, true)
 
 	return &InjectableEndpoint{endpoint: endpoint{
-		fds:  []int{fd},
-		mtu:  mtu,
-		caps: capabilities,
+		fds:           []int{fd},
+		mtu:           mtu,
+		caps:          capabilities,
+		writevMaxIovs: rawfile.MaxIovs,
 	}}
 }
diff --git a/pkg/tcpip/link/rawfile/rawfile_unsafe.go b/pkg/tcpip/link/rawfile/rawfile_unsafe.go
index ba92aedbc..43fe57830 100644
--- a/pkg/tcpip/link/rawfile/rawfile_unsafe.go
+++ b/pkg/tcpip/link/rawfile/rawfile_unsafe.go
@@ -19,12 +19,66 @@
 package rawfile
 
 import (
+	"reflect"
 	"unsafe"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/tcpip"
 )
 
+// SizeofIovec is the size of a unix.Iovec in bytes.
+const SizeofIovec = unsafe.Sizeof(unix.Iovec{})
+
+// MaxIovs is UIO_MAXIOV, the maximum number of iovecs that may be passed to a
+// host system call in a single array.
+const MaxIovs = 1024
+
+// IovecFromBytes returns a unix.Iovec representing bs.
+//
+// Preconditions: len(bs) > 0.
+func IovecFromBytes(bs []byte) unix.Iovec {
+	iov := unix.Iovec{
+		Base: &bs[0],
+	}
+	iov.SetLen(len(bs))
+	return iov
+}
+
+func bytesFromIovec(iov unix.Iovec) (bs []byte) {
+	sh := (*reflect.SliceHeader)(unsafe.Pointer(&bs))
+	sh.Data = uintptr(unsafe.Pointer(iov.Base))
+	sh.Len = int(iov.Len)
+	sh.Cap = int(iov.Len)
+	return
+}
+
+// AppendIovecFromBytes returns append(iovs, IovecFromBytes(bs)). If len(bs) ==
+// 0, AppendIovecFromBytes returns iovs without modification. If len(iovs) >=
+// max, AppendIovecFromBytes replaces the final iovec in iovs with one that
+// also includes the contents of bs. Note that this implies that
+// AppendIovecFromBytes is only usable when the returned iovec slice is used as
+// the source of a write.
+func AppendIovecFromBytes(iovs []unix.Iovec, bs []byte, max int) []unix.Iovec {
+	if len(bs) == 0 {
+		return iovs
+	}
+	if len(iovs) < max {
+		return append(iovs, IovecFromBytes(bs))
+	}
+	iovs[len(iovs)-1] = IovecFromBytes(append(bytesFromIovec(iovs[len(iovs)-1]), bs...))
+	return iovs
+}
+
+// MMsgHdr represents the mmsg_hdr structure required by recvmmsg() on linux.
+type MMsgHdr struct {
+	Msg unix.Msghdr
+	Len uint32
+	_   [4]byte
+}
+
+// SizeofMMsgHdr is the size of a MMsgHdr in bytes.
+const SizeofMMsgHdr = unsafe.Sizeof(MMsgHdr{})
+
 // GetMTU determines the MTU of a network interface device.
 func GetMTU(name string) (uint32, error) {
 	fd, err := unix.Socket(unix.AF_UNIX, unix.SOCK_DGRAM, 0)
@@ -137,13 +191,6 @@ func BlockingReadv(fd int, iovecs []unix.Iovec) (int, tcpip.Error) {
 	}
 }
 
-// MMsgHdr represents the mmsg_hdr structure required by recvmmsg() on linux.
-type MMsgHdr struct {
-	Msg unix.Msghdr
-	Len uint32
-	_   [4]byte
-}
-
 // BlockingRecvMMsg reads from a file descriptor that is set up as non-blocking
 // and stores the received messages in a slice of MMsgHdr structures. If no data
 // is available, it will block in a poll() syscall until the file descriptor