Refactor sharedmem tx.transmit to use bufferv2.

This reduces the amount of garbage produced by the methods.

PiperOrigin-RevId: 471141815

pkg/bufferv2/buffer.go

@@ -571,6 +571,11 @@ func (br *BufferReader) Close() {
 	br.b.Release()
 }
 
+// Len returns the number of bytes in the unread portion of the buffer.
+func (br *BufferReader) Len() int {
+	return int(br.b.Size())
+}
+
 // Range specifies a range of buffer.
 type Range struct {
 	begin int

pkg/tcpip/link/sharedmem/server_tx.go

@@ -20,10 +20,12 @@ package sharedmem
 import (
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/atomicbitops"
+	"gvisor.dev/gvisor/pkg/bufferv2"
 	"gvisor.dev/gvisor/pkg/cleanup"
 	"gvisor.dev/gvisor/pkg/eventfd"
 	"gvisor.dev/gvisor/pkg/tcpip/link/sharedmem/pipe"
 	"gvisor.dev/gvisor/pkg/tcpip/link/sharedmem/queue"
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
 // serverTx represents the server end of the sharedmem queue and is used to send
@@ -113,12 +115,9 @@ func (s *serverTx) cleanup() {
 
 // acquireBuffers acquires enough buffers to hold all the data in views or
 // returns nil if not enough buffers are currently available.
-func (s *serverTx) acquireBuffers(views [][]byte, buffers []queue.RxBuffer) (acquiredBuffers []queue.RxBuffer) {
+func (s *serverTx) acquireBuffers(pktBuffer bufferv2.Buffer, buffers []queue.RxBuffer) (acquiredBuffers []queue.RxBuffer) {
 	acquiredBuffers = buffers[:0]
-	wantBytes := 0
-	for i := range views {
-		wantBytes += len(views[i])
-	}
+	wantBytes := int(pktBuffer.Size())
 	for wantBytes > 0 {
 		var b []byte
 		if b = s.fillPipe.Pull(); b == nil {
@@ -137,45 +136,32 @@ func (s *serverTx) acquireBuffers(views [][]byte, buffers []queue.RxBuffer) (acq
 // well as the total number of bytes copied.
 //
 // To avoid allocations the filledBuffers are appended to the buffers slice
-// which will be grown as required.
-func (s *serverTx) fillPacket(views [][]byte, buffers []queue.RxBuffer) (filledBuffers []queue.RxBuffer, totalCopied uint32) {
-	// fillBuffer copies as much of the views as possible into the provided buffer
-	// and returns any left over views (if any).
-	fillBuffer := func(buffer *queue.RxBuffer, views [][]byte) (left [][]byte) {
-		if len(views) == 0 {
-			return nil
-		}
-		copied := uint64(0)
-		availBytes := buffer.Size
-		for availBytes > 0 && len(views) > 0 {
-			n := copy(s.data[buffer.Offset+copied:][:uint64(buffer.Size)-copied], views[0])
-			copied += uint64(n)
-			availBytes -= uint32(n)
-			views[0] = views[0][n:]
-			if len(views[0]) != 0 {
-				break
-			}
-			views = views[1:]
-		}
-		buffer.Size = uint32(copied)
-		return views
-	}
-	bufs := s.acquireBuffers(views, buffers)
+// which will be grown as required. This method takes ownership of pktBuffer.
+func (s *serverTx) fillPacket(pktBuffer bufferv2.Buffer, buffers []queue.RxBuffer) (filledBuffers []queue.RxBuffer, totalCopied uint32) {
+	bufs := s.acquireBuffers(pktBuffer, buffers)
 	if bufs == nil {
+		pktBuffer.Release()
 		return nil, 0
 	}
-	for i := 0; len(views) > 0 && i < len(bufs); i++ {
-		// Copy the packet into the posted buffer.
-		views = fillBuffer(&bufs[i], views)
-		totalCopied += bufs[i].Size
-	}
+	br := pktBuffer.AsBufferReader()
+	defer br.Close()
 
+	for i := 0; br.Len() > 0 && i < len(bufs); i++ {
+		buf := bufs[i]
+		copied, err := br.Read(s.data[buf.Offset:][:buf.Size])
+		buf.Size = uint32(copied)
+		// Copy the packet into the posted buffer.
+		totalCopied += bufs[i].Size
+		if err != nil {
+			return bufs, totalCopied
+		}
+	}
 	return bufs, totalCopied
 }
 
-func (s *serverTx) transmit(views [][]byte) bool {
+func (s *serverTx) transmit(pkt *stack.PacketBuffer) bool {
 	buffers := make([]queue.RxBuffer, 8)
-	buffers, totalCopied := s.fillPacket(views, buffers)
+	buffers, totalCopied := s.fillPacket(pkt.ToBuffer(), buffers)
 	if totalCopied == 0 {
 		// drop the packet as not enough buffers were probably available
 		// to send.

pkg/tcpip/link/sharedmem/sharedmem.go

@@ -343,11 +343,10 @@ func (e *endpoint) writePacketLocked(r stack.RouteInfo, protocol tcpip.NetworkPr
 		e.AddVirtioNetHeader(pkt)
 	}
 
-	views := pkt.AsSlices()
 	// Transmit the packet.
-	// TODO(b/231582970): Change transmit() to take a bufferv2.Buffer instead of a
-	// collection of slices.
-	ok := e.tx.transmit(views...)
+	b := pkt.ToBuffer()
+	defer b.Release()
+	ok := e.tx.transmit(b)
 	if !ok {
 		return &tcpip.ErrWouldBlock{}
 	}

pkg/tcpip/link/sharedmem/sharedmem_server.go

@@ -228,8 +228,7 @@ func (e *serverEndpoint) writePacketLocked(r stack.RouteInfo, protocol tcpip.Net
 		e.AddVirtioNetHeader(pkt)
 	}
 
-	views := pkt.AsSlices()
-	ok := e.tx.transmit(views)
+	ok := e.tx.transmit(pkt)
 	if !ok {
 		return &tcpip.ErrWouldBlock{}
 	}

pkg/tcpip/link/sharedmem/tx.go

@@ -18,6 +18,7 @@ import (
 	"math"
 
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/bufferv2"
 	"gvisor.dev/gvisor/pkg/eventfd"
 	"gvisor.dev/gvisor/pkg/tcpip/link/sharedmem/queue"
 )
@@ -92,7 +93,7 @@ func (t *tx) cleanup() {
 
 // transmit sends a packet made of bufs. Returns a boolean that specifies
 // whether the packet was successfully transmitted.
-func (t *tx) transmit(bufs ...[]byte) bool {
+func (t *tx) transmit(buffer bufferv2.Buffer) bool {
 	// Pull completions from the tx queue and add their buffers back to the
 	// pool so that we can reuse them.
 	for {
@@ -107,10 +108,7 @@ func (t *tx) transmit(bufs ...[]byte) bool {
 	}
 
 	bSize := t.bufs.entrySize
-	total := uint32(0)
-	for _, data := range bufs {
-		total += uint32(len(data))
-	}
+	total := uint32(buffer.Size())
 	bufCount := (total + bSize - 1) / bSize
 
 	// Allocate enough buffers to hold all the data.
@@ -132,17 +130,17 @@ func (t *tx) transmit(bufs ...[]byte) bool {
 	// Copy data into allocated buffers.
 	nBuf := buf
 	var dBuf []byte
-	for _, data := range bufs {
-		for len(data) > 0 {
+	buffer.Apply(func(v *bufferv2.View) {
+		for v.Size() > 0 {
 			if len(dBuf) == 0 {
 				dBuf = t.data[nBuf.Offset:][:nBuf.Size]
 				nBuf = nBuf.Next
 			}
-			n := copy(dBuf, data)
-			data = data[n:]
+			n := copy(dBuf, v.AsSlice())
+			v.TrimFront(n)
 			dBuf = dBuf[n:]
 		}
-	}
+	})
 
 	// Get an id for this packet and send it out.
 	id := t.ids.add(buf)