Move connection termination from beforeSave() to restore.

pkg/tcpip/faketime/faketime.go

@@ -25,6 +25,8 @@ import (
 )
 
 // NullClock implements a clock that never advances.
+//
+// +stateify savable
 type NullClock struct{}
 
 var _ tcpip.Clock = (*NullClock)(nil)
@@ -92,8 +94,9 @@ func (n *notificationChannels) wait() {
 	}
 }
 
+// +stateify savable
 type manualClockMutex struct {
-	sync.RWMutex
+	sync.RWMutex `state:"nosave"`
 
 	// now is the current (fake) time of the clock.
 	now time.Time
@@ -102,7 +105,7 @@ type manualClockMutex struct {
 	times timeHeap
 
 	// timers holds the timers scheduled for each time.
-	timers map[time.Time]map[*manualTimer]struct{}
+	timers map[time.Time]map[*manualTimer]struct{} `state:"nosave"`
 }
 
 // ManualClock implements tcpip.Clock and only advances manually with Advance

pkg/tcpip/stack/addressable_endpoint_state.go

@@ -738,6 +738,8 @@ func (a *AddressableEndpointState) Cleanup() {
 var _ AddressEndpoint = (*addressState)(nil)
 
 // addressState holds state for an address.
+//
+// +stateify savable
 type addressState struct {
 	addressableEndpointState *AddressableEndpointState
 	addr                     tcpip.AddressWithPrefix
@@ -748,7 +750,7 @@ type addressState struct {
 	//
 	// AddressableEndpointState.mu
 	//   addressState.mu
-	mu   addressStateRWMutex
+	mu   addressStateRWMutex `state:"nosave"`
 	refs addressStateRefs
 	// checklocks:mu
 	kind AddressKind

pkg/tcpip/stack/route.go

@@ -24,6 +24,8 @@ import (
 // Route represents a route through the networking stack to a given destination.
 //
 // It is safe to call Route's methods from multiple goroutines.
+//
+// +stateify savable
 type Route struct {
 	routeInfo routeInfo
 
@@ -33,7 +35,7 @@ type Route struct {
 	localAddressNIC *nic
 
 	// mu protects annotated fields below.
-	mu routeRWMutex
+	mu routeRWMutex `state:"nosave"`
 
 	// localAddressEndpoint is the local address this route is associated with.
 	// +checklocks:mu

pkg/tcpip/transport/tcp/endpoint.go

@@ -418,7 +418,8 @@ type Endpoint struct {
 	isPortReserved    bool
 	isRegistered      bool
 	boundNICID        tcpip.NICID
-	route             *stack.Route `state:"nosave"`
+	route             *stack.Route
+	restoreConn       bool
 	ipv4TTL           uint8
 	ipv6HopLimit      int16
 	isConnectNotified bool

pkg/tcpip/transport/tcp/endpoint_state.go

@@ -33,7 +33,7 @@ var logDisconnectOnce sync.Once
 
 func logDisconnect() {
 	logDisconnectOnce.Do(func() {
-		log.Infof("One or more TCP connections terminated during save")
+		log.Infof("One or more TCP connections terminated during restore")
 	})
 }
 
@@ -49,13 +49,7 @@ func (e *Endpoint) beforeSave() {
 	switch {
 	case epState == StateInitial || epState == StateBound:
 	case epState.connected() || epState.handshake():
-		if !e.route.HasSaveRestoreCapability() {
-			logDisconnect()
-			e.resetConnectionLocked(&tcpip.ErrConnectionAborted{})
-			e.mu.Unlock()
-			e.Close()
-			e.mu.Lock()
-		}
+		e.restoreConn = e.route.HasSaveRestoreCapability()
 		fallthrough
 	case epState == StateListen:
 		// Nothing to do.
@@ -133,6 +127,28 @@ func (e *Endpoint) afterLoad(ctx context.Context) {
 	}
 }
 
+func (e *Endpoint) releaseOldRoute() {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	if e.route != nil {
+		e.route.Release()
+		e.route = nil
+	}
+}
+
+func (e *Endpoint) closeAtRestore() {
+	e.mu.Lock()
+	logDisconnect()
+	e.resetConnectionLocked(&tcpip.ErrConnectionAborted{})
+	// Update the origEndpointState to error as it will be used during restore.
+	e.origEndpointState = uint32(StateError)
+	e.mu.Unlock()
+	e.Close()
+	e.stack.CompleteTransportEndpointCleanup(e)
+	tcpip.DeleteDanglingEndpoint(e)
+}
+
 // Restore implements tcpip.RestoredEndpoint.Restore.
 func (e *Endpoint) Restore(s *stack.Stack) {
 	if !e.EndpointState().closed() {
@@ -154,6 +170,7 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 
 	e.mu.Lock()
 	id := e.ID
+	restoreConn := e.restoreConn
 	e.mu.Unlock()
 
 	bind := func() {
@@ -186,6 +203,13 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 	epState := EndpointState(e.origEndpointState)
 	switch {
 	case epState.connected():
+		if !restoreConn {
+			e.closeAtRestore()
+			connectedLoading.Done()
+			return
+		}
+		e.releaseOldRoute()
+
 		bind()
 		if e.connectingAddress.BitLen() == 0 {
 			e.connectingAddress = e.TransportEndpointInfo.ID.RemoteAddress
@@ -254,6 +278,7 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 				tcpip.AsyncLoading.Done()
 			}()
 		} else {
+			e.releaseOldRoute()
 			go func() {
 				connectedLoading.Wait()
 				e.LockUser()
@@ -269,6 +294,7 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 			}()
 		}
 	case epState == StateConnecting:
+		e.releaseOldRoute()
 		// Initial SYN hasn't been sent yet so initiate a connect.
 		tcpip.AsyncLoading.Add(1)
 		go func() {
@@ -288,6 +314,16 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 		go func() {
 			connectedLoading.Wait()
 			listenLoading.Wait()
+
+			// Handshake phase.
+			if !restoreConn {
+				e.closeAtRestore()
+				connectingLoading.Done()
+				tcpip.AsyncLoading.Done()
+				return
+			}
+			e.releaseOldRoute()
+
 			// Initial SYN has been sent/received so we should bind the
 			// ports start the retransmit timer for the SYNs and let it
 			// naturally complete the connection.
@@ -314,6 +350,7 @@ func (e *Endpoint) Restore(s *stack.Stack) {
 			e.mu.Unlock()
 		}()
 	case epState == StateBound:
+		e.releaseOldRoute()
 		tcpip.AsyncLoading.Add(1)
 		go func() {
 			connectedLoading.Wait()