netstack: release rcv lock after ping socket save is done.

PiperOrigin-RevId: 196324694
Change-Id: Ia3a48976433f21622eacb4a38fefe7143ca5e31b

2 files changed, 20 insertions, 1 deletions

diff --git a/pkg/tcpip/transport/ping/endpoint.go b/pkg/tcpip/transport/ping/endpoint.go
index 609e7d947..2b80881bb 100644
--- a/pkg/tcpip/transport/ping/endpoint.go
+++ b/pkg/tcpip/transport/ping/endpoint.go
@@ -52,7 +52,7 @@ type endpoint struct {
 	rcvMu         sync.Mutex `state:"nosave"`
 	rcvReady      bool
 	rcvList       pingPacketList
-	rcvBufSizeMax int
+	rcvBufSizeMax int `state:".(int)"`
 	rcvBufSize    int
 	rcvClosed     bool
 	rcvTimestamp  bool
diff --git a/pkg/tcpip/transport/ping/endpoint_state.go b/pkg/tcpip/transport/ping/endpoint_state.go
index e1664f049..29fde2585 100644
--- a/pkg/tcpip/transport/ping/endpoint_state.go
+++ b/pkg/tcpip/transport/ping/endpoint_state.go
@@ -29,9 +29,28 @@ func (p *pingPacket) loadData(data buffer.VectorisedView) {
 // beforeSave is invoked by stateify.
 func (e *endpoint) beforeSave() {
 	// Stop incoming packets from being handled (and mutate endpoint state).
+	// The lock will be released after savercvBufSizeMax(), which would have
+	// saved e.rcvBufSizeMax and set it to 0 to continue blocking incoming
+	// packets.
 	e.rcvMu.Lock()
 }
 
+// saveRcvBufSizeMax is invoked by stateify.
+func (e *endpoint) saveRcvBufSizeMax() int {
+	max := e.rcvBufSizeMax
+	// Make sure no new packets will be handled regardless of the lock.
+	e.rcvBufSizeMax = 0
+	// Release the lock acquired in beforeSave() so regular endpoint closing
+	// logic can proceed after save.
+	e.rcvMu.Unlock()
+	return max
+}
+
+// loadRcvBufSizeMax is invoked by stateify.
+func (e *endpoint) loadRcvBufSizeMax(max int) {
+	e.rcvBufSizeMax = max
+}
+
 // afterLoad is invoked by stateify.
 func (e *endpoint) afterLoad() {
 	e.stack = stack.StackFromEnv