From 07af79fc654202c758e399287ee723b32ac98232 Mon Sep 17 00:00:00 2001
From: Matt Joras <mjoras@fb.com>
Date: Thu, 8 Jul 2021 10:29:39 -0700
Subject: [PATCH] Don't remove stream from readable if the read callback isn't
 set.

Summary: This probably never made much sense as a policy. The application can reasonably expect to be able to set a read callback after the stream first becomes readable.

Reviewed By: afrind

Differential Revision: D29563483

fbshipit-source-id: 215a6aed130108248a0b4cbc36c02c8da70bf721
---
 quic/api/QuicTransportBase.cpp          |  2 +-
 quic/api/test/QuicTransportBaseTest.cpp | 46 +++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/quic/api/QuicTransportBase.cpp b/quic/api/QuicTransportBase.cpp
index 7f223342a..eebda2af7 100644
--- a/quic/api/QuicTransportBase.cpp
+++ b/quic/api/QuicTransportBase.cpp
@@ -798,7 +798,7 @@ void QuicTransportBase::invokeReadDataAndCallbacks() {
   for (StreamId streamId : readableStreamsCopy) {
     auto callback = self->readCallbacks_.find(streamId);
     if (callback == self->readCallbacks_.end()) {
-      self->conn_->streamManager->readableStreams().erase(streamId);
+      // Stream doesn't have a read callback set, skip it.
       continue;
     }
     auto readCb = callback->second.readCb;
diff --git a/quic/api/test/QuicTransportBaseTest.cpp b/quic/api/test/QuicTransportBaseTest.cpp
index 564ed3c9e..045f2eb1e 100644
--- a/quic/api/test/QuicTransportBaseTest.cpp
+++ b/quic/api/test/QuicTransportBaseTest.cpp
@@ -652,6 +652,52 @@ TEST_F(QuicTransportImplTest, ReadCallbackDataAvailable) {
   transport.reset();
 }
 
+TEST_F(QuicTransportImplTest, ReadCallbackDataAvailableNoReap) {
+  auto stream1 = transport->createBidirectionalStream().value();
+  auto stream2 = transport->createBidirectionalStream().value();
+  StreamId stream3 = 0x6;
+
+  NiceMock<MockReadCallback> readCb1;
+  NiceMock<MockReadCallback> readCb2;
+  NiceMock<MockReadCallback> readCb3;
+
+  transport->setReadCallback(stream1, &readCb1);
+  transport->setReadCallback(stream2, &readCb2);
+
+  transport->addDataToStream(
+      stream1, StreamBuffer(folly::IOBuf::copyBuffer("actual stream data"), 0));
+
+  transport->addDataToStream(
+      stream2,
+      StreamBuffer(folly::IOBuf::copyBuffer("actual stream data"), 10));
+
+  transport->addDataToStream(
+      stream3, StreamBuffer(folly::IOBuf::copyBuffer("actual stream data"), 0));
+
+  EXPECT_CALL(readCb1, readAvailable(stream1));
+  transport->driveReadCallbacks();
+
+  transport->setReadCallback(stream3, &readCb3);
+  transport->addDataToStream(
+      stream2, StreamBuffer(folly::IOBuf::copyBuffer("actual stream data"), 0));
+
+  EXPECT_CALL(readCb1, readAvailable(stream1));
+  EXPECT_CALL(readCb2, readAvailable(stream2));
+  EXPECT_CALL(readCb3, readAvailable(stream3));
+  transport->driveReadCallbacks();
+
+  EXPECT_CALL(readCb1, readAvailable(stream1));
+  EXPECT_CALL(readCb2, readAvailable(stream2));
+  EXPECT_CALL(readCb3, readAvailable(stream3));
+  transport->driveReadCallbacks();
+
+  EXPECT_CALL(readCb2, readAvailable(stream2));
+  EXPECT_CALL(readCb3, readAvailable(stream3));
+  transport->setReadCallback(stream1, nullptr);
+  transport->driveReadCallbacks();
+  transport.reset();
+}
+
 TEST_F(QuicTransportImplTest, ReadCallbackDataAvailableOrdered) {
   auto transportSettings = transport->getTransportSettings();
   transportSettings.orderedReadCallbacks = true;