From 31f9d93f45d57b145dab2b66a20fe56e77cc96ab Mon Sep 17 00:00:00 2001
From: Yang Chi <yangchi@fb.com>
Date: Thu, 1 Apr 2021 10:40:39 -0700
Subject: [PATCH] Remove WriteBufferMeta from retransmission buffer upon ack

Summary: as title

Reviewed By: mjoras

Differential Revision: D27361934

fbshipit-source-id: 81618c039ff026a8b3dc2caba5ce335ee431e15e
---
 quic/state/stream/StreamSendHandlers.cpp      | 42 +++++++++++++------
 .../stream/test/StreamStateMachineTest.cpp    | 29 +++++++++++++
 2 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/quic/state/stream/StreamSendHandlers.cpp b/quic/state/stream/StreamSendHandlers.cpp
index aadae772f..b444137a2 100644
--- a/quic/state/stream/StreamSendHandlers.cpp
+++ b/quic/state/stream/StreamSendHandlers.cpp
@@ -97,18 +97,36 @@ void sendAckSMHandler(
     const WriteStreamFrame& ackedFrame) {
   switch (stream.sendState) {
     case StreamSendState::Open: {
-      // Clean up the acked buffers from the retransmissionBuffer.
-      auto ackedBuffer = stream.retransmissionBuffer.find(ackedFrame.offset);
-      if (ackedBuffer != stream.retransmissionBuffer.end()) {
-        VLOG(10) << "Open: acked stream data stream=" << stream.id
-                 << " offset=" << ackedBuffer->second->offset
-                 << " len=" << ackedBuffer->second->data.chainLength()
-                 << " eof=" << ackedBuffer->second->eof << " " << stream.conn;
-        stream.ackedIntervals.insert(
-            ackedBuffer->second->offset,
-            ackedBuffer->second->offset +
-                ackedBuffer->second->data.chainLength());
-        stream.retransmissionBuffer.erase(ackedBuffer);
+      if (!ackedFrame.fromBufMeta) {
+        // Clean up the acked buffers from the retransmissionBuffer.
+        auto ackedBuffer = stream.retransmissionBuffer.find(ackedFrame.offset);
+        if (ackedBuffer != stream.retransmissionBuffer.end()) {
+          VLOG(10) << "Open: acked stream data stream=" << stream.id
+                   << " offset=" << ackedBuffer->second->offset
+                   << " len=" << ackedBuffer->second->data.chainLength()
+                   << " eof=" << ackedBuffer->second->eof << " " << stream.conn;
+          stream.ackedIntervals.insert(
+              ackedBuffer->second->offset,
+              ackedBuffer->second->offset +
+                  ackedBuffer->second->data.chainLength());
+          stream.retransmissionBuffer.erase(ackedBuffer);
+        }
+      } else {
+        auto ackedBuffer =
+            stream.retransmissionBufMetas.find(ackedFrame.offset);
+        if (ackedBuffer != stream.retransmissionBufMetas.end()) {
+          CHECK_EQ(ackedFrame.offset, ackedBuffer->second.offset);
+          CHECK_EQ(ackedFrame.len, ackedBuffer->second.length);
+          CHECK_EQ(ackedFrame.fin, ackedBuffer->second.eof);
+          VLOG(10) << "Open: acked stream data bufmeta=" << stream.id
+                   << " offset=" << ackedBuffer->second.offset
+                   << " len=" << ackedBuffer->second.length
+                   << " eof=" << ackedBuffer->second.eof << " " << stream.conn;
+          stream.ackedIntervals.insert(
+              ackedBuffer->second.offset,
+              ackedBuffer->second.offset + ackedBuffer->second.length);
+          stream.retransmissionBufMetas.erase(ackedBuffer);
+        }
       }
 
       // This stream may be able to invoke some deliveryCallbacks:
diff --git a/quic/state/stream/test/StreamStateMachineTest.cpp b/quic/state/stream/test/StreamStateMachineTest.cpp
index f8ccdf7bd..0a7e1271e 100644
--- a/quic/state/stream/test/StreamStateMachineTest.cpp
+++ b/quic/state/stream/test/StreamStateMachineTest.cpp
@@ -829,5 +829,34 @@ TEST_F(QuicUnidirectionalStreamTest, ResetSentInvalidRstStream) {
       QuicTransportException);
 }
 
+TEST_F(QuicOpenStateTest, DSRStreamAcked) {
+  auto conn = createConn();
+  conn->clientConnectionId = getTestConnectionId(0);
+  conn->serverConnectionId = getTestConnectionId(1);
+  auto stream = conn->streamManager->createNextBidirectionalStream().value();
+  writeDataToQuicStream(
+      *stream,
+      folly::IOBuf::copyBuffer("Big ship stucks in small water"),
+      false);
+  writeBufMetaToQuicStream(*stream, BufferMeta(1000), true);
+  auto bufMetaStartingOffset = stream->writeBufMeta.offset;
+  handleStreamWritten(
+      *conn,
+      *stream,
+      bufMetaStartingOffset,
+      300,
+      false,
+      1,
+      PacketNumberSpace::AppData,
+      true);
+  ASSERT_NE(
+      stream->retransmissionBufMetas.end(),
+      stream->retransmissionBufMetas.find(bufMetaStartingOffset));
+  WriteStreamFrame frame(stream->id, bufMetaStartingOffset, 300, false);
+  frame.fromBufMeta = true;
+  sendAckSMHandler(*stream, frame);
+  EXPECT_TRUE(stream->retransmissionBufMetas.empty());
+}
+
 } // namespace test
 } // namespace quic