From ee34fd3b9d13022226ae883c3dc728e34d3ee444 Mon Sep 17 00:00:00 2001
From: Nayana Bidari <nybidari@google.com>
Date: Fri, 16 Aug 2024 10:43:33 -0700
Subject: [PATCH] Mark more structs in netstack as savable.

PiperOrigin-RevId: 663780274
---
 pkg/rawfile/rawfile_unsafe.go                |  2 ++
 pkg/tcpip/link/fdbased/endpoint.go           |  6 +++---
 pkg/tcpip/link/fdbased/mmap.go               |  2 ++
 pkg/tcpip/link/fdbased/packet_dispatchers.go |  5 +++++
 pkg/tcpip/link/fdbased/processors.go         |  7 +++++--
 pkg/tcpip/link/qdisc/fifo/fifo.go            | 12 ++++++++----
 pkg/tcpip/link/sharedmem/pipe/pipe.go        |  2 ++
 pkg/tcpip/link/sharedmem/pipe/rx.go          |  2 ++
 pkg/tcpip/link/sharedmem/pipe/tx.go          |  2 ++
 pkg/tcpip/link/sharedmem/server_rx.go        |  1 +
 pkg/tcpip/link/sharedmem/server_tx.go        |  2 ++
 pkg/tcpip/link/sharedmem/sharedmem_server.go |  7 ++++---
 pkg/tcpip/link/stopfd/stopfd.go              |  2 ++
 pkg/tcpip/link/tun/device.go                 |  2 ++
 pkg/tcpip/stack/bridge.go                    |  5 ++++-
 pkg/tcpip/stack/gro/gro.go                   |  6 ++++++
 16 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/pkg/rawfile/rawfile_unsafe.go b/pkg/rawfile/rawfile_unsafe.go
index 4ec671ae4..9443c19ab 100644
--- a/pkg/rawfile/rawfile_unsafe.go
+++ b/pkg/rawfile/rawfile_unsafe.go
@@ -69,6 +69,8 @@ func AppendIovecFromBytes(iovs []unix.Iovec, bs []byte, max int) []unix.Iovec {
 }
 
 // MMsgHdr represents the mmsg_hdr structure required by recvmmsg() on linux.
+//
+// +stateify savable
 type MMsgHdr struct {
 	Msg unix.Msghdr
 	Len uint32
diff --git a/pkg/tcpip/link/fdbased/endpoint.go b/pkg/tcpip/link/fdbased/endpoint.go
index abab35959..93f8cc8dd 100644
--- a/pkg/tcpip/link/fdbased/endpoint.go
+++ b/pkg/tcpip/link/fdbased/endpoint.go
@@ -129,7 +129,7 @@ type endpoint struct {
 
 	// closed is a function to be called when the FD's peer (if any) closes
 	// its end of the communication pipe.
-	closed func(tcpip.Error)
+	closed func(tcpip.Error) `state:"nosave"`
 
 	inboundDispatchers []linkDispatcher
 
@@ -146,7 +146,7 @@ type endpoint struct {
 	gsoMaxSize uint32
 
 	// wg keeps track of running goroutines.
-	wg sync.WaitGroup
+	wg sync.WaitGroup `state:"nosave"`
 
 	// gsoKind is the supported kind of GSO.
 	gsoKind stack.SupportedGSO
@@ -837,7 +837,7 @@ func (*endpoint) SetOnCloseAction(func()) {}
 // InjectableEndpoint is an injectable fd-based endpoint. The endpoint writes
 // to the FD, but does not read from it. All reads come from injected packets.
 //
-// +satetify savable
+// +stateify savable
 type InjectableEndpoint struct {
 	endpoint
 
diff --git a/pkg/tcpip/link/fdbased/mmap.go b/pkg/tcpip/link/fdbased/mmap.go
index 802d36cc7..9817130d2 100644
--- a/pkg/tcpip/link/fdbased/mmap.go
+++ b/pkg/tcpip/link/fdbased/mmap.go
@@ -115,6 +115,8 @@ func (t tPacketHdr) Payload() []byte {
 
 // packetMMapDispatcher uses PACKET_RX_RING's to read/dispatch inbound packets.
 // See: mmap_amd64_unsafe.go for implementation details.
+//
+// +stateify savable
 type packetMMapDispatcher struct {
 	stopfd.StopFD
 	// fd is the file descriptor used to send and receive packets.
diff --git a/pkg/tcpip/link/fdbased/packet_dispatchers.go b/pkg/tcpip/link/fdbased/packet_dispatchers.go
index b2883b838..6e091243d 100644
--- a/pkg/tcpip/link/fdbased/packet_dispatchers.go
+++ b/pkg/tcpip/link/fdbased/packet_dispatchers.go
@@ -31,6 +31,7 @@ import (
 // BufConfig defines the shape of the buffer used to read packets from the NIC.
 var BufConfig = []int{128, 256, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768}
 
+// +stateify savable
 type iovecBuffer struct {
 	// buffer is the actual buffer that holds the packet contents. Some contents
 	// are reused across calls to pullBuffer if number of requested bytes is
@@ -145,6 +146,8 @@ func (b *iovecBuffer) release() {
 
 // readVDispatcher uses readv() system call to read inbound packets and
 // dispatches them.
+//
+// +stateify savable
 type readVDispatcher struct {
 	stopfd.StopFD
 	// fd is the file descriptor used to send and receive packets.
@@ -207,6 +210,8 @@ func (d *readVDispatcher) dispatch() (bool, tcpip.Error) {
 
 // recvMMsgDispatcher uses the recvmmsg system call to read inbound packets and
 // dispatches them.
+//
+// +stateify savable
 type recvMMsgDispatcher struct {
 	stopfd.StopFD
 	// fd is the file descriptor used to send and receive packets.
diff --git a/pkg/tcpip/link/fdbased/processors.go b/pkg/tcpip/link/fdbased/processors.go
index 86a38e153..877c5b49e 100644
--- a/pkg/tcpip/link/fdbased/processors.go
+++ b/pkg/tcpip/link/fdbased/processors.go
@@ -30,8 +30,9 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/stack/gro"
 )
 
+// +stateify savable
 type processor struct {
-	mu sync.Mutex
+	mu sync.Mutex `state:"nosave"`
 	// +checklocks:mu
 	pkts stack.PacketBufferList
 
@@ -83,10 +84,12 @@ func (p *processor) deliverPackets() {
 
 // processorManager handles starting, closing, and queuing packets on processor
 // goroutines.
+//
+// +stateify savable
 type processorManager struct {
 	processors []processor
 	seed       uint32
-	wg         sync.WaitGroup
+	wg         sync.WaitGroup `state:"nosave"`
 	e          *endpoint
 	ready      []bool
 }
diff --git a/pkg/tcpip/link/qdisc/fifo/fifo.go b/pkg/tcpip/link/qdisc/fifo/fifo.go
index eb1c22709..76d8dbe5b 100644
--- a/pkg/tcpip/link/qdisc/fifo/fifo.go
+++ b/pkg/tcpip/link/qdisc/fifo/fifo.go
@@ -41,8 +41,10 @@ const (
 // queueDispatchers. All outgoing packets are consistently hashed to a single
 // underlying queue using the PacketBuffer.Hash if set, otherwise all packets
 // are queued to the first queue to avoid reordering in case of missing hash.
+//
+// +stateify savable
 type discipline struct {
-	wg          sync.WaitGroup
+	wg          sync.WaitGroup `state:"nosave"`
 	dispatchers []queueDispatcher
 
 	closed atomicbitops.Int32
@@ -51,15 +53,17 @@ type discipline struct {
 // queueDispatcher is responsible for dispatching all outbound packets in its
 // queue. It will also smartly batch packets when possible and write them
 // through the lower LinkWriter.
+//
+// +stateify savable
 type queueDispatcher struct {
 	lower stack.LinkWriter
 
-	mu sync.Mutex
+	mu sync.Mutex `state:"nosave"`
 	// +checklocks:mu
 	queue packetBufferCircularList
 
-	newPacketWaker sleep.Waker
-	closeWaker     sleep.Waker
+	newPacketWaker sleep.Waker `state:"nosave"`
+	closeWaker     sleep.Waker `state:"nosave"`
 }
 
 // New creates a new fifo queuing discipline with the n queues with maximum
diff --git a/pkg/tcpip/link/sharedmem/pipe/pipe.go b/pkg/tcpip/link/sharedmem/pipe/pipe.go
index 74c9f0311..75a8d8fc4 100644
--- a/pkg/tcpip/link/sharedmem/pipe/pipe.go
+++ b/pkg/tcpip/link/sharedmem/pipe/pipe.go
@@ -61,6 +61,8 @@ func slotToPayloadSize(offset uint64) uint64 {
 // pipe. Indices into this pipe are split into two fields: offset, which counts
 // the number of bytes from the beginning of the buffer, and revolution, which
 // counts the number of times the index has wrapped around.
+//
+// +stateify savable
 type pipe struct {
 	buffer []byte
 }
diff --git a/pkg/tcpip/link/sharedmem/pipe/rx.go b/pkg/tcpip/link/sharedmem/pipe/rx.go
index b4d2d1e74..525701019 100644
--- a/pkg/tcpip/link/sharedmem/pipe/rx.go
+++ b/pkg/tcpip/link/sharedmem/pipe/rx.go
@@ -15,6 +15,8 @@
 package pipe
 
 // Rx is the receive side of the shared memory ring buffer.
+//
+// +stateify savable
 type Rx struct {
 	p pipe
 
diff --git a/pkg/tcpip/link/sharedmem/pipe/tx.go b/pkg/tcpip/link/sharedmem/pipe/tx.go
index 5ab0e211e..2a14d5d49 100644
--- a/pkg/tcpip/link/sharedmem/pipe/tx.go
+++ b/pkg/tcpip/link/sharedmem/pipe/tx.go
@@ -15,6 +15,8 @@
 package pipe
 
 // Tx is the transmit side of the shared memory ring buffer.
+//
+// +stateify savable
 type Tx struct {
 	p              pipe
 	maxPayloadSize uint64
diff --git a/pkg/tcpip/link/sharedmem/server_rx.go b/pkg/tcpip/link/sharedmem/server_rx.go
index 701d44284..a3c8e386a 100644
--- a/pkg/tcpip/link/sharedmem/server_rx.go
+++ b/pkg/tcpip/link/sharedmem/server_rx.go
@@ -27,6 +27,7 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/link/sharedmem/queue"
 )
 
+// +stateify savable
 type serverRx struct {
 	// packetPipe represents the receive end of the pipe that carries the packet
 	// descriptors sent by the client.
diff --git a/pkg/tcpip/link/sharedmem/server_tx.go b/pkg/tcpip/link/sharedmem/server_tx.go
index 6c94b5ec7..4704b4a24 100644
--- a/pkg/tcpip/link/sharedmem/server_tx.go
+++ b/pkg/tcpip/link/sharedmem/server_tx.go
@@ -30,6 +30,8 @@ import (
 
 // serverTx represents the server end of the sharedmem queue and is used to send
 // packets to the peer in the buffers posted by the peer in the fillPipe.
+//
+// +stateify savable
 type serverTx struct {
 	// fillPipe represents the receive end of the pipe that carries the RxBuffers
 	// posted by the peer.
diff --git a/pkg/tcpip/link/sharedmem/sharedmem_server.go b/pkg/tcpip/link/sharedmem/sharedmem_server.go
index 2321425c8..578d6c0b4 100644
--- a/pkg/tcpip/link/sharedmem/sharedmem_server.go
+++ b/pkg/tcpip/link/sharedmem/sharedmem_server.go
@@ -27,6 +27,7 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
+// +stateify savable
 type serverEndpoint struct {
 	// bufferSize is the size of each individual buffer.
 	// bufferSize is immutable.
@@ -39,7 +40,7 @@ type serverEndpoint struct {
 	stopRequested atomicbitops.Uint32
 
 	// Wait group used to indicate that all workers have stopped.
-	completed sync.WaitGroup
+	completed sync.WaitGroup `state:"nosave"`
 
 	// peerFD is an fd to the peer that can be used to detect when the peer is
 	// gone.
@@ -59,10 +60,10 @@ type serverEndpoint struct {
 
 	// onClosed is a function to be called when the FD's peer (if any) closes its
 	// end of the communication pipe.
-	onClosed func(tcpip.Error)
+	onClosed func(tcpip.Error) `state:"nosave"`
 
 	// mu protects the following fields.
-	mu sync.RWMutex
+	mu sync.RWMutex `state:"nosave"`
 
 	// tx is the transmit queue.
 	// +checklocks:mu
diff --git a/pkg/tcpip/link/stopfd/stopfd.go b/pkg/tcpip/link/stopfd/stopfd.go
index 9cada9e26..b761aad47 100644
--- a/pkg/tcpip/link/stopfd/stopfd.go
+++ b/pkg/tcpip/link/stopfd/stopfd.go
@@ -25,6 +25,8 @@ import (
 )
 
 // StopFD is an eventfd used to signal the stop of a dispatcher.
+//
+// +stateify savable
 type StopFD struct {
 	EFD int
 }
diff --git a/pkg/tcpip/link/tun/device.go b/pkg/tcpip/link/tun/device.go
index 9b0903183..13c757748 100644
--- a/pkg/tcpip/link/tun/device.go
+++ b/pkg/tcpip/link/tun/device.go
@@ -332,6 +332,8 @@ func (d *Device) WriteNotify() {
 //
 // It is ref-counted as multiple opening files can attach to the same NIC.
 // The last owner is responsible for deleting the NIC.
+//
+// +stateify savable
 type tunEndpoint struct {
 	tunEndpointRefs
 	*channel.Endpoint
diff --git a/pkg/tcpip/stack/bridge.go b/pkg/tcpip/stack/bridge.go
index 72cd5913f..50c8f9640 100644
--- a/pkg/tcpip/stack/bridge.go
+++ b/pkg/tcpip/stack/bridge.go
@@ -22,6 +22,7 @@ import (
 
 var _ NetworkLinkEndpoint = (*BridgeEndpoint)(nil)
 
+// +stateify savable
 type bridgePort struct {
 	bridge *BridgeEndpoint
 	nic    *nic
@@ -75,8 +76,10 @@ func NewBridgeEndpoint(mtu uint32) *BridgeEndpoint {
 }
 
 // BridgeEndpoint is a bridge endpoint.
+//
+// +stateify savable
 type BridgeEndpoint struct {
-	mu bridgeRWMutex
+	mu bridgeRWMutex `state:"nosave"`
 	// +checklocks:mu
 	ports map[tcpip.NICID]*bridgePort
 	// +checklocks:mu
diff --git a/pkg/tcpip/stack/gro/gro.go b/pkg/tcpip/stack/gro/gro.go
index 89856d9a1..2276d2874 100644
--- a/pkg/tcpip/stack/gro/gro.go
+++ b/pkg/tcpip/stack/gro/gro.go
@@ -50,6 +50,8 @@ const (
 )
 
 // A groBucket holds packets that are undergoing GRO.
+//
+// +stateify savable
 type groBucket struct {
 	// count is the number of packets in the bucket.
 	count int
@@ -265,6 +267,8 @@ func (gb *groBucket) found(gd *GRO, groPkt *groPacket, flushGROPkt bool, pkt *st
 
 // A groPacket is packet undergoing GRO. It may be several packets coalesced
 // together.
+//
+// +stateify savable
 type groPacket struct {
 	// groPacketEntry is an intrusive list.
 	groPacketEntry
@@ -303,6 +307,8 @@ func (pk *groPacket) payloadSize() int {
 }
 
 // GRO coalesces incoming packets to increase throughput.
+//
+// +stateify savable
 type GRO struct {
 	enabled bool
 	buckets [groNBuckets]groBucket