Modify udpPacket to hold a PacketBuffer reference instead of a VectorizedView.

PiperOrigin-RevId: 411896048

pkg/tcpip/stack/packet_buffer.go

@@ -15,6 +15,7 @@ package stack
 
 import (
 	"fmt"
+	"io"
 
 	"gvisor.dev/gvisor/pkg/buffer"
 	"gvisor.dev/gvisor/pkg/sync"
@@ -91,6 +92,8 @@ type PacketBufferOptions struct {
 // `consumed` value is stored for each header, and it gets incremented by the
 // consumed length. PacketBuffer adds this value to `reserved` to compute the
 // starting offset of each header in `buf`.
+//
+// +stateify savable
 type PacketBuffer struct {
 	_ sync.NoCopy
 
@@ -432,6 +435,8 @@ func (pk *PacketBufferList) DecRef() {
 }
 
 // headerInfo stores metadata about a header in a packet.
+//
+// +stateify savable
 type headerInfo struct {
 	// offset is the offset of the header in pk.buf relative to
 	// pk.buf[pk.reserved]. See the PacketBuffer struct for details.
@@ -469,6 +474,8 @@ func (h PacketHeader) Consume(size int) (v tcpipbuffer.View, consumed bool) {
 }
 
 // PacketData represents the data portion of a PacketBuffer.
+//
+// +stateify savable
 type PacketData struct {
 	pk *PacketBuffer
 }
@@ -489,6 +496,28 @@ func (d PacketData) Consume(size int) (tcpipbuffer.View, bool) {
 	return v, ok
 }
 
+// ReadTo reads bytes from d to dst. It also removes these bytes from d
+// unless peek is true.
+func (d PacketData) ReadTo(dst io.Writer, peek bool) (int, error) {
+	var err error
+	done := 0
+	for _, v := range d.Views() {
+		var n int
+		n, err = dst.Write(v)
+		done += n
+		if err != nil {
+			break
+		}
+		if n != len(v) {
+			panic(fmt.Sprintf("io.Writer.Write succeeded with incomplete write: %d != %d", n, len(v)))
+		}
+	}
+	if !peek {
+		d.pk.buf.TrimFront(int64(done))
+	}
+	return done, err
+}
+
 // CapLength reduces d to at most length bytes.
 func (d PacketData) CapLength(length int) {
 	if length < 0 {

pkg/tcpip/stack/registration.go

@@ -51,6 +51,8 @@ type TransportEndpointID struct {
 }
 
 // NetworkPacketInfo holds information about a network layer packet.
+//
+// +stateify savable
 type NetworkPacketInfo struct {
 	// LocalAddressBroadcast is true if the packet's local address is a broadcast
 	// address.

pkg/tcpip/stack/route.go

@@ -52,6 +52,7 @@ type Route struct {
 	linkRes *linkResolver
 }
 
+// +stateify savable
 type routeInfo struct {
 	RemoteAddress tcpip.Address
 
@@ -97,6 +98,8 @@ func (r *Route) Loop() PacketLooping {
 }
 
 // RouteInfo contains all of Route's exported fields.
+//
+// +stateify savable
 type RouteInfo struct {
 	routeInfo
 

pkg/tcpip/transport/udp/BUILD

@@ -27,6 +27,7 @@ go_library(
     imports = ["gvisor.dev/gvisor/pkg/tcpip/buffer"],
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/buffer",
         "//pkg/sleep",
         "//pkg/sync",
         "//pkg/tcpip",

pkg/tcpip/transport/udp/endpoint.go

@@ -37,8 +37,8 @@ type udpPacket struct {
 	senderAddress      tcpip.FullAddress
 	destinationAddress tcpip.FullAddress
 	packetInfo         tcpip.IPPacketInfo
-	data               buffer.VectorisedView `state:".(buffer.VectorisedView)"`
-	receivedAt         time.Time             `state:".(int64)"`
+	pkt                *stack.PacketBuffer
+	receivedAt         time.Time `state:".(int64)"`
 	// tos stores either the receiveTOS or receiveTClass value.
 	tos uint8
 }
@@ -191,6 +191,7 @@ func (e *endpoint) Close() {
 	for !e.rcvList.Empty() {
 		p := e.rcvList.Front()
 		e.rcvList.Remove(p)
+		p.pkt.DecRef()
 	}
 	e.rcvMu.Unlock()
 
@@ -226,7 +227,8 @@ func (e *endpoint) Read(dst io.Writer, opts tcpip.ReadOptions) (tcpip.ReadResult
 	p := e.rcvList.Front()
 	if !opts.Peek {
 		e.rcvList.Remove(p)
-		e.rcvBufSize -= p.data.Size()
+		defer p.pkt.DecRef()
+		e.rcvBufSize -= p.pkt.Data().Size()
 	}
 	e.rcvMu.Unlock()
 
@@ -272,14 +274,14 @@ func (e *endpoint) Read(dst io.Writer, opts tcpip.ReadOptions) (tcpip.ReadResult
 
 	// Read Result
 	res := tcpip.ReadResult{
-		Total:           p.data.Size(),
+		Total:           p.pkt.Data().Size(),
 		ControlMessages: cm,
 	}
 	if opts.NeedRemoteAddr {
 		res.RemoteAddr = p.senderAddress
 	}
 
-	n, err := p.data.ReadTo(dst, opts.Peek)
+	n, err := p.pkt.Data().ReadTo(dst, opts.Peek)
 	if n == 0 && err != nil {
 		return res, &tcpip.ErrBadBuffer{}
 	}
@@ -513,7 +515,7 @@ func (e *endpoint) GetSockOptInt(opt tcpip.SockOptInt) (int, tcpip.Error) {
 		e.rcvMu.Lock()
 		if !e.rcvList.Empty() {
 			p := e.rcvList.Front()
-			v = p.data.Size()
+			v = p.pkt.Data().Size()
 		}
 		e.rcvMu.Unlock()
 		return v, nil
@@ -917,10 +919,11 @@ func (e *endpoint) HandlePacket(id stack.TransportEndpointID, pkt *stack.PacketB
 			Addr: id.LocalAddress,
 			Port: hdr.DestinationPort(),
 		},
-		data: pkt.Data().ExtractVV(),
+		pkt: pkt,
 	}
+	pkt.IncRef()
 	e.rcvList.PushBack(packet)
-	e.rcvBufSize += packet.data.Size()
+	e.rcvBufSize += pkt.Data().Size()
 
 	// Save any useful information from the network header to the packet.
 	switch pkt.NetworkProtocolNumber {

pkg/tcpip/transport/udp/endpoint_state.go

@@ -19,7 +19,6 @@ import (
 	"time"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
-	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"gvisor.dev/gvisor/pkg/tcpip/transport"
 )
@@ -34,16 +33,6 @@ func (p *udpPacket) loadReceivedAt(nsec int64) {
 	p.receivedAt = time.Unix(0, nsec)
 }
 
-// saveData saves udpPacket.data field.
-func (p *udpPacket) saveData() buffer.VectorisedView {
-	return p.data.Clone(nil)
-}
-
-// loadData loads udpPacket.data field.
-func (p *udpPacket) loadData(data buffer.VectorisedView) {
-	p.data = data
-}
-
 // afterLoad is invoked by stateify.
 func (e *endpoint) afterLoad() {
 	stack.StackFromEnv.RegisterRestoredEndpoint(e)