Fix a race where keepalives could be sent while there is pending data

PiperOrigin-RevId: 219571556
Change-Id: I5a1042c1cb05eb2711eb01627fd298bad6c543a6

3 files changed, 14 insertions, 2 deletions

diff --git a/pkg/tcpip/transport/tcp/BUILD b/pkg/tcpip/transport/tcp/BUILD
index 936d1064e..5a77ee232 100644
--- a/pkg/tcpip/transport/tcp/BUILD
+++ b/pkg/tcpip/transport/tcp/BUILD
@@ -62,8 +62,6 @@ go_test(
         "tcp_test.go",
         "tcp_timestamp_test.go",
     ],
-    # FIXME
-    tags = ["flaky"],
     deps = [
         ":tcp",
         "//pkg/tcpip",
diff --git a/pkg/tcpip/transport/tcp/connect.go b/pkg/tcpip/transport/tcp/connect.go
index 800d2409e..965779a68 100644
--- a/pkg/tcpip/transport/tcp/connect.go
+++ b/pkg/tcpip/transport/tcp/connect.go
@@ -829,6 +829,13 @@ func (e *endpoint) resetKeepaliveTimer(receivedData bool) {
 	}
 }
 
+// disableKeepaliveTimer stops the keepalive timer.
+func (e *endpoint) disableKeepaliveTimer() {
+	e.keepalive.Lock()
+	e.keepalive.timer.disable()
+	e.keepalive.Unlock()
+}
+
 // protocolMainLoop is the main loop of the TCP protocol. It runs in its own
 // goroutine and is responsible for sending segments and handling received
 // segments.
diff --git a/pkg/tcpip/transport/tcp/snd.go b/pkg/tcpip/transport/tcp/snd.go
index eefe93d48..b260c0e57 100644
--- a/pkg/tcpip/transport/tcp/snd.go
+++ b/pkg/tcpip/transport/tcp/snd.go
@@ -405,6 +405,7 @@ func (s *sender) sendData() {
 	// eventually.
 	var seg *segment
 	end := s.sndUna.Add(s.sndWnd)
+	var dataSent bool
 	for seg = s.writeNext; seg != nil && s.outstanding < s.sndCwnd; seg = seg.Next() {
 		// We abuse the flags field to determine if we have already
 		// assigned a sequence number to this segment.
@@ -448,6 +449,12 @@ func (s *sender) sendData() {
 			segEnd = seg.sequenceNumber.Add(seqnum.Size(seg.data.Size()))
 		}
 
+		if !dataSent {
+			dataSent = true
+			// We are sending data, so we should stop the keepalive timer to
+			// ensure that no keepalives are sent while there is pending data.
+			s.ep.disableKeepaliveTimer()
+		}
 		s.sendSegment(seg.data, seg.flags, seg.sequenceNumber)
 
 		// Update sndNxt if we actually sent new data (as opposed to