From 027fedad5bd44541c9e18aa84c889f6d848ec669 Mon Sep 17 00:00:00 2001
From: Matt Joras <mjoras@fb.com>
Date: Tue, 29 Oct 2019 12:32:11 -0700
Subject: [PATCH] Search loss buffer before retransmission buffer when
 determining if a write is a clone.

Summary:
We maintain the invariant that a buffer cannot be in the loss buffer and retransmission buffers at the same time. As the retransmission buffer holds all unacknowledged data that isn't marked lost, it is very likely to be larger than the loss buffer. This makes the existing case to check for cloning very expensive.

Instead search the loss buffer first, change the search of the loss buffer to a binary search, and elide the double search.

Reviewed By: yangchi

Differential Revision: D18203444

fbshipit-source-id: 66a4e424d61c4b0e3cad12c7eca009ad3d6c5a0d
---
 quic/api/QuicTransportFunctions.cpp | 42 ++++++++++++++---------------
 quic/api/QuicTransportFunctions.h   |  3 ++-
 quic/api/test/QuicTransportTest.cpp |  2 +-
 3 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/quic/api/QuicTransportFunctions.cpp b/quic/api/QuicTransportFunctions.cpp
index f1e4ba51e..5fe5bcf10 100644
--- a/quic/api/QuicTransportFunctions.cpp
+++ b/quic/api/QuicTransportFunctions.cpp
@@ -111,13 +111,9 @@ void handleRetransmissionWritten(
     uint64_t frameOffset,
     uint64_t frameLen,
     bool frameFin,
-    PacketNum packetNum) {
+    PacketNum packetNum,
+    std::deque<StreamBuffer>::iterator lossBufferIter) {
   conn.lossState.totalBytesRetransmitted += frameLen;
-  auto lossBufferIter = std::find_if(
-      stream.lossBuffer.begin(),
-      stream.lossBuffer.end(),
-      [&](const auto& buffer) { return frameOffset == buffer.offset; });
-  CHECK(lossBufferIter != stream.lossBuffer.end());
   VLOG(10) << nodeToString(conn.nodeType) << " sent retransmission"
            << " packetNum=" << packetNum << " " << conn;
   auto bufferLen = lossBufferIter->data.chainLength();
@@ -163,24 +159,28 @@ bool handleStreamWritten(
     return true;
   }
 
-  // If the data is in retx buffer, this is a clone write
-  auto retxBufferIter = std::find_if(
-      stream.retransmissionBuffer.begin(),
-      stream.retransmissionBuffer.end(),
-      [&](const auto& buffer) {
-        return frameOffset == buffer.offset &&
-            frameLen == buffer.data.chainLength() && frameFin == buffer.eof;
-      });
-  if (retxBufferIter != stream.retransmissionBuffer.end()) {
-    conn.lossState.totalStreamBytesCloned += frameLen;
+  // If the data is in the loss buffer, it is a retransmission.
+  auto lossBufferIter = std::lower_bound(
+      stream.lossBuffer.begin(),
+      stream.lossBuffer.end(),
+      frameOffset,
+      [](const auto& buf, auto off) { return buf.offset < off; });
+  if (lossBufferIter != stream.lossBuffer.end() &&
+      lossBufferIter->offset == frameOffset) {
+    handleRetransmissionWritten(
+        conn,
+        stream,
+        frameOffset,
+        frameLen,
+        frameFin,
+        packetNum,
+        lossBufferIter);
+    QUIC_STATS(conn.infoCallback, onPacketRetransmission);
     return false;
   }
 
-  // If it's neither new data nor clone data, then it is a retransmission and
-  // the data has to be in loss buffer.
-  handleRetransmissionWritten(
-      conn, stream, frameOffset, frameLen, frameFin, packetNum);
-  QUIC_STATS(conn.infoCallback, onPacketRetransmission);
+  // Otherwise it must be a clone write.
+  conn.lossState.totalStreamBytesCloned += frameLen;
   return false;
 }
 
diff --git a/quic/api/QuicTransportFunctions.h b/quic/api/QuicTransportFunctions.h
index 0518f25b2..caf253c96 100644
--- a/quic/api/QuicTransportFunctions.h
+++ b/quic/api/QuicTransportFunctions.h
@@ -119,7 +119,8 @@ void handleRetransmissionWritten(
     uint64_t frameOffset,
     uint64_t frameLen,
     bool frameFin,
-    PacketNum packetNum);
+    PacketNum packetNum,
+    std::deque<StreamBuffer>::iterator lossBufferIter);
 
 /**
  * Update the connection and stream state after stream data is written and deal
diff --git a/quic/api/test/QuicTransportTest.cpp b/quic/api/test/QuicTransportTest.cpp
index 1a6c6fe6e..cc7b55edf 100644
--- a/quic/api/test/QuicTransportTest.cpp
+++ b/quic/api/test/QuicTransportTest.cpp
@@ -247,7 +247,7 @@ void dropPackets(QuicServerConnectionState& conn) {
           [&streamFrame](const auto& buffer) {
             return streamFrame->offset == buffer.offset;
           });
-      EXPECT_TRUE(itr != stream->retransmissionBuffer.end());
+      ASSERT_TRUE(itr != stream->retransmissionBuffer.end());
       stream->lossBuffer.insert(
           std::upper_bound(
               stream->lossBuffer.begin(),