Remove TCP endpoint goroutines.

This change removes all endpoint goroutines and all TCP processing is now done
inline in the TCP processor loop. TCP timers directly invoke handlers as
required rather than assert a waker.

UnlockUser is also simplified to just queue the endpoint to the processor
instead of trying to process segments inline. This allows us to centralize logic
for TCP state handling in the processor. This potentially could involve an extra
wakeup but now that endpoint goroutines do not exist this is not such a big
concern as in case of busy servers the processor goroutines will already be
running anyway.

This change also allows us to clean up S/R as now restoring a TCP endpoint does
not require restarting a goroutine and moving it to the right logical point but
only requires that we restart any timers that may have been running when the
save was done and restore any port bindings as required.

Endpoint.Release is now removed in favor of Endpoint.Abort by using Abort in
places where we use Endpoint.Release.

Updates #231

PiperOrigin-RevId: 442673015

pkg/refsvfs2/refs_map.go

@@ -47,9 +47,9 @@ func init() {
 	liveObjects = make(map[CheckedObject]struct{})
 }
 
-// leakCheckEnabled returns whether leak checking is enabled. The following
+// LeakCheckEnabled returns whether leak checking is enabled. The following
 // functions should only be called if it returns true.
-func leakCheckEnabled() bool {
+func LeakCheckEnabled() bool {
 	return refs_vfs1.GetLeakMode() != refs_vfs1.NoLeakChecking
 }
 
@@ -61,14 +61,14 @@ func leakCheckPanicEnabled() bool {
 
 // Register adds obj to the live object map.
 func Register(obj CheckedObject) {
-	if leakCheckEnabled() {
+	if LeakCheckEnabled() {
 		liveObjectsMu.Lock()
 		if _, ok := liveObjects[obj]; ok {
 			panic(fmt.Sprintf("Unexpected entry in leak checking map: reference %p already added", obj))
 		}
 		liveObjects[obj] = struct{}{}
 		liveObjectsMu.Unlock()
-		if leakCheckEnabled() && obj.LogRefs() {
+		if LeakCheckEnabled() && obj.LogRefs() {
 			logEvent(obj, "registered")
 		}
 	}
@@ -76,14 +76,14 @@ func Register(obj CheckedObject) {
 
 // Unregister removes obj from the live object map.
 func Unregister(obj CheckedObject) {
-	if leakCheckEnabled() {
+	if LeakCheckEnabled() {
 		liveObjectsMu.Lock()
 		defer liveObjectsMu.Unlock()
 		if _, ok := liveObjects[obj]; !ok {
 			panic(fmt.Sprintf("Expected to find entry in leak checking map for reference %p", obj))
 		}
 		delete(liveObjects, obj)
-		if leakCheckEnabled() && obj.LogRefs() {
+		if LeakCheckEnabled() && obj.LogRefs() {
 			logEvent(obj, "unregistered")
 		}
 	}
@@ -91,21 +91,21 @@ func Unregister(obj CheckedObject) {
 
 // LogIncRef logs a reference increment.
 func LogIncRef(obj CheckedObject, refs int64) {
-	if leakCheckEnabled() && obj.LogRefs() {
+	if LeakCheckEnabled() && obj.LogRefs() {
 		logEvent(obj, fmt.Sprintf("IncRef to %d", refs))
 	}
 }
 
 // LogTryIncRef logs a successful TryIncRef call.
 func LogTryIncRef(obj CheckedObject, refs int64) {
-	if leakCheckEnabled() && obj.LogRefs() {
+	if LeakCheckEnabled() && obj.LogRefs() {
 		logEvent(obj, fmt.Sprintf("TryIncRef to %d", refs))
 	}
 }
 
 // LogDecRef logs a reference decrement.
 func LogDecRef(obj CheckedObject, refs int64) {
-	if leakCheckEnabled() && obj.LogRefs() {
+	if LeakCheckEnabled() && obj.LogRefs() {
 		logEvent(obj, fmt.Sprintf("DecRef to %d", refs))
 	}
 }
@@ -128,7 +128,7 @@ var checkOnce sync.Once
 // anymore, at which point anything left in the map is considered a leak. On
 // multiple calls, only the first call will perform the leak check.
 func DoLeakCheck() {
-	if leakCheckEnabled() {
+	if LeakCheckEnabled() {
 		checkOnce.Do(doLeakCheck)
 	}
 }
@@ -136,7 +136,7 @@ func DoLeakCheck() {
 // DoRepeatedLeakCheck is the same as DoLeakCheck except that it can be called
 // multiple times by the caller to incrementally perform leak checking.
 func DoRepeatedLeakCheck() {
-	if leakCheckEnabled() {
+	if LeakCheckEnabled() {
 		doLeakCheck()
 	}
 }

pkg/sentry/inet/inet.go

@@ -79,6 +79,9 @@ type Stack interface {
 	// RouteTable returns the network stack's route table.
 	RouteTable() []Route
 
+	// Pause pauses the network stack before save.
+	Pause()
+
 	// Resume restarts the network stack after restore.
 	Resume()
 

pkg/sentry/inet/test_stack.go

@@ -144,6 +144,9 @@ func (s *TestStack) RouteTable() []Route {
 	return s.RouteList
 }
 
+// Pause implements Stack.
+func (s *TestStack) Pause() {}
+
 // Resume implements Stack.
 func (s *TestStack) Resume() {}
 

pkg/sentry/kernel/kernel.go

@@ -552,6 +552,15 @@ func (k *Kernel) SaveTo(ctx context.Context, w wire.Writer) error {
 
 	// Save the timekeeper's state.
 
+	if rootNS := k.rootNetworkNamespace; rootNS != nil && rootNS.Stack() != nil {
+		// Pause the network stack.
+		netstackPauseStart := time.Now()
+		log.Infof("Pausing root network namespace")
+		k.rootNetworkNamespace.Stack().Pause()
+		defer k.rootNetworkNamespace.Stack().Resume()
+		log.Infof("Pausing root network namespace took [%s].", time.Since(netstackPauseStart))
+	}
+
 	// Save the kernel state.
 	kernelStart := time.Now()
 	stats, err := state.Save(ctx, w, k)

pkg/sentry/socket/hostinet/stack.go

@@ -481,6 +481,9 @@ func (s *Stack) RouteTable() []inet.Route {
 	return append([]inet.Route(nil), s.routes...)
 }
 
+// Pause implements inet.Stack.Pause.
+func (*Stack) Pause() {}
+
 // Resume implements inet.Stack.Resume.
 func (*Stack) Resume() {}
 

pkg/sentry/socket/netstack/stack.go

@@ -444,6 +444,11 @@ func (s *Stack) IPTables() (*stack.IPTables, error) {
 	return s.Stack.IPTables(), nil
 }
 
+// Pause implements inet.Stack.Pause.
+func (s *Stack) Pause() {
+	s.Stack.Pause()
+}
+
 // Resume implements inet.Stack.Resume.
 func (s *Stack) Resume() {
 	s.Stack.Resume()

pkg/tcpip/header/tcp.go

@@ -134,6 +134,8 @@ type TCPFields struct {
 
 // TCPSynOptions is used to return the parsed TCP Options in a syn
 // segment.
+//
+// +stateify savable
 type TCPSynOptions struct {
 	// MSS is the maximum segment size provided by the peer in the SYN.
 	MSS uint16

pkg/tcpip/socketops.go

@@ -55,8 +55,10 @@ type SocketOptionsHandler interface {
 	// buffer size. It also returns the newly set value.
 	OnSetSendBufferSize(v int64) (newSz int64)
 
-	// OnSetReceiveBufferSize is invoked by SO_RCVBUF and SO_RCVBUFFORCE.
-	OnSetReceiveBufferSize(v, oldSz int64) (newSz int64)
+	// OnSetReceiveBufferSize is invoked by SO_RCVBUF and SO_RCVBUFFORCE. The
+	// handler can optionally return a callback which will be called after
+	// the buffer size is updated to newSz.
+	OnSetReceiveBufferSize(v, oldSz int64) (newSz int64, postSet func())
 
 	// WakeupWriters is invoked when the send buffer size for an endpoint is
 	// changed. The handler notifies the writers if the send buffer size is
@@ -107,8 +109,8 @@ func (*DefaultSocketOptionsHandler) OnSetSendBufferSize(v int64) (newSz int64) {
 func (*DefaultSocketOptionsHandler) WakeupWriters() {}
 
 // OnSetReceiveBufferSize implements SocketOptionsHandler.OnSetReceiveBufferSize.
-func (*DefaultSocketOptionsHandler) OnSetReceiveBufferSize(v, oldSz int64) (newSz int64) {
-	return v
+func (*DefaultSocketOptionsHandler) OnSetReceiveBufferSize(v, oldSz int64) (newSz int64, postSet func()) {
+	return v, nil
 }
 
 // StackHandler holds methods to access the stack options. These must be
@@ -700,11 +702,15 @@ func (so *SocketOptions) ReceiveBufferLimits() (min, max int64) {
 // SetReceiveBufferSize sets the value of the SO_RCVBUF option, optionally
 // notifying the owning endpoint.
 func (so *SocketOptions) SetReceiveBufferSize(receiveBufferSize int64, notify bool) {
+	var postSet func()
 	if notify {
 		oldSz := so.receiveBufferSize.Load()
-		receiveBufferSize = so.handler.OnSetReceiveBufferSize(receiveBufferSize, oldSz)
+		receiveBufferSize, postSet = so.handler.OnSetReceiveBufferSize(receiveBufferSize, oldSz)
 	}
 	so.receiveBufferSize.Store(receiveBufferSize)
+	if postSet != nil {
+		postSet()
+	}
 }
 
 // GetRcvlowat gets value for SO_RCVLOWAT option.

pkg/tcpip/stack/registration.go

@@ -225,6 +225,13 @@ type TransportProtocol interface {
 	// Wait waits for any worker goroutines owned by the protocol to stop.
 	Wait()
 
+	// Pause requests that any protocol level background workers pause.
+	Pause()
+
+	// Resume resumes any protocol level background workers that were
+	// previously paused by Pause.
+	Resume()
+
 	// 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

@@ -1606,6 +1606,13 @@ func (s *Stack) Wait() {
 	}
 }
 
+// Pause pauses any protocol level background workers.
+func (s *Stack) Pause() {
+	for _, p := range s.transportProtocols {
+		p.proto.Pause()
+	}
+}
+
 // Resume restarts the stack after a restore. This must be called after the
 // entire system has been restored.
 func (s *Stack) Resume() {
@@ -1618,6 +1625,10 @@ func (s *Stack) Resume() {
 	for _, e := range eps {
 		e.Resume(s)
 	}
+	// Now resume any protocol level background workers.
+	for _, p := range s.transportProtocols {
+		p.proto.Resume()
+	}
 }
 
 // RegisterPacketEndpoint registers ep with the stack, causing it to receive

pkg/tcpip/stack/transport_demuxer.go

@@ -181,22 +181,22 @@ func (epsByNIC *endpointsByNIC) handlePacket(id TransportEndpointID, pkt *Packet
 		epsByNIC.mu.RUnlock()
 		return true
 	}
+	epsByNIC.mu.RUnlock()
 
 	transEP.HandlePacket(id, pkt)
-	epsByNIC.mu.RUnlock() // Don't use defer for performance reasons.
 	return true
 }
 
 // handleError delivers an error to the transport endpoint identified by id.
 func (epsByNIC *endpointsByNIC) handleError(n *nic, id TransportEndpointID, transErr TransportError, pkt *PacketBuffer) {
 	epsByNIC.mu.RLock()
-	defer epsByNIC.mu.RUnlock()
 
 	mpep, ok := epsByNIC.endpoints[n.ID()]
 	if !ok {
 		mpep, ok = epsByNIC.endpoints[0]
 	}
 	if !ok {
+		epsByNIC.mu.RUnlock()
 		return
 	}
 
@@ -204,7 +204,10 @@ func (epsByNIC *endpointsByNIC) handleError(n *nic, id TransportEndpointID, tran
 	// broadcast like we are doing with handlePacket above?
 
 	// multiPortEndpoints are guaranteed to have at least one element.
-	mpep.selectEndpoint(id, epsByNIC.seed).HandleError(transErr, pkt)
+	transEP := mpep.selectEndpoint(id, epsByNIC.seed)
+	epsByNIC.mu.RUnlock()
+
+	transEP.HandleError(transErr, pkt)
 }
 
 // registerEndpoint returns true if it succeeds. It fails and returns

pkg/tcpip/stack/transport_test.go

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

pkg/tcpip/tcpip.go

@@ -674,9 +674,6 @@ type Endpoint interface {
 	// SocketOptions returns the structure which contains all the socket
 	// level options.
 	SocketOptions() *SocketOptions
-
-	// Release releases all reference counted objects held by the endpoint.
-	Release()
 }
 
 // LinkPacketInfo holds Link layer information for a received packet.
@@ -2499,7 +2496,7 @@ func ReleaseDanglingEndpoints() {
 	// Calling Release on a dangling endpoint that has been deleted is a noop.
 	eps := GetDanglingEndpoints()
 	for _, ep := range eps {
-		ep.Release()
+		ep.Abort()
 	}
 }
 

pkg/tcpip/transport/icmp/endpoint.go

@@ -767,9 +767,6 @@ func (e *endpoint) HandlePacket(id stack.TransportEndpointID, pkt *stack.PacketB
 // HandleError implements stack.TransportEndpoint.
 func (*endpoint) HandleError(stack.TransportError, *stack.PacketBuffer) {}
 
-// Release implements stack.TransportEndpoint.
-func (*endpoint) Release() {}
-
 // State implements tcpip.Endpoint.State. The ICMP endpoint currently doesn't
 // expose internal socket state.
 func (e *endpoint) State() uint32 {

pkg/tcpip/transport/icmp/protocol.go

@@ -121,6 +121,12 @@ func (*protocol) Close() {}
 // Wait implements stack.TransportProtocol.Wait.
 func (*protocol) Wait() {}
 
+// Pause implements stack.TransportProtocol.Pause.
+func (*protocol) Pause() {}
+
+// Resume implements stack.TransportProtocol.Resume.
+func (*protocol) Resume() {}
+
 // 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/packet/endpoint.go

@@ -492,9 +492,6 @@ func (ep *endpoint) Stats() tcpip.EndpointStats {
 // SetOwner implements tcpip.Endpoint.SetOwner.
 func (*endpoint) SetOwner(tcpip.PacketOwner) {}
 
-// Release implements tcpip.Release.
-func (*endpoint) Release() {}
-
 // SocketOptions implements tcpip.Endpoint.SocketOptions.
 func (ep *endpoint) SocketOptions() *tcpip.SocketOptions {
 	return &ep.ops

pkg/tcpip/transport/raw/endpoint.go

@@ -735,9 +735,6 @@ func (*endpoint) LastError() tcpip.Error {
 	return nil
 }
 
-// Release implements stack.TransportEndpoint.Release.
-func (*endpoint) Release() {}
-
 // SocketOptions implements tcpip.Endpoint.SocketOptions.
 func (e *endpoint) SocketOptions() *tcpip.SocketOptions {
 	return &e.ops

pkg/tcpip/transport/tcp/accept.go

@@ -23,7 +23,6 @@ import (
 	"io"
 	"time"
 
-	"gvisor.dev/gvisor/pkg/sleep"
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
@@ -225,10 +224,12 @@ func (l *listenContext) createConnectingEndpoint(s *segment, rcvdSynOpts header.
 // handshake in progress, which includes the new endpoint in the SYN-RCVD
 // state.
 //
-// On success, a handshake h is returned with h.ep.mu held.
+// On success, a handshake h is returned.
+//
+// NOTE: h.ep.mu is not held and must be acquired if any state needs to be
+// modified.
 //
 // Precondition: if l.listenEP != nil, l.listenEP.mu must be locked.
-// +checklocksacquire:h.ep.mu
 func (l *listenContext) startHandshake(s *segment, opts header.TCPSynOptions, queue *waiter.Queue, owner tcpip.PacketOwner) (h *handshake, _ tcpip.Error) {
 	// Create new endpoint.
 	irs := s.sequenceNumber
@@ -290,57 +291,56 @@ func (l *listenContext) startHandshake(s *segment, opts header.TCPSynOptions, qu
 	h = ep.newPassiveHandshake(isn, irs, opts, deferAccept)
 	h.listenEP = l.listenEP
 	h.start()
+	h.ep.mu.Unlock()
 	return h, nil
 }
 
 // performHandshake performs a TCP 3-way handshake. On success, the new
-// established endpoint is returned with e.mu held.
+// established endpoint is returned.
 //
 // Precondition: if l.listenEP != nil, l.listenEP.mu must be locked.
 func (l *listenContext) performHandshake(s *segment, opts header.TCPSynOptions, queue *waiter.Queue, owner tcpip.PacketOwner) (*endpoint, tcpip.Error) {
+	waitEntry, notifyCh := waiter.NewChannelEntry(waiter.WritableEvents)
+	queue.EventRegister(&waitEntry)
+	defer queue.EventUnregister(&waitEntry)
+
 	h, err := l.startHandshake(s, opts, queue, owner)
 	if err != nil {
 		return nil, err
 	}
-	ep := h.ep
 
-	// N.B. the endpoint is generated above by startHandshake, and will be
-	// returned locked. This first call is forced.
-	if err := h.complete(); err != nil { // +checklocksforce
+	// performHandshake is used by the Forwarder which will block till the
+	// handshake either succeeds or fails. We do this by registering for
+	// events above and block on the notification channel.
+	<-notifyCh
+
+	ep := h.ep
+	ep.mu.Lock()
+	if !ep.EndpointState().connected() {
 		ep.stack.Stats().TCP.FailedConnectionAttempts.Increment()
 		ep.stats.FailedConnectionAttempts.Increment()
-		l.cleanupFailedHandshake(h)
+		ep.h = nil
+		ep.mu.Unlock()
+		ep.Close()
+		ep.notifyAborted()
+		ep.drainClosingSegmentQueue()
 		return nil, err
 	}
-	l.cleanupCompletedHandshake(h)
-	return ep, nil
-}
 
-// +checklocks:h.ep.mu
-func (l *listenContext) cleanupFailedHandshake(h *handshake) {
-	e := h.ep
-	e.mu.Unlock()
-	e.Close()
-	e.notifyAborted()
-	e.drainClosingSegmentQueue()
-	e.h = nil
-}
-
-// cleanupCompletedHandshake transfers any state from the completed handshake to
-// the new endpoint.
-//
-// +checklocks:h.ep.mu
-func (l *listenContext) cleanupCompletedHandshake(h *handshake) {
-	e := h.ep
-	e.isConnectNotified = true
+	ep.isConnectNotified = true
 
+	// Transfer any state from the completed handshake to the endpoint.
+	//
 	// Update the receive window scaling. We can't do it before the
 	// handshake because it's possible that the peer doesn't support window
 	// scaling.
-	e.rcv.RcvWndScale = e.h.effectiveRcvWndScale()
+	ep.rcv.RcvWndScale = ep.h.effectiveRcvWndScale()
 
-	// Clean up handshake state stored in the endpoint so that it can be GCed.
-	e.h = nil
+	// Clean up handshake state stored in the endpoint so that it can be
+	// GCed.
+	ep.h = nil
+	ep.mu.Unlock()
+	return ep, nil
 }
 
 // propagateInheritableOptionsLocked propagates any options set on the listening
@@ -418,7 +418,7 @@ type acceptQueue struct {
 }
 
 func (a *acceptQueue) isFull() bool {
-	return a.endpoints.Len() == a.capacity
+	return a.endpoints.Len() >= a.capacity
 }
 
 // handleListenSegment is called when a listening endpoint receives a segment
@@ -478,59 +478,7 @@ func (e *endpoint) handleListenSegment(ctx *listenContext, s *segment) tcpip.Err
 				e.stats.FailedConnectionAttempts.Increment()
 				return false, err
 			}
-
 			e.acceptQueue.pendingEndpoints[h.ep] = struct{}{}
-			e.pendingAccepted.Add(1)
-
-			go func() {
-				defer func() {
-					e.pendingAccepted.Done()
-
-					e.acceptMu.Lock()
-					defer e.acceptMu.Unlock()
-					delete(e.acceptQueue.pendingEndpoints, h.ep)
-				}()
-
-				// Note that startHandshake returns a locked endpoint. The force call
-				// here just makes it so.
-				if err := h.complete(); err != nil { // +checklocksforce
-					e.stack.Stats().TCP.FailedConnectionAttempts.Increment()
-					e.stats.FailedConnectionAttempts.Increment()
-					ctx.cleanupFailedHandshake(h)
-					return
-				}
-				ctx.cleanupCompletedHandshake(h)
-				h.ep.startAcceptedLoop()
-				e.stack.Stats().TCP.PassiveConnectionOpenings.Increment()
-
-				// Deliver the endpoint to the accept queue.
-				//
-				// Drop the lock before notifying to avoid deadlock in user-specified
-				// callbacks.
-				delivered := func() bool {
-					e.acceptMu.Lock()
-					defer e.acceptMu.Unlock()
-					for {
-						// The listener is transitioning out of the Listen state; bail.
-						if e.acceptQueue.capacity == 0 {
-							return false
-						}
-						if e.acceptQueue.isFull() {
-							e.acceptCond.Wait()
-							continue
-						}
-
-						e.acceptQueue.endpoints.PushBack(h.ep)
-						return true
-					}
-				}()
-
-				if delivered {
-					e.waiterQueue.Notify(waiter.ReadableEvents)
-				} else {
-					h.ep.notifyProtocolGoroutine(notifyReset)
-				}
-			}()
 
 			return false, nil
 		}()
@@ -711,15 +659,14 @@ func (e *endpoint) handleListenSegment(ctx *listenContext, s *segment) tcpip.Err
 		}
 		h.ep.AssertLockHeld(n)
 		h.transitionToStateEstablishedLocked(s)
+		n.mu.Unlock()
 
 		// Requeue the segment if the ACK completing the handshake has more info
 		// to be procesed by the newly established endpoint.
 		if (s.flags.Contains(header.TCPFlagFin) || s.data.Size() > 0) && n.enqueueSegment(s) {
-			n.newSegmentWaker.Assert()
+			n.notifyProcessor()
 		}
 
-		// Start the protocol goroutine.
-		n.startAcceptedLoop()
 		e.stack.Stats().TCP.PassiveConnectionOpenings.Increment()
 
 		// Deliver the endpoint to the accept queue.
@@ -734,80 +681,3 @@ func (e *endpoint) handleListenSegment(ctx *listenContext, s *segment) tcpip.Err
 		return nil
 	}
 }
-
-// protocolListenLoop is the main loop of a listening TCP endpoint. It runs in
-// its own goroutine and is responsible for handling connection requests.
-func (e *endpoint) protocolListenLoop(rcvWnd seqnum.Size) {
-	e.mu.Lock()
-	v6Only := e.ops.GetV6Only()
-	ctx := newListenContext(e.stack, e.protocol, e, rcvWnd, v6Only, e.NetProto)
-
-	defer func() {
-		e.setEndpointState(StateClose)
-
-		// Do cleanup if needed.
-		e.completeWorkerLocked()
-
-		if e.drainDone != nil {
-			close(e.drainDone)
-		}
-		e.mu.Unlock()
-
-		e.drainClosingSegmentQueue()
-
-		// Notify waiters that the endpoint is shutdown.
-		e.waiterQueue.Notify(waiter.ReadableEvents | waiter.WritableEvents | waiter.EventHUp | waiter.EventErr)
-	}()
-
-	var s sleep.Sleeper
-	s.AddWaker(&e.notificationWaker)
-	s.AddWaker(&e.newSegmentWaker)
-	defer s.Done()
-	for {
-		e.mu.Unlock()
-		w := s.Fetch(true)
-		e.mu.Lock()
-		switch w {
-		case &e.notificationWaker:
-			n := e.fetchNotifications()
-			if n&notifyClose != 0 {
-				return
-			}
-			if n&notifyDrain != 0 {
-				for !e.segmentQueue.empty() {
-					s := e.segmentQueue.dequeue()
-					// TODO(gvisor.dev/issue/4690): Better handle errors instead of
-					// silently dropping.
-					_ = e.handleListenSegment(ctx, s)
-					s.DecRef()
-				}
-				close(e.drainDone)
-				e.mu.Unlock()
-				<-e.undrain
-				e.mu.Lock()
-			}
-
-		case &e.newSegmentWaker:
-			// Process at most maxSegmentsPerWake segments.
-			mayRequeue := true
-			for i := 0; i < maxSegmentsPerWake; i++ {
-				s := e.segmentQueue.dequeue()
-				if s == nil {
-					mayRequeue = false
-					break
-				}
-
-				// TODO(gvisor.dev/issue/4690): Better handle errors instead of
-				// silently dropping.
-				_ = e.handleListenSegment(ctx, s)
-				s.DecRef()
-			}
-
-			// If the queue is not empty, make sure we'll wake up
-			// in the next iteration.
-			if mayRequeue && !e.segmentQueue.empty() {
-				e.newSegmentWaker.Assert()
-			}
-		}
-	}
-}

pkg/tcpip/transport/tcp/dispatcher.go

@@ -16,6 +16,7 @@ package tcp
 
 import (
 	"encoding/binary"
+	"fmt"
 	"math/rand"
 
 	"gvisor.dev/gvisor/pkg/sleep"
@@ -24,6 +25,7 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/hash/jenkins"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
+	"gvisor.dev/gvisor/pkg/waiter"
 )
 
 // epQueue is a queue of endpoints.
@@ -35,13 +37,15 @@ type epQueue struct {
 // enqueue adds e to the queue if the endpoint is not already on the queue.
 func (q *epQueue) enqueue(e *endpoint) {
 	q.mu.Lock()
+	defer q.mu.Unlock()
+	e.pendingProcessingMu.Lock()
+	defer e.pendingProcessingMu.Unlock()
+
 	if e.pendingProcessing {
-		q.mu.Unlock()
 		return
 	}
 	q.list.PushBack(e)
 	e.pendingProcessing = true
-	q.mu.Unlock()
 }
 
 // dequeue removes and returns the first element from the queue if available,
@@ -50,7 +54,9 @@ func (q *epQueue) dequeue() *endpoint {
 	q.mu.Lock()
 	if e := q.list.Front(); e != nil {
 		q.list.Remove(e)
+		e.pendingProcessingMu.Lock()
 		e.pendingProcessing = false
+		e.pendingProcessingMu.Unlock()
 		q.mu.Unlock()
 		return e
 	}
@@ -72,6 +78,9 @@ type processor struct {
 	sleeper          sleep.Sleeper
 	newEndpointWaker sleep.Waker
 	closeWaker       sleep.Waker
+	pauseWaker       sleep.Waker
+	pauseChan        chan struct{}
+	resumeChan       chan struct{}
 }
 
 func (p *processor) close() {
@@ -84,58 +93,258 @@ func (p *processor) queueEndpoint(ep *endpoint) {
 	p.newEndpointWaker.Assert()
 }
 
-const (
-	newEndpointWaker = 1
-	closeWaker       = 2
-)
+// deliverAccepted delivers a passively connected endpoint to the accept queue
+// of its associated listening endpoint.
+//
+// +checklocks:ep.mu
+func deliverAccepted(ep *endpoint) bool {
+	lEP := ep.h.listenEP
+	lEP.acceptMu.Lock()
 
+	// Remove endpoint from list of pendingEndpoints as the handshake is now
+	// complete.
+	delete(lEP.acceptQueue.pendingEndpoints, ep)
+	// Deliver this endpoint to the listening socket's accept queue.
+	if lEP.acceptQueue.capacity == 0 {
+		lEP.acceptMu.Unlock()
+		return false
+	}
+
+	// NOTE: We always queue the endpoint and on purpose do not check if
+	// accept queue is full at this point. This is similar to linux because
+	// two racing incoming ACK's can both pass the acceptQueue.isFull check
+	// and proceed to ESTABLISHED state. In such a case its better to
+	// deliver both even if it temporarily exceeds the queue limit rather
+	// than drop a connection that is fully connected.
+	//
+	// For reference see:
+	//    https://github.com/torvalds/linux/blob/169e77764adc041b1dacba84ea90516a895d43b2/net/ipv4/tcp_minisocks.c#L764
+	//    https://github.com/torvalds/linux/blob/169e77764adc041b1dacba84ea90516a895d43b2/net/ipv4/tcp_ipv4.c#L1500
+	lEP.acceptQueue.endpoints.PushBack(ep)
+	lEP.acceptMu.Unlock()
+	ep.h.listenEP.waiterQueue.Notify(waiter.ReadableEvents)
+
+	return true
+}
+
+// handleConnecting is responsible for TCP processing for an endpoint in one of
+// the connecting states.
+func (p *processor) handleConnecting(ep *endpoint) {
+	if !ep.TryLock() {
+		return
+	}
+	cleanup := func() {
+		ep.mu.Unlock()
+		ep.drainClosingSegmentQueue()
+		ep.waiterQueue.Notify(waiter.EventHUp | waiter.EventErr | waiter.ReadableEvents | waiter.WritableEvents)
+	}
+	if !ep.EndpointState().connecting() {
+		// If the endpoint has already transitioned out of a connecting
+		// stage then just return (only possible if it was closed or
+		// timed out by the time we got around to processing the wakeup.
+		ep.mu.Unlock()
+		return
+	}
+	if err := ep.h.processSegments(); err != nil { // +checklocksforce:ep.h.ep.mu
+		// handshake failed. clean up the tcp endpoint and handshake
+		// state.
+		ep.handshakeFailed(err)
+		cleanup()
+		return
+	}
+
+	if ep.EndpointState() == StateEstablished && ep.h.listenEP != nil {
+		ep.isConnectNotified = true
+		ep.stack.Stats().TCP.PassiveConnectionOpenings.Increment()
+		if !deliverAccepted(ep) {
+			ep.resetConnectionLocked(&tcpip.ErrConnectionAborted{})
+			cleanup()
+			return
+		}
+	}
+	ep.mu.Unlock()
+}
+
+// handleConnected is responsible for TCP processing for an endpoint in one of
+// the connected states(StateEstablished, StateFinWait1 etc.)
+func (p *processor) handleConnected(ep *endpoint) {
+	if !ep.TryLock() {
+		return
+	}
+
+	if !ep.EndpointState().connected() {
+		// If the endpoint has already transitioned out of a connected
+		// state then just return (only possible if it was closed or
+		// timed out by the time we got around to processing the wakeup.
+		ep.mu.Unlock()
+		return
+	}
+
+	// NOTE: We read this outside of e.mu lock which means that by the time
+	// we get to handleSegments the endpoint may not be in ESTABLISHED. But
+	// this should be fine as all normal shutdown states are handled by
+	// handleSegmentsLocked.
+	switch err := ep.handleSegmentsLocked(); {
+	case err != nil:
+		// Send any active resets if required.
+		ep.resetConnectionLocked(err)
+		fallthrough
+	case ep.EndpointState() == StateClose:
+		ep.mu.Unlock()
+		ep.stack.Stats().TCP.CurrentConnected.Decrement()
+		ep.drainClosingSegmentQueue()
+		ep.waiterQueue.Notify(waiter.EventHUp | waiter.EventErr | waiter.ReadableEvents | waiter.WritableEvents)
+		return
+	case ep.EndpointState() == StateTimeWait:
+		p.startTimeWait(ep)
+	}
+	ep.mu.Unlock()
+}
+
+// startTimeWait starts a new goroutine to handle TIME-WAIT.
+//
+// +checklocks:ep.mu
+func (p *processor) startTimeWait(ep *endpoint) {
+	// Disable close timer as we are now entering real TIME_WAIT.
+	if ep.finWait2Timer != nil {
+		ep.finWait2Timer.Stop()
+	}
+	// Wake up any waiters before we start TIME-WAIT.
+	ep.waiterQueue.Notify(waiter.EventHUp | waiter.EventErr | waiter.ReadableEvents | waiter.WritableEvents)
+	timeWaitDuration := ep.getTimeWaitDuration()
+	ep.timeWaitTimer = ep.stack.Clock().AfterFunc(timeWaitDuration, ep.timeWaitTimerExpired)
+}
+
+// handleTimeWait is responsible for TCP processing for an endpoint in TIME-WAIT
+// state.
+func (p *processor) handleTimeWait(ep *endpoint) {
+	if !ep.TryLock() {
+		return
+	}
+
+	if ep.EndpointState() != StateTimeWait {
+		// If the endpoint has already transitioned out of a TIME-WAIT
+		// state then just return (only possible if it was closed or
+		// timed out by the time we got around to processing the wakeup.
+		ep.mu.Unlock()
+		return
+	}
+
+	extendTimeWait, reuseTW := ep.handleTimeWaitSegments()
+	if reuseTW != nil {
+		ep.transitionToStateCloseLocked()
+		ep.mu.Unlock()
+		ep.drainClosingSegmentQueue()
+		ep.waiterQueue.Notify(waiter.EventHUp | waiter.EventErr | waiter.ReadableEvents | waiter.WritableEvents)
+		reuseTW()
+		return
+	}
+	if extendTimeWait {
+		ep.timeWaitTimer.Reset(ep.getTimeWaitDuration())
+	}
+	ep.mu.Unlock()
+}
+
+// handleListen is responsible for TCP processing for an endpoint in LISTEN
+// state.
+func (p *processor) handleListen(ep *endpoint) {
+	if !ep.TryLock() {
+		return
+	}
+	defer ep.mu.Unlock()
+
+	if ep.EndpointState() != StateListen {
+		// If the endpoint has already transitioned out of a LISTEN
+		// state then just return (only possible if it was closed or
+		// shutdown).
+		return
+	}
+
+	for i := 0; i < maxSegmentsPerWake; i++ {
+		s := ep.segmentQueue.dequeue()
+		if s == nil {
+			break
+		}
+
+		// TODO(gvisor.dev/issue/4690): Better handle errors instead of
+		// silently dropping.
+		_ = ep.handleListenSegment(ep.listenCtx, s)
+		s.DecRef()
+	}
+}
+
+// start runs the main loop for a processor which is responsible for all TCP
+// processing for TCP endpoints.
 func (p *processor) start(wg *sync.WaitGroup) {
 	defer wg.Done()
 	defer p.sleeper.Done()
 
 	for {
-		if w := p.sleeper.Fetch(true); w == &p.closeWaker {
-			break
-		}
-		// If not the closeWaker, it must be &p.newEndpointWaker.
-		for {
-			ep := p.epQ.dequeue()
-			if ep == nil {
-				break
-			}
-			if ep.segmentQueue.empty() {
+		switch w := p.sleeper.Fetch(true); {
+		case w == &p.closeWaker:
+			return
+		case w == &p.pauseWaker:
+			if !p.epQ.empty() {
+				p.newEndpointWaker.Assert()
+				p.pauseWaker.Assert()
 				continue
+			} else {
+				p.pauseChan <- struct{}{}
+				<-p.resumeChan
 			}
-
-			// If socket has transitioned out of connected state then just let the
-			// worker handle the packet.
-			//
-			// NOTE: We read this outside of e.mu lock which means that by the time
-			// we get to handleSegments the endpoint may not be in ESTABLISHED. But
-			// this should be fine as all normal shutdown states are handled by
-			// handleSegments and if the endpoint moves to a CLOSED/ERROR state
-			// then handleSegments is a noop.
-			if ep.EndpointState() == StateEstablished && ep.TryLock() {
-				// If the endpoint is in a connected state then we do direct delivery
-				// to ensure low latency and avoid scheduler interactions.
-				switch err := ep.handleSegmentsLocked(true /* fastPath */); {
-				case err != nil:
-					// Send any active resets if required.
-					ep.resetConnectionLocked(err)
-					fallthrough
-				case ep.EndpointState() == StateClose:
-					ep.notifyProtocolGoroutine(notifyTickleWorker)
-				case !ep.segmentQueue.empty():
+		case w == &p.newEndpointWaker:
+			for {
+				ep := p.epQ.dequeue()
+				if ep == nil {
+					break
+				}
+				if ep.segmentQueue.empty() {
+					continue
+				}
+				switch state := ep.EndpointState(); {
+				case state.connecting():
+					p.handleConnecting(ep)
+				case state.connected() && state != StateTimeWait:
+					p.handleConnected(ep)
+				case state == StateTimeWait:
+					p.handleTimeWait(ep)
+				case state == StateListen:
+					p.handleListen(ep)
+				case state == StateError || state == StateClose:
+					// Try to redeliver any still queued
+					// packets to another endpoint or send a
+					// RST if it can't be delivered.
+					ep.mu.Lock()
+					if st := ep.EndpointState(); st == StateError || st == StateClose {
+						ep.drainClosingSegmentQueue()
+					}
+					ep.mu.Unlock()
+				default:
+					panic(fmt.Sprintf("unexpected tcp state in processor: %v", state))
+				}
+				// If there are more segments to process then
+				// requeue this endpoint for processing.
+				if !ep.segmentQueue.empty() {
 					p.epQ.enqueue(ep)
 				}
-				ep.mu.Unlock()
-			} else {
-				ep.newSegmentWaker.Assert()
 			}
 		}
 	}
 }
 
+// pause pauses the processor loop.
+func (p *processor) pause() chan struct{} {
+	p.pauseWaker.Assert()
+	return p.pauseChan
+}
+
+// resume resumes a previously paused loop.
+//
+// Precondition: Pause must have been called previously.
+func (p *processor) resume() {
+	p.resumeChan <- struct{}{}
+}
+
 // dispatcher manages a pool of TCP endpoint processors which are responsible
 // for the processing of inbound segments. This fixed pool of processor
 // goroutines do full tcp processing. The processor is selected based on the
@@ -143,20 +352,33 @@ func (p *processor) start(wg *sync.WaitGroup) {
 // in-order.
 type dispatcher struct {
 	processors []processor
-	// seed is a random secret for a jenkins hash.
-	seed uint32
-	wg   sync.WaitGroup
+	wg         sync.WaitGroup
+	hasher     jenkinsHasher
+	mu         sync.Mutex
+	// +checklocks:mu
+	paused bool
+	// +checklocks:mu
+	closed bool
 }
 
+// init initializes a dispatcher and starts the main loop for all the processors
+// owned by this dispatcher.
 func (d *dispatcher) init(rng *rand.Rand, nProcessors int) {
 	d.close()
 	d.wait()
+
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	d.closed = false
 	d.processors = make([]processor, nProcessors)
-	d.seed = rng.Uint32()
+	d.hasher = jenkinsHasher{seed: rng.Uint32()}
 	for i := range d.processors {
 		p := &d.processors[i]
 		p.sleeper.AddWaker(&p.newEndpointWaker)
 		p.sleeper.AddWaker(&p.closeWaker)
+		p.sleeper.AddWaker(&p.pauseWaker)
+		p.pauseChan = make(chan struct{})
+		p.resumeChan = make(chan struct{})
 		d.wg.Add(1)
 		// NB: sleeper-waker registration must happen synchronously to avoid races
 		// with `close`.  It's possible to pull all this logic into `start`, but
@@ -165,17 +387,32 @@ func (d *dispatcher) init(rng *rand.Rand, nProcessors int) {
 	}
 }
 
+// close closes a dispatcher and its processors.
 func (d *dispatcher) close() {
+	d.mu.Lock()
+	d.closed = true
+	d.mu.Unlock()
 	for i := range d.processors {
 		d.processors[i].close()
 	}
 }
 
+// wait waits for all processor goroutines to end.
 func (d *dispatcher) wait() {
 	d.wg.Wait()
 }
 
+// queuePacket queues an incoming packet to the matching tcp endpoint and
+// also queues the endpoint to a processor queue for processing.
 func (d *dispatcher) queuePacket(stackEP stack.TransportEndpoint, id stack.TransportEndpointID, clock tcpip.Clock, pkt *stack.PacketBuffer) {
+	d.mu.Lock()
+	closed := d.closed
+	d.mu.Unlock()
+
+	if closed {
+		return
+	}
+
 	ep := stackEP.(*endpoint)
 
 	s := newIncomingSegment(id, clock, pkt)
@@ -202,25 +439,62 @@ func (d *dispatcher) queuePacket(stackEP stack.TransportEndpoint, id stack.Trans
 		return
 	}
 
-	// For sockets not in established state let the worker goroutine
-	// handle the packets.
-	if ep.EndpointState() != StateEstablished {
-		ep.newSegmentWaker.Assert()
-		return
-	}
-
 	d.selectProcessor(id).queueEndpoint(ep)
 }
 
+// selectProcessor uses a hash of the transport endpoint ID to queue the
+// endpoint to a specific processor. This is required to main TCP ordering as
+// queueing the same endpoint to multiple processors can *potentially* result in
+// out of order processing of incoming segments. It also ensures that a dispatcher
+// evenly loads the processor goroutines.
 func (d *dispatcher) selectProcessor(id stack.TransportEndpointID) *processor {
+	return &d.processors[d.hasher.hash(id)%uint32(len(d.processors))]
+}
+
+// pause pauses a dispatcher and all its processor goroutines.
+func (d *dispatcher) pause() {
+	d.mu.Lock()
+	d.paused = true
+	d.mu.Unlock()
+	for i := range d.processors {
+		<-d.processors[i].pause()
+	}
+}
+
+// resume resumes a previously paused dispatcher and its processor goroutines.
+// Calling resume on a dispatcher that was never paused is a no-op.
+func (d *dispatcher) resume() {
+	d.mu.Lock()
+
+	if !d.paused {
+		// If this was a restore run the stack is a new instance and
+		// it was never paused, so just return as there is nothing to
+		// resume.
+		d.mu.Unlock()
+		return
+	}
+	d.paused = false
+	d.mu.Unlock()
+	for i := range d.processors {
+		d.processors[i].resume()
+	}
+}
+
+// jenkinsHasher contains state needed to for a jenkins hash.
+type jenkinsHasher struct {
+	seed uint32
+}
+
+// hash hashes the provided TransportEndpointID using the jenkins hash
+// algorithm.
+func (j jenkinsHasher) hash(id stack.TransportEndpointID) uint32 {
 	var payload [4]byte
 	binary.LittleEndian.PutUint16(payload[0:], id.LocalPort)
 	binary.LittleEndian.PutUint16(payload[2:], id.RemotePort)
 
-	h := jenkins.Sum32(d.seed)
+	h := jenkins.Sum32(j.seed)
 	h.Write(payload[:])
 	h.Write([]byte(id.LocalAddress))
 	h.Write([]byte(id.RemoteAddress))
-
-	return &d.processors[h.Sum32()%uint32(len(d.processors))]
+	return h.Sum32()
 }