From 39406b00bf6642f37ba3f9800bd651f1020f7c15 Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Mon, 9 Dec 2024 09:31:38 -0800
Subject: [PATCH] netstack: add checklocks to TCPSenderState and write list

The annotations are a bit messy, but:

- It's worth it for the safety.
- Annotations could be removed by flattening tcp.Endpoint. Its current layout
  is, I believe, related to the old "protocol goroutine" architecture.
- The annotations could be cut down by improving checklocksalias.

PiperOrigin-RevId: 704315794
---
 pkg/tcpip/transport/tcp/connect.go    |  5 ++++
 pkg/tcpip/transport/tcp/cubic.go      | 17 ++++++++++++
 pkg/tcpip/transport/tcp/cubic_test.go | 16 +++++++++++
 pkg/tcpip/transport/tcp/endpoint.go   | 27 ++++++++++++++-----
 pkg/tcpip/transport/tcp/rack.go       |  8 ++++++
 pkg/tcpip/transport/tcp/rcv.go        |  1 +
 pkg/tcpip/transport/tcp/reno.go       | 12 +++++++++
 pkg/tcpip/transport/tcp/snd.go        | 38 ++++++++++++++++++++++++++-
 8 files changed, 117 insertions(+), 7 deletions(-)

diff --git a/pkg/tcpip/transport/tcp/connect.go b/pkg/tcpip/transport/tcp/connect.go
index b7e2db465..e3f87ffe9 100644
--- a/pkg/tcpip/transport/tcp/connect.go
+++ b/pkg/tcpip/transport/tcp/connect.go
@@ -647,6 +647,7 @@ func (h *handshake) retransmitHandlerLocked() tcpip.Error {
 // to an established state given the last segment received from peer. It also
 // initializes sender/receiver.
 // +checklocks:h.ep.mu
+// +checklocksalias:h.ep.snd.ep.mu=h.ep.mu
 func (h *handshake) transitionToStateEstablishedLocked(s *segment) {
 	// Stop the SYN retransmissions now that handshake is complete.
 	if h.retransmitTimer != nil {
@@ -1064,6 +1065,7 @@ func (e *Endpoint) sendData(next *segment) {
 // error code and sends a RST if and only if the error is not ErrConnectionReset
 // indicating that the connection is being reset due to receiving a RST.
 // +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *Endpoint) resetConnectionLocked(err tcpip.Error) {
 	// Only send a reset if the connection is being aborted for a reason
 	// other than receiving a reset.
@@ -1371,6 +1373,9 @@ func (e *Endpoint) keepaliveTimerExpired() tcpip.Error {
 
 // resetKeepaliveTimer restarts or stops the keepalive timer, depending on
 // whether it is enabled for this endpoint.
+//
+// +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *Endpoint) resetKeepaliveTimer(receivedData bool) {
 	e.keepalive.Lock()
 	defer e.keepalive.Unlock()
diff --git a/pkg/tcpip/transport/tcp/cubic.go b/pkg/tcpip/transport/tcp/cubic.go
index 46a052e81..ca459f08a 100644
--- a/pkg/tcpip/transport/tcp/cubic.go
+++ b/pkg/tcpip/transport/tcp/cubic.go
@@ -68,6 +68,8 @@ type cubicState struct {
 
 // newCubicCC returns a partially initialized cubic state with the constants
 // beta and c set and t set to current time.
+//
+// +checklocks:s.ep.mu
 func newCubicCC(s *sender) *cubicState {
 	now := s.ep.stack.Clock().NowMonotonic()
 	return &cubicState{
@@ -93,6 +95,8 @@ func newCubicCC(s *sender) *cubicState {
 // previously lowered ssThresh without experiencing packet loss.
 //
 // Refer: https://tools.ietf.org/html/rfc8312#section-4.8
+//
+// +checklocks:c.s.ep.mu
 func (c *cubicState) enterCongestionAvoidance() {
 	// See: https://tools.ietf.org/html/rfc8312#section-4.7 &
 	// https://tools.ietf.org/html/rfc8312#section-4.8
@@ -116,6 +120,8 @@ func (c *cubicState) enterCongestionAvoidance() {
 // increase').  The RFC version includes only the latter algorithm and adds an
 // intermediate phase called Conservative Slow Start, which is not implemented
 // here.
+//
+// +checklocks:c.s.ep.mu
 func (c *cubicState) updateHyStart(rtt time.Duration) {
 	if rtt < 0 {
 		// negative indicates unknown
@@ -151,6 +157,7 @@ func (c *cubicState) updateHyStart(rtt time.Duration) {
 	}
 }
 
+// +checklocks:c.s.ep.mu
 func (c *cubicState) beginHyStartRound(now tcpip.MonotonicTime) {
 	c.EndSeq = c.s.SndNxt
 	c.SampleCount = 0
@@ -164,6 +171,8 @@ func (c *cubicState) beginHyStartRound(now tcpip.MonotonicTime) {
 // algorithm used by NewReno. If after adjusting the congestion window we cross
 // the ssThresh then it will return the number of packets that must be consumed
 // in congestion avoidance mode.
+//
+// +checklocks:c.s.ep.mu
 func (c *cubicState) updateSlowStart(packetsAcked int) int {
 	// Don't let the congestion window cross into the congestion
 	// avoidance range.
@@ -186,6 +195,8 @@ func (c *cubicState) updateSlowStart(packetsAcked int) int {
 // Update updates cubic's internal state variables. It must be called on every
 // ACK received.
 // Refer: https://tools.ietf.org/html/rfc8312#section-4
+//
+// +checklocks:c.s.ep.mu
 func (c *cubicState) Update(packetsAcked int, rtt time.Duration) {
 	if c.s.Ssthresh == InitialSsthresh && c.s.SndCwnd < c.s.Ssthresh {
 		c.updateHyStart(rtt)
@@ -246,6 +257,8 @@ func (c *cubicState) getCwnd(packetsAcked, sndCwnd int, srtt time.Duration) int
 }
 
 // HandleLossDetected implements congestionControl.HandleLossDetected.
+//
+// +checklocks:c.s.ep.mu
 func (c *cubicState) HandleLossDetected() {
 	// See: https://tools.ietf.org/html/rfc8312#section-4.5
 	c.numCongestionEvents++
@@ -258,6 +271,8 @@ func (c *cubicState) HandleLossDetected() {
 }
 
 // HandleRTOExpired implements congestionContrl.HandleRTOExpired.
+//
+// +checklocks:c.s.ep.mu
 func (c *cubicState) HandleRTOExpired() {
 	// See: https://tools.ietf.org/html/rfc8312#section-4.6
 	c.T = c.s.ep.stack.Clock().NowMonotonic()
@@ -296,6 +311,8 @@ func (c *cubicState) PostRecovery() {
 
 // reduceSlowStartThreshold returns new SsThresh as described in
 // https://tools.ietf.org/html/rfc8312#section-4.7.
+//
+// +checklocks:c.s.ep.mu
 func (c *cubicState) reduceSlowStartThreshold() {
 	c.s.Ssthresh = int(math.Max(float64(c.s.SndCwnd)*c.Beta, 2.0))
 }
diff --git a/pkg/tcpip/transport/tcp/cubic_test.go b/pkg/tcpip/transport/tcp/cubic_test.go
index af31386c9..dbd427653 100644
--- a/pkg/tcpip/transport/tcp/cubic_test.go
+++ b/pkg/tcpip/transport/tcp/cubic_test.go
@@ -46,7 +46,9 @@ func TestHyStartAckTrainOK(t *testing.T) {
 			Ssthresh: InitialSsthresh,
 		},
 	}
+	snd.ep.mu.Lock()
 	uut := newCubicCC(snd)
+	snd.ep.mu.Unlock()
 	snd.cc = uut
 
 	if uut.LastRTT != effectivelyInfinity {
@@ -57,6 +59,8 @@ func TestHyStartAckTrainOK(t *testing.T) {
 	}
 
 	d0 := 4 * time.Millisecond
+	uut.s.ep.mu.Lock()
+	defer uut.s.ep.mu.Unlock()
 	uut.updateHyStart(d0)
 	if uut.CurrRTT != d0 {
 		t.Fatal()
@@ -127,7 +131,9 @@ func TestHyStartAckTrainTooSpread(t *testing.T) {
 			Ssthresh: InitialSsthresh,
 		},
 	}
+	snd.ep.mu.Lock()
 	uut := newCubicCC(snd)
+	snd.ep.mu.Unlock()
 	snd.cc = uut
 
 	if uut.LastRTT != effectivelyInfinity {
@@ -137,6 +143,8 @@ func TestHyStartAckTrainTooSpread(t *testing.T) {
 		t.Fatal()
 	}
 	d0 := 4 * time.Millisecond
+	uut.s.ep.mu.Lock()
+	defer uut.s.ep.mu.Unlock()
 	uut.updateHyStart(d0)
 	if uut.CurrRTT != d0 {
 		t.Fatal()
@@ -196,10 +204,14 @@ func TestHyStartDelayOK(t *testing.T) {
 			Ssthresh: InitialSsthresh,
 		},
 	}
+	snd.ep.mu.Lock()
 	uut := newCubicCC(snd)
+	snd.ep.mu.Unlock()
 	snd.cc = uut
 
 	d0 := 4 * time.Millisecond
+	uut.s.ep.mu.Lock()
+	defer uut.s.ep.mu.Unlock()
 	uut.updateHyStart(d0)
 
 	// Move SndNext and SndUna to advance to a new round.
@@ -247,10 +259,14 @@ func TestHyStartDelay_BelowThresh(t *testing.T) {
 			Ssthresh: InitialSsthresh,
 		},
 	}
+	snd.ep.mu.Lock()
 	uut := newCubicCC(snd)
+	snd.ep.mu.Unlock()
 	snd.cc = uut
 
 	d0 := 4 * time.Millisecond
+	uut.s.ep.mu.Lock()
+	defer uut.s.ep.mu.Unlock()
 	uut.updateHyStart(d0)
 
 	// Move SndNext and SndUna to advance to a new round.
diff --git a/pkg/tcpip/transport/tcp/endpoint.go b/pkg/tcpip/transport/tcp/endpoint.go
index f92ba1e90..6ef254ff4 100644
--- a/pkg/tcpip/transport/tcp/endpoint.go
+++ b/pkg/tcpip/transport/tcp/endpoint.go
@@ -338,6 +338,9 @@ func (sq *sndQueueInfo) CloneState(other *TCPSndBufState) {
 // For more details please see the detailed documentation on
 // e.LockUser/e.UnlockUser methods.
 //
+// TODO(b/339664055): Checklocks should be used more extensively here. Coverage
+// is currently sparse.
+//
 // +stateify savable
 type Endpoint struct {
 	TCPEndpointStateInner
@@ -521,7 +524,8 @@ type Endpoint struct {
 	acceptQueue acceptQueue
 
 	rcv *receiver `state:"wait"`
-	snd *sender   `state:"wait"`
+
+	snd *sender `state:"wait"`
 
 	// The goroutine drain completion notification channel.
 	drainDone chan struct{} `state:"nosave"`
@@ -632,6 +636,7 @@ func (e *Endpoint) isOwnedByUser() bool {
 // should not be holding the lock for long and spinning reduces latency as we
 // avoid an expensive sleep/wakeup of the syscall goroutine).
 // +checklocksacquire:e.mu
+// +checklocksacquire:e.snd.ep.mu
 func (e *Endpoint) LockUser() {
 	const iterations = 5
 	for i := 0; i < iterations; i++ {
@@ -644,14 +649,14 @@ func (e *Endpoint) LockUser() {
 			if e.ownedByUser.Load() == 1 {
 				e.mu.Lock()
 				e.ownedByUser.Store(1)
-				return
+				return // +checklocksforce: this locks e.snd.ep.mu
 			}
 			// Spin but don't yield the processor since the lower half
 			// should yield the lock soon.
 			continue
 		}
 		e.ownedByUser.Store(1)
-		return
+		return // +checklocksforce: this locks e.snd.ep.mu
 	}
 
 	for i := 0; i < iterations; i++ {
@@ -664,7 +669,7 @@ func (e *Endpoint) LockUser() {
 			if e.ownedByUser.Load() == 1 {
 				e.mu.Lock()
 				e.ownedByUser.Store(1)
-				return
+				return // +checklocksforce: this locks e.snd.ep.mu
 			}
 			// Spin but yield the processor since the lower half
 			// should yield the lock soon.
@@ -672,7 +677,7 @@ func (e *Endpoint) LockUser() {
 			continue
 		}
 		e.ownedByUser.Store(1)
-		return
+		return // +checklocksforce: this locks e.snd.ep.mu
 	}
 
 	// Finally just give up and wait for the Lock.
@@ -1005,6 +1010,7 @@ func (e *Endpoint) purgeReadQueue() {
 }
 
 // +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *Endpoint) purgeWriteQueue() {
 	if e.snd != nil {
 		e.sndQueueInfo.sndQueueMu.Lock()
@@ -1582,6 +1588,7 @@ func (e *Endpoint) readFromPayloader(p tcpip.Payloader, opts tcpip.WriteOptions,
 
 // queueSegment reads data from the payloader and returns a segment to be sent.
 // +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *Endpoint) queueSegment(p tcpip.Payloader, opts tcpip.WriteOptions) (*segment, int, tcpip.Error) {
 	e.sndQueueInfo.sndQueueMu.Lock()
 	defer e.sndQueueInfo.sndQueueMu.Unlock()
@@ -2372,6 +2379,7 @@ func (e *Endpoint) registerEndpoint(addr tcpip.FullAddress, netProto tcpip.Netwo
 
 // connect connects the endpoint to its peer.
 // +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *Endpoint) connect(addr tcpip.FullAddress, handshake bool) tcpip.Error {
 	connectingAddr := addr.Addr
 
@@ -2465,7 +2473,6 @@ func (e *Endpoint) connect(addr tcpip.FullAddress, handshake bool) tcpip.Error {
 			}
 		}
 		e.segmentQueue.mu.Unlock()
-		e.snd.ep.AssertLockHeld(e)
 		e.snd.updateMaxPayloadSize(int(e.route.MTU()), 0)
 		e.setEndpointState(StateEstablished)
 		// Set the new auto tuned send buffer size after entering
@@ -2508,6 +2515,7 @@ func (e *Endpoint) Shutdown(flags tcpip.ShutdownFlags) tcpip.Error {
 }
 
 // +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *Endpoint) shutdownLocked(flags tcpip.ShutdownFlags) tcpip.Error {
 	e.shutdownFlags |= flags
 	switch {
@@ -2950,6 +2958,9 @@ func (e *Endpoint) HandleError(transErr stack.TransportError, pkt *stack.PacketB
 
 // updateSndBufferUsage is called by when room opens up in the send buffer. The
 // number of newly available bytes is v.
+//
+// +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *Endpoint) updateSndBufferUsage(v int) {
 	sendBufferSize := e.getSendBufferSize()
 	e.sndQueueInfo.sndQueueMu.Lock()
@@ -3131,6 +3142,7 @@ func (e *Endpoint) maxOptionSize() (size int) {
 // used before invoking the probe.
 //
 // +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *Endpoint) completeStateLocked(s *TCPEndpointState) {
 	s.TCPEndpointStateInner = e.TCPEndpointStateInner
 	s.ID = TCPEndpointID(e.TransportEndpointInfo.ID)
@@ -3280,6 +3292,9 @@ func GetTCPReceiveBufferLimits(s tcpip.StackHandler) tcpip.ReceiveBufferSizeOpti
 
 // computeTCPSendBufferSize implements auto tuning of send buffer size and
 // returns the new send buffer size.
+//
+// +checklocks:e.mu
+// +checklocksalias:e.snd.ep.mu=e.mu
 func (e *Endpoint) computeTCPSendBufferSize() int64 {
 	curSndBufSz := int64(e.getSendBufferSize())
 
diff --git a/pkg/tcpip/transport/tcp/rack.go b/pkg/tcpip/transport/tcp/rack.go
index 77f50fa92..fd80dd77e 100644
--- a/pkg/tcpip/transport/tcp/rack.go
+++ b/pkg/tcpip/transport/tcp/rack.go
@@ -161,6 +161,8 @@ func (s *sender) shouldSchedulePTO() bool {
 
 // schedulePTO schedules the probe timeout as defined in
 // https://tools.ietf.org/html/draft-ietf-tcpm-rack-08#section-7.5.1.
+//
+// +checklocks:s.ep.mu
 func (s *sender) schedulePTO() {
 	pto := time.Second
 	s.rtt.Lock()
@@ -236,6 +238,8 @@ func (s *sender) probeTimerExpired() tcpip.Error {
 // detectTLPRecovery detects if recovery was accomplished by the loss probes
 // and updates TLP state accordingly.
 // See https://tools.ietf.org/html/draft-ietf-tcpm-rack-08#section-7.6.3.
+//
+// +checklocks:s.ep.mu
 func (s *sender) detectTLPRecovery(ack seqnum.Value, rcvdSeg *segment) {
 	if !(s.ep.SACKPermitted && s.rc.tlpRxtOut) {
 		return
@@ -279,6 +283,8 @@ func (s *sender) detectTLPRecovery(ack seqnum.Value, rcvdSeg *segment) {
 //     been observed RACK uses reo_wnd of zero during loss recovery, in order to
 //     retransmit quickly, or when the number of DUPACKs exceeds the classic
 //     DUPACKthreshold.
+//
+// +checklocks:rc.snd.ep.mu
 func (rc *rackControl) updateRACKReorderWindow() {
 	dsackSeen := rc.DSACKSeen
 	snd := rc.snd
@@ -352,6 +358,8 @@ func (rc *rackControl) exitRecovery() {
 // detectLoss marks the segment as lost if the reordering window has elapsed
 // and the ACK is not received. It will also arm the reorder timer.
 // See: https://tools.ietf.org/html/draft-ietf-tcpm-rack-08#section-7.2 Step 5.
+//
+// +checklocks:rc.snd.ep.mu
 func (rc *rackControl) detectLoss(rcvTime tcpip.MonotonicTime) int {
 	var timeout time.Duration
 	numLost := 0
diff --git a/pkg/tcpip/transport/tcp/rcv.go b/pkg/tcpip/transport/tcp/rcv.go
index ba984afe1..64045c690 100644
--- a/pkg/tcpip/transport/tcp/rcv.go
+++ b/pkg/tcpip/transport/tcp/rcv.go
@@ -96,6 +96,7 @@ func (r *receiver) currentWindow() (curWnd seqnum.Size) {
 // getSendParams returns the parameters needed by the sender when building
 // segments to send.
 // +checklocks:r.ep.mu
+// +checklocksalias:r.ep.snd.ep.mu=r.ep.mu
 func (r *receiver) getSendParams() (RcvNxt seqnum.Value, rcvWnd seqnum.Size) {
 	newWnd := r.ep.selectWindow()
 	curWnd := r.currentWindow()
diff --git a/pkg/tcpip/transport/tcp/reno.go b/pkg/tcpip/transport/tcp/reno.go
index 2d1b011db..0f137a327 100644
--- a/pkg/tcpip/transport/tcp/reno.go
+++ b/pkg/tcpip/transport/tcp/reno.go
@@ -35,6 +35,8 @@ func newRenoCC(s *sender) *renoState {
 // algorithm used by NewReno. If after adjusting the congestion window
 // we cross the SSthreshold then it will return the number of packets that
 // must be consumed in congestion avoidance mode.
+//
+// +checklocks:r.s.ep.mu
 func (r *renoState) updateSlowStart(packetsAcked int) int {
 	// Don't let the congestion window cross into the congestion
 	// avoidance range.
@@ -51,6 +53,8 @@ func (r *renoState) updateSlowStart(packetsAcked int) int {
 
 // updateCongestionAvoidance will update congestion window in congestion
 // avoidance mode as described in RFC5681 section 3.1
+//
+// +checklocks:r.s.ep.mu
 func (r *renoState) updateCongestionAvoidance(packetsAcked int) {
 	// Consume the packets in congestion avoidance mode.
 	r.s.SndCAAckCount += packetsAcked
@@ -62,6 +66,8 @@ func (r *renoState) updateCongestionAvoidance(packetsAcked int) {
 
 // reduceSlowStartThreshold reduces the slow-start threshold per RFC 5681,
 // page 6, eq. 4. It is called when we detect congestion in the network.
+//
+// +checklocks:r.s.ep.mu
 func (r *renoState) reduceSlowStartThreshold() {
 	r.s.Ssthresh = r.s.Outstanding / 2
 	if r.s.Ssthresh < 2 {
@@ -73,6 +79,8 @@ func (r *renoState) reduceSlowStartThreshold() {
 // Update updates the congestion state based on the number of packets that
 // were acknowledged.
 // Update implements congestionControl.Update.
+//
+// +checklocks:r.s.ep.mu
 func (r *renoState) Update(packetsAcked int, _ time.Duration) {
 	if r.s.SndCwnd < r.s.Ssthresh {
 		packetsAcked = r.updateSlowStart(packetsAcked)
@@ -84,6 +92,8 @@ func (r *renoState) Update(packetsAcked int, _ time.Duration) {
 }
 
 // HandleLossDetected implements congestionControl.HandleLossDetected.
+//
+// +checklocks:r.s.ep.mu
 func (r *renoState) HandleLossDetected() {
 	// A retransmit was triggered due to nDupAckThreshold or when RACK
 	// detected loss. Reduce our slow start threshold.
@@ -91,6 +101,8 @@ func (r *renoState) HandleLossDetected() {
 }
 
 // HandleRTOExpired implements congestionControl.HandleRTOExpired.
+//
+// +checklocks:r.s.ep.mu
 func (r *renoState) HandleRTOExpired() {
 	// We lost a packet, so reduce ssthresh.
 	r.reduceSlowStartThreshold()
diff --git a/pkg/tcpip/transport/tcp/snd.go b/pkg/tcpip/transport/tcp/snd.go
index 440316b16..da52fecb4 100644
--- a/pkg/tcpip/transport/tcp/snd.go
+++ b/pkg/tcpip/transport/tcp/snd.go
@@ -99,7 +99,9 @@ type lossRecovery interface {
 //
 // +stateify savable
 type sender struct {
+	// +checklocks:ep.mu
 	TCPSenderState
+
 	ep *Endpoint
 
 	// lr is the loss recovery algorithm used by the sender.
@@ -124,6 +126,8 @@ type sender struct {
 	// writeList holds all writable data: both unsent data and
 	// sent-but-unacknowledged data. Alternatively: it holds all bytes
 	// starting from SND.UNA.
+	//
+	// +checklocks:ep.mu
 	writeList protectedWriteList
 
 	// resendTimer is used for RTOs.
@@ -268,7 +272,14 @@ func newSender(ep *Endpoint, iss, irs seqnum.Value, sndWnd seqnum.Size, mss uint
 			set: make(map[*segment]struct{}),
 		},
 	}
+	return newSenderHelper(ep, iss, irs, sndWnd, mss, sndWndScale, maxPayloadSize, s)
+}
 
+// newSenderHelper exists to sate checklocks.
+//
+// +checklocks:ep.mu
+// +checklocksalias:s.ep.mu=ep.mu
+func newSenderHelper(ep *Endpoint, iss, irs seqnum.Value, sndWnd seqnum.Size, mss uint16, sndWndScale int, maxPayloadSize int, s *sender) *sender {
 	if s.gso {
 		s.ep.gso.MSS = uint16(maxPayloadSize)
 	}
@@ -288,7 +299,6 @@ func newSender(ep *Endpoint, iss, irs seqnum.Value, sndWnd seqnum.Size, mss uint
 	s.probeTimer.init(s.ep.stack.Clock(), timerHandler(s.ep, s.probeTimerExpired))
 	s.corkTimer.init(s.ep.stack.Clock(), timerHandler(s.ep, s.corkTimerExpired))
 
-	s.ep.AssertLockHeld(ep)
 	s.updateMaxPayloadSize(int(ep.route.MTU()), 0)
 	// Initialize SACK Scoreboard after updating max payload size as we use
 	// the maxPayloadSize as the smss when determining if a segment is lost
@@ -320,6 +330,8 @@ func newSender(ep *Endpoint, iss, irs seqnum.Value, sndWnd seqnum.Size, mss uint
 // initCongestionControl initializes the specified congestion control module and
 // returns a handle to it. It also initializes the sndCwnd and sndSsThresh to
 // their initial values.
+//
+// +checklocks:s.ep.mu
 func (s *sender) initCongestionControl(congestionControlName tcpip.CongestionControlOption) congestionControl {
 	s.SndCwnd = InitialCwnd
 	s.Ssthresh = InitialSsthresh
@@ -420,6 +432,8 @@ func (s *sender) sendAck() {
 
 // updateRTO updates the retransmit timeout when a new roud-trip time is
 // available. This is done in accordance with section 2 of RFC 6298.
+//
+// +checklocks:s.ep.mu
 func (s *sender) updateRTO(rtt time.Duration) {
 	s.rtt.Lock()
 	if !s.rtt.TCPRTTState.SRTTInited {
@@ -666,6 +680,8 @@ func (s *sender) pCount(seg *segment, maxPayloadSize int) int {
 
 // splitSeg splits a given segment at the size specified and inserts the
 // remainder as a new segment after the current one in the write list.
+//
+// +checklocks:s.ep.mu
 func (s *sender) splitSeg(seg *segment, size int) {
 	if seg.payloadSize() <= size {
 		return
@@ -701,6 +717,8 @@ func (s *sender) splitSeg(seg *segment, size int) {
 //
 // rescueRtx will be true only if nextSeg is a rescue retransmission as
 // described by Step 4) of the NextSeg algorithm.
+//
+// +checklocks:s.ep.mu
 func (s *sender) NextSeg(nextSegHint *segment) (nextSeg, hint *segment, rescueRtx bool) {
 	var s3 *segment
 	var s4 *segment
@@ -1004,6 +1022,7 @@ func (s *sender) sendZeroWindowProbe() {
 	s.resendTimer.enable(s.RTO)
 }
 
+// +checklocks:s.ep.mu
 func (s *sender) enableZeroWindowProbing() {
 	s.zeroWindowProbing = true
 	// We piggyback the probing on the retransmit timer with the
@@ -1022,6 +1041,7 @@ func (s *sender) disableZeroWindowProbing() {
 	s.resendTimer.disable()
 }
 
+// +checklocks:s.ep.mu
 func (s *sender) postXmit(dataSent bool, shouldScheduleProbe bool) {
 	if dataSent {
 		// We sent data, so we should stop the keepalive timer to ensure
@@ -1098,6 +1118,7 @@ func (s *sender) sendData() {
 	s.postXmit(dataSent, true /* shouldScheduleProbe */)
 }
 
+// +checklocks:s.ep.mu
 func (s *sender) enterRecovery() {
 	// Initialize the variables used to detect spurious recovery after
 	// entering recovery.
@@ -1141,6 +1162,7 @@ func (s *sender) enterRecovery() {
 	s.ep.stack.Stats().TCP.FastRecovery.Increment()
 }
 
+// +checklocks:s.ep.mu
 func (s *sender) leaveRecovery() {
 	s.FastRecovery.Active = false
 	s.FastRecovery.MaxCwnd = 0
@@ -1163,6 +1185,8 @@ func (s *sender) isAssignedSequenceNumber(seg *segment) bool {
 // maintains the congestion window in number of packets and not bytes, so
 // SetPipe() here measures number of outstanding packets rather than actual
 // outstanding bytes in the network.
+//
+// +checklocks:s.ep.mu
 func (s *sender) SetPipe() {
 	// If SACK isn't permitted or it is permitted but recovery is not active
 	// then ignore pipe calculations.
@@ -1216,6 +1240,8 @@ func (s *sender) SetPipe() {
 // shouldEnterRecovery returns true if the sender should enter fast recovery
 // based on dupAck count and sack scoreboard.
 // See RFC 6675 section 5.
+//
+// +checklocks:s.ep.mu
 func (s *sender) shouldEnterRecovery() bool {
 	return s.DupAckCount >= nDupAckThreshold ||
 		(s.ep.SACKPermitted && s.ep.tcpRecovery&tcpip.TCPRACKLossDetection == 0 && s.ep.scoreboard.IsLost(s.SndUna))
@@ -1224,6 +1250,8 @@ func (s *sender) shouldEnterRecovery() bool {
 // detectLoss is called when an ack is received and returns whether a loss is
 // detected. It manages the state related to duplicate acks and determines if
 // a retransmit is needed according to the rules in RFC 6582 (NewReno).
+//
+// +checklocks:s.ep.mu
 func (s *sender) detectLoss(seg *segment) (fastRetransmit bool) {
 	// We're not in fast recovery yet.
 
@@ -1274,6 +1302,8 @@ func (s *sender) detectLoss(seg *segment) (fastRetransmit bool) {
 
 // isDupAck determines if seg is a duplicate ack as defined in
 // https://tools.ietf.org/html/rfc5681#section-2.
+//
+// +checklocks:s.ep.mu
 func (s *sender) isDupAck(seg *segment) bool {
 	// A TCP that utilizes selective acknowledgments (SACKs) [RFC2018, RFC2883]
 	// can leverage the SACK information to determine when an incoming ACK is a
@@ -1304,6 +1334,8 @@ func (s *sender) isDupAck(seg *segment) bool {
 //
 // See: https://tools.ietf.org/html/draft-ietf-tcpm-rack-08#section-7.2
 // steps 2 and 3.
+//
+// +checklocks:s.ep.mu
 func (s *sender) walkSACK(rcvdSeg *segment) bool {
 	s.rc.setDSACKSeen(false)
 
@@ -1421,6 +1453,7 @@ func (s *sender) recordRetransmitTS() {
 	s.retransmitTS = s.ep.tsValNow()
 }
 
+// +checklocks:s.ep.mu
 func (s *sender) detectSpuriousRecovery(hasDSACK bool, tsEchoReply uint32) {
 	// Return if the sender has already detected spurious recovery.
 	if s.spuriousRecovery {
@@ -1790,6 +1823,7 @@ func (s *sender) sendSegment(seg *segment) tcpip.Error {
 // flags and sequence number.
 // +checklocks:s.ep.mu
 // +checklocksalias:s.ep.rcv.ep.mu=s.ep.mu
+// +checklocksalias:s.ep.rcv.ep.snd.ep.mu=s.ep.mu
 func (s *sender) sendSegmentFromPacketBuffer(pkt *stack.PacketBuffer, flags header.TCPFlags, seq seqnum.Value) tcpip.Error {
 	s.LastSendTime = s.ep.stack.Clock().NowMonotonic()
 	if seq == s.RTTMeasureSeqNum {
@@ -1811,7 +1845,9 @@ func (s *sender) sendSegmentFromPacketBuffer(pkt *stack.PacketBuffer, flags head
 
 // sendEmptySegment sends a new empty segment, flags and sequence number.
 // +checklocks:s.ep.mu
+// +checklocksalias:s.ep.rcv.ep.snd.ep.mu=s.ep.mu
 // +checklocksalias:s.ep.rcv.ep.mu=s.ep.mu
+// +checklocksalias:s.ep.snd.ep.mu=s.ep.mu
 func (s *sender) sendEmptySegment(flags header.TCPFlags, seq seqnum.Value) tcpip.Error {
 	s.LastSendTime = s.ep.stack.Clock().NowMonotonic()
 	if seq == s.RTTMeasureSeqNum {