sentry: skip waiting for undrain for netstack TCP endpoints in error state.

PiperOrigin-RevId: 224214981
Change-Id: I4c1dd5b1c856f7a4f9866a5dda44a5297e92486a

1 files changed, 18 insertions, 8 deletions

diff --git a/pkg/tcpip/transport/tcp/connect.go b/pkg/tcpip/transport/tcp/connect.go
index 00cb39560..0e6bb6763 100644
--- a/pkg/tcpip/transport/tcp/connect.go
+++ b/pkg/tcpip/transport/tcp/connect.go
@@ -976,25 +976,35 @@ func (e *endpoint) protocolMainLoop(handshake bool) *tcpip.Error {
 					e.mu.Unlock()
 				}
 				if n&notifyClose != 0 && closeTimer == nil {
-					// Reset the connection 3 seconds after the
-					// endpoint has been closed.
+					// Reset the connection 3 seconds after
+					// the endpoint has been closed.
+					//
+					// The timer could fire in background
+					// when the endpoint is drained. That's
+					// OK as the loop here will not honor
+					// the firing until the undrain arrives.
 					closeTimer = time.AfterFunc(3*time.Second, func() {
 						closeWaker.Assert()
 					})
 				}
 
+				if n&notifyKeepaliveChanged != 0 {
+					// The timer could fire in background
+					// when the endpoint is drained. That's
+					// OK. See above.
+					e.resetKeepaliveTimer(true)
+				}
+
 				if n&notifyDrain != 0 {
 					for !e.segmentQueue.empty() {
 						if err := e.handleSegments(); err != nil {
 							return err
 						}
 					}
-					close(e.drainDone)
-					<-e.undrain
-				}
-
-				if n&notifyKeepaliveChanged != 0 {
-					e.resetKeepaliveTimer(true)
+					if e.state != stateError {
+						close(e.drainDone)
+						<-e.undrain
+					}
 				}
 
 				return nil