Restore listening connections when netstack s/r is enabled.

This CL restores the listening connections when netstack s/r is enabled.
The changes include:
- New method as a workaround to replace the new routes and nics to the loaded
stack after restore.
- New Restore() for transport layer protocols to restore the protocol level
background workers.
- Adds afterLoad() method for fdbased processors.
- Adds a test to verify listening connection is restored after checkpointing
with netstack s/r enabled.
- Few other changes to save restore fields to enable netstack s/r.

PiperOrigin-RevId: 698453124

Makefile

@@ -353,7 +353,8 @@ docker-tests: load-basic $(RUNTIME_BIN)
 	@$(call install_runtime,$(RUNTIME)-dcache,--fdlimit=2000 --dcache=100) # Used by TestDentryCacheLimit.
 	@$(call install_runtime,$(RUNTIME)-host-uds,--host-uds=all) # Used by TestHostSocketConnect.
 	@$(call install_runtime,$(RUNTIME)-overlay,--overlay2=all:self) # Used by TestOverlay*.
-	@$(call test_runtime,$(RUNTIME),$(INTEGRATION_TARGETS) //test/e2e:integration_runtime_test //test/e2e:runtime_in_docker_test)
+	@$(call install_runtime,$(RUNTIME)-TESTONLY-save-restore-netstack,--TESTONLY-save-restore-netstack=true) # Used by TestRestoreListenConnWithNetstackSR.
+	@$(call test_runtime,$(RUNTIME),$(INTEGRATION_TARGETS) --test_env=TEST_SAVE_RESTORE_NETSTACK=true //test/e2e:integration_runtime_test //test/e2e:runtime_in_docker_test)
 .PHONY: docker-tests
 
 plugin-network-tests: load-basic $(RUNTIME_BIN)

pkg/sentry/inet/inet.go

@@ -100,6 +100,14 @@ type Stack interface {
 	// Restore restarts the network stack after restore.
 	Restore()
 
+	// ReplaceConfig replaces the new network stack configuration to the
+	// loaded or saved network stack after restore.
+	// TODO(b/379115439): This method is a workaround to update netstack config
+	// during restore. It should be removed after a new method is added to
+	// extract the complete config from the spec and update it in the loaded
+	// stack during restore.
+	ReplaceConfig(st Stack)
+
 	// Destroy the network stack.
 	Destroy()
 
@@ -126,6 +134,9 @@ type Stack interface {
 
 	// EnableSaveRestore enables netstack s/r.
 	EnableSaveRestore() error
+
+	// IsSaveRestoreEnabled returns true when netstack s/r is enabled.
+	IsSaveRestoreEnabled() bool
 }
 
 // Interface contains information about a network interface.

pkg/sentry/inet/test_stack.go

@@ -175,6 +175,9 @@ func (s *TestStack) Pause() {}
 // Restore implements Stack.
 func (s *TestStack) Restore() {}
 
+// ReplaceConfig implements Stack.
+func (s *TestStack) ReplaceConfig(_ Stack) {}
+
 // Resume implements Stack.
 func (s *TestStack) Resume() {}
 
@@ -226,3 +229,9 @@ func (*TestStack) EnableSaveRestore() error {
 	// No-op.
 	return nil
 }
+
+// IsSaveRestoreEnabled implements Stack.
+func (*TestStack) IsSaveRestoreEnabled() bool {
+	// No-op.
+	return false
+}

pkg/sentry/kernel/kernel.go

@@ -845,7 +845,17 @@ func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pages
 		close(timeReady)
 	}
 
-	if net != nil {
+	if saveRestoreNet {
+		log.Infof("netstack save restore is enabled")
+		s := k.rootNetworkNamespace.Stack()
+		if s == nil {
+			panic("inet.Stack cannot be nil when netstack s/r is enabled")
+		}
+		if net != nil {
+			s.ReplaceConfig(net)
+		}
+		s.Restore()
+	} else if net != nil {
 		net.Restore()
 	}
 

pkg/sentry/socket/hostinet/stack.go

@@ -398,6 +398,9 @@ func (*Stack) Pause() {}
 // Restore implements inet.Stack.Restore.
 func (*Stack) Restore() {}
 
+// ReplaceConfig implements inet.Stack.ReplaceConfig.
+func (s *Stack) ReplaceConfig(_ inet.Stack) {}
+
 // Resume implements inet.Stack.Resume.
 func (*Stack) Resume() {}
 
@@ -430,3 +433,8 @@ func (*Stack) SetPortRange(uint16, uint16) error {
 func (*Stack) EnableSaveRestore() error {
 	return fmt.Errorf("s/r is not supported for hostinet")
 }
+
+// IsSaveRestoreEnabled implements inet.Stack.IsSaveRestoreEnabled.
+func (s *Stack) IsSaveRestoreEnabled() bool {
+	return false
+}

pkg/sentry/socket/netstack/netstack_state.go

@@ -34,7 +34,7 @@ func (s *sock) loadTimestamp(_ context.Context, nsec int64) {
 }
 
 func (s *Stack) saveStack() *stack.Stack {
-	if s.shouldSaveRestoreStack {
+	if s.IsSaveRestoreEnabled() {
 		return s.Stack
 	}
 

pkg/sentry/socket/netstack/save_restore.go

@@ -22,7 +22,7 @@ import (
 
 // afterLoad is invoked by stateify.
 func (s *Stack) afterLoad(ctx context.Context) {
-	if s.shouldSaveRestoreStack {
+	if s.IsSaveRestoreEnabled() {
 		// This indicates that netstack s/r is enabled and the stack
 		// should not be replaced with the new stack from context.
 		return

pkg/sentry/socket/netstack/stack.go

@@ -40,16 +40,23 @@ import (
 //
 // +stateify savable
 type Stack struct {
-	Stack                  *stack.Stack `state:".(*stack.Stack)"`
-	shouldSaveRestoreStack bool
+	Stack *stack.Stack `state:".(*stack.Stack)"`
 }
 
 // EnableSaveRestore enables netstack s/r.
 func (s *Stack) EnableSaveRestore() error {
-	s.shouldSaveRestoreStack = true
+	s.Stack.EnableSaveRestore()
 	return nil
 }
 
+// IsSaveRestoreEnabled implements inet.Stack.IsSaveRestoreEnabled.
+func (s *Stack) IsSaveRestoreEnabled() bool {
+	if s.Stack == nil {
+		return false
+	}
+	return s.Stack.IsSaveRestoreEnabled()
+}
+
 // Destroy implements inet.Stack.Destroy.
 func (s *Stack) Destroy() {
 	s.Stack.Close()
@@ -912,6 +919,14 @@ func (s *Stack) Restore() {
 	s.Stack.Restore()
 }
 
+// ReplaceConfig implements inet.Stack.ReplaceConfig.
+func (s *Stack) ReplaceConfig(st inet.Stack) {
+	if _, ok := st.(*Stack); !ok {
+		panic("netstack.Stack cannot be nil when netstack s/r is enabled")
+	}
+	s.Stack.ReplaceConfig(st.(*Stack).Stack)
+}
+
 // Resume implements inet.Stack.Resume.
 func (s *Stack) Resume() {
 	s.Stack.Resume()

pkg/tcpip/link/fdbased/processors.go

@@ -18,6 +18,7 @@
 package fdbased
 
 import (
+	"context"
 	"encoding/binary"
 
 	"gvisor.dev/gvisor/pkg/rand"
@@ -126,6 +127,12 @@ func (m *processorManager) start() {
 	}
 }
 
+// afterLoad is invoked by stateify.
+func (m *processorManager) afterLoad(context.Context) {
+	m.wg.Add(len(m.processors))
+	m.start()
+}
+
 func (m *processorManager) connectionHash(cid *connectionID) uint32 {
 	var payload [4]byte
 	binary.LittleEndian.PutUint16(payload[0:], cid.srcPort)

pkg/tcpip/stack/neighbor_entry.go

@@ -29,6 +29,8 @@ const (
 )
 
 // NeighborEntry describes a neighboring device in the local network.
+//
+// +stateify savable
 type NeighborEntry struct {
 	Addr      tcpip.Address
 	LinkAddr  tcpip.LinkAddress
@@ -76,17 +78,38 @@ const (
 	Unreachable
 )
 
+// +stateify savable
 type timer struct {
 	// done indicates to the timer that the timer was stopped.
 	done *bool
 
-	timer tcpip.Timer
+	timer tcpip.Timer `state:"nosave"`
+}
+
+// +stateify savable
+type neighborEntryMu struct {
+	neighborEntryRWMutex `state:"nosave"`
+
+	neigh NeighborEntry
+
+	// done is closed when address resolution is complete. It is nil iff s is
+	// incomplete and resolution is not yet in progress.
+	done chan struct{} `state:"nosave"`
+
+	// onResolve is called with the result of address resolution.
+	onResolve []func(LinkResolutionResult) `state:"nosave"`
+
+	isRouter bool
+
+	timer timer
 }
 
 // neighborEntry implements a neighbor entry's individual node behavior, as per
 // RFC 4861 section 7.3.3. Neighbor Unreachability Detection operates in
 // parallel with the sending of packets to a neighbor, necessitating the
 // entry's lock to be acquired for all operations.
+//
+// +stateify savable
 type neighborEntry struct {
 	neighborEntryEntry
 
@@ -95,22 +118,7 @@ type neighborEntry struct {
 	// nudState points to the Neighbor Unreachability Detection configuration.
 	nudState *NUDState
 
-	mu struct {
-		neighborEntryRWMutex
-
-		neigh NeighborEntry
-
-		// done is closed when address resolution is complete. It is nil iff s is
-		// incomplete and resolution is not yet in progress.
-		done chan struct{}
-
-		// onResolve is called with the result of address resolution.
-		onResolve []func(LinkResolutionResult)
-
-		isRouter bool
-
-		timer timer
-	}
+	mu neighborEntryMu
 }
 
 // newNeighborEntry creates a neighbor cache entry starting at the default

pkg/tcpip/stack/pending_packets.go

@@ -33,9 +33,8 @@ type pendingPacket struct {
 	pkt       *PacketBuffer
 }
 
-// +stateify savable
 type packetsPendingLinkResolutionMu struct {
-	packetsPendingLinkResolutionMutex `state:"nosave"`
+	packetsPendingLinkResolutionMutex
 
 	// The packets to send once the resolver completes.
 	//
@@ -56,7 +55,7 @@ type packetsPendingLinkResolutionMu struct {
 // +stateify savable
 type packetsPendingLinkResolution struct {
 	nic *nic
-	mu  packetsPendingLinkResolutionMu
+	mu  packetsPendingLinkResolutionMu `state:"nosave"`
 }
 
 func (f *packetsPendingLinkResolution) incrementOutgoingPacketErrors(pkt *PacketBuffer) {

pkg/tcpip/stack/registration.go

@@ -244,6 +244,9 @@ type TransportProtocol interface {
 	// previously paused by Pause.
 	Resume()
 
+	// Restore starts any protocol level background workers during restore.
+	Restore()
+
 	// Parse sets pkt.TransportHeader and trims pkt.Data appropriately. It does
 	// neither and returns false if pkt.Data is too small, i.e. pkt.Data.Size() <
 	// MinimumPacketSize()

pkg/tcpip/stack/stack.go

@@ -90,16 +90,16 @@ type Stack struct {
 
 	// routeTable is a list of routes sorted by prefix length, longest (most specific) first.
 	// +checklocks:routeMu
-	routeTable tcpip.RouteList
+	routeTable tcpip.RouteList `state:"nosave"`
 
 	mu stackRWMutex `state:"nosave"`
 	// +checklocks:mu
-	nics map[tcpip.NICID]*nic
+	nics map[tcpip.NICID]*nic `state:"nosave"`
 	// +checklocks:mu
 	defaultForwardingEnabled map[tcpip.NetworkProtocolNumber]struct{}
 
 	// nicIDGen is used to generate NIC IDs.
-	nicIDGen atomicbitops.Int32
+	nicIDGen atomicbitops.Int32 `state:"nosave"`
 
 	// cleanupEndpointsMu protects cleanupEndpoints.
 	cleanupEndpointsMu cleanupEndpointsMutex `state:"nosave"`
@@ -180,6 +180,9 @@ type Stack struct {
 	// tsOffsetSecret is the secret key for generating timestamp offsets
 	// initialized at stack startup.
 	tsOffsetSecret uint32
+
+	// saveRestoreEnabled indicates whether the stack is saved and restored.
+	saveRestoreEnabled bool
 }
 
 // NetworkProtocolFactory instantiates a network protocol.
@@ -1966,6 +1969,28 @@ func (s *Stack) Pause() {
 	}
 }
 
+// ReplaceConfig replaces config in the loaded stack.
+func (s *Stack) ReplaceConfig(st *Stack) {
+	if st == nil {
+		panic("stack.Stack cannot be nil when netstack s/r is enabled")
+	}
+
+	// Update route table.
+	s.SetRouteTable(st.GetRouteTable())
+
+	// Update NICs.
+	st.mu.Lock()
+	defer st.mu.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.nics = make(map[tcpip.NICID]*nic)
+	for id, nic := range st.nics {
+		nic.stack = s
+		s.nics[id] = nic
+		_ = s.NextNICID()
+	}
+}
+
 // Restore restarts the stack after a restore. This must be called after the
 // entire system has been restored.
 func (s *Stack) Restore() {
@@ -1974,13 +1999,18 @@ func (s *Stack) Restore() {
 	s.mu.Lock()
 	eps := s.restoredEndpoints
 	s.restoredEndpoints = nil
+	saveRestoreEnabled := s.saveRestoreEnabled
 	s.mu.Unlock()
 	for _, e := range eps {
 		e.Restore(s)
 	}
 	// Now resume any protocol level background workers.
 	for _, p := range s.transportProtocols {
-		p.proto.Resume()
+		if saveRestoreEnabled {
+			p.proto.Restore()
+		} else {
+			p.proto.Resume()
+		}
 	}
 }
 
@@ -2406,3 +2436,19 @@ func (s *Stack) SetNICStack(id tcpip.NICID, peer *Stack) (tcpip.NICID, tcpip.Err
 	id = tcpip.NICID(peer.NextNICID())
 	return id, peer.CreateNICWithOptions(id, ne, NICOptions{Name: nic.Name()})
 }
+
+// EnableSaveRestore marks the saveRestoreEnabled to true.
+func (s *Stack) EnableSaveRestore() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.saveRestoreEnabled = true
+}
+
+// IsSaveRestoreEnabled returns true if save restore is enabled for the stack.
+func (s *Stack) IsSaveRestoreEnabled() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	return s.saveRestoreEnabled
+}

pkg/tcpip/stack/transport_test.go

@@ -337,6 +337,9 @@ func (*fakeTransportProtocol) Pause() {}
 // Resume implements TransportProtocol.Resume.
 func (*fakeTransportProtocol) Resume() {}
 
+// Restore implements TransportProtocol.Restore.
+func (*fakeTransportProtocol) Restore() {}
+
 // Parse implements TransportProtocol.Parse.
 func (*fakeTransportProtocol) Parse(pkt *stack.PacketBuffer) bool {
 	if _, ok := pkt.TransportHeader().Consume(fakeTransHeaderLen); ok {

pkg/tcpip/transport/icmp/protocol.go

@@ -128,6 +128,9 @@ func (*protocol) Pause() {}
 // Resume implements stack.TransportProtocol.Resume.
 func (*protocol) Resume() {}
 
+// Restore implements stack.TransportProtocol.Restore.
+func (*protocol) Restore() {}
+
 // Parse implements stack.TransportProtocol.Parse.
 func (*protocol) Parse(pkt *stack.PacketBuffer) bool {
 	// Right now, the Parse() method is tied to enabled protocols passed into

pkg/tcpip/transport/tcp/dispatcher.go

@@ -79,7 +79,7 @@ func (q *epQueue) empty() bool {
 // +stateify savable
 type processor struct {
 	epQ     epQueue
-	sleeper sleep.Sleeper
+	sleeper sleep.Sleeper `state:"nosave"`
 	// TODO(b/341946753): Restore them when netstack is savable.
 	newEndpointWaker sleep.Waker   `state:"nosave"`
 	closeWaker       sleep.Waker   `state:"nosave"`
@@ -381,9 +381,18 @@ func (d *dispatcher) init(rng *rand.Rand, nProcessors int) {
 
 	d.mu.Lock()
 	defer d.mu.Unlock()
+
 	d.closed = false
 	d.processors = make([]processor, nProcessors)
 	d.hasher = jenkinsHasher{seed: rng.Uint32()}
+	d.startLocked()
+}
+
+// +checklocks:d.mu
+func (d *dispatcher) startLocked() {
+	if d.closed {
+		return
+	}
 	for i := range d.processors {
 		p := &d.processors[i]
 		p.sleeper.AddWaker(&p.newEndpointWaker)
@@ -399,6 +408,13 @@ func (d *dispatcher) init(rng *rand.Rand, nProcessors int) {
 	}
 }
 
+func (d *dispatcher) start() {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	d.startLocked()
+}
+
 // close closes a dispatcher and its processors.
 func (d *dispatcher) close() {
 	d.mu.Lock()

pkg/tcpip/transport/tcp/endpoint.go

@@ -364,8 +364,8 @@ type Endpoint struct {
 
 	// The following fields are initialized at creation time and do not
 	// change throughout the lifetime of the endpoint.
-	stack       *stack.Stack  `state:"manual"`
-	protocol    *protocol     `state:"manual"`
+	stack       *stack.Stack
+	protocol    *protocol
 	waiterQueue *waiter.Queue `state:"wait"`
 
 	// hardError is meaningful only when state is stateError. It stores the
@@ -416,8 +416,8 @@ type Endpoint struct {
 	// state.
 	origEndpointState uint32 `state:"nosave"`
 
-	isPortReserved    bool `state:"manual"`
-	isRegistered      bool `state:"manual"`
+	isPortReserved    bool
+	isRegistered      bool
 	boundNICID        tcpip.NICID
 	route             *stack.Route `state:"manual"`
 	ipv4TTL           uint8

pkg/tcpip/transport/tcp/endpoint_state.go

@@ -24,6 +24,7 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/ports"
+	"gvisor.dev/gvisor/pkg/tcpip/seqnum"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
@@ -129,7 +130,11 @@ func (e *Endpoint) afterLoad(ctx context.Context) {
 	// Restore the endpoint to InitialState as it will be moved to
 	// its origEndpointState during Restore.
 	e.state = atomicbitops.FromUint32(uint32(StateInitial))
-	stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e)
+	if e.stack.IsSaveRestoreEnabled() {
+		e.stack.RegisterRestoredEndpoint(e)
+	} else {
+		stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e)
+	}
 }
 
 // Restore implements tcpip.RestoredEndpoint.Restore.
@@ -143,8 +148,11 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 		snd.probeTimer.init(s.Clock(), timerHandler(e, e.snd.probeTimerExpired))
 		snd.corkTimer.init(s.Clock(), timerHandler(e, e.snd.corkTimerExpired))
 	}
-	e.stack = s
-	e.protocol = protocolFromStack(s)
+	saveRestoreEnabled := e.stack.IsSaveRestoreEnabled()
+	if !saveRestoreEnabled {
+		e.stack = s
+		e.protocol = protocolFromStack(s)
+	}
 	e.ops.InitHandler(e, e.stack, GetTCPSendBufferLimits, GetTCPReceiveBufferLimits)
 	e.segmentQueue.thaw()
 
@@ -216,24 +224,35 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 		e.mu.Unlock()
 		connectedLoading.Done()
 	case epState == StateListen:
-		tcpip.AsyncLoading.Add(1)
-		go func() {
-			connectedLoading.Wait()
-			bind()
-			e.acceptMu.Lock()
-			backlog := e.acceptQueue.capacity
-			e.acceptMu.Unlock()
-			if err := e.Listen(backlog); err != nil {
-				panic("endpoint listening failed: " + err.String())
-			}
+		if !saveRestoreEnabled {
+			tcpip.AsyncLoading.Add(1)
+			go func() {
+				connectedLoading.Wait()
+				bind()
+				e.acceptMu.Lock()
+				backlog := e.acceptQueue.capacity
+				e.acceptMu.Unlock()
+				if err := e.Listen(backlog); err != nil {
+					panic("endpoint listening failed: " + err.String())
+				}
+				e.LockUser()
+				if e.shutdownFlags != 0 {
+					e.shutdownLocked(e.shutdownFlags)
+				}
+				e.UnlockUser()
+				listenLoading.Done()
+				tcpip.AsyncLoading.Done()
+			}()
+		} else {
 			e.LockUser()
-			if e.shutdownFlags != 0 {
-				e.shutdownLocked(e.shutdownFlags)
-			}
+			// All endpoints will be moved to initial state after
+			// restore. Set endpoint to its originial listen state.
+			e.setEndpointState(StateListen)
+			// Initialize the listening context.
+			rcvWnd := seqnum.Size(e.receiveBufferAvailable())
+			e.listenCtx = newListenContext(e.stack, e.protocol, e, rcvWnd, e.ops.GetV6Only(), e.NetProto)
 			e.UnlockUser()
-			listenLoading.Done()
-			tcpip.AsyncLoading.Done()
-		}()
+		}
 	case epState == StateConnecting:
 		// Initial SYN hasn't been sent yet so initiate a connect.
 		tcpip.AsyncLoading.Add(1)

pkg/tcpip/transport/tcp/protocol.go

@@ -508,6 +508,11 @@ func (p *protocol) Resume() {
 	p.dispatcher.resume()
 }
 
+// Restore implements stack.TransportProtocol.Restore.
+func (p *protocol) Restore() {
+	p.dispatcher.start()
+}
+
 // Parse implements stack.TransportProtocol.Parse.
 func (*protocol) Parse(pkt *stack.PacketBuffer) bool {
 	return parse.TCP(pkt)

pkg/tcpip/transport/udp/protocol.go

@@ -124,6 +124,9 @@ func (*protocol) Pause() {}
 // Resume implements stack.TransportProtocol.Resume.
 func (*protocol) Resume() {}
 
+// Restore implements stack.TransportProtocol.Restore.
+func (*protocol) Restore() {}
+
 // Parse implements stack.TransportProtocol.Parse.
 func (*protocol) Parse(pkt *stack.PacketBuffer) bool {
 	return parse.UDP(pkt)