From c94521f51124b6e86a47b0a8c169e2b246725add Mon Sep 17 00:00:00 2001
From: Matt Joras <mjoras@fb.com>
Date: Wed, 8 Jul 2020 18:25:59 -0700
Subject: [PATCH] Optionally order read callbacks.

Summary: Without this option the read callbacks come in an  somewhat random order based on stream ID. This enforces ordering by stream ID, which is useful for some applications.

Reviewed By: yangchi

Differential Revision: D22442837

fbshipit-source-id: cd97de9760deff7b19841f46e23da2004df31492
---
 quic/api/QuicTransportBase.cpp          |  3 ++
 quic/api/test/QuicTransportBaseTest.cpp | 53 +++++++++++++++++++++++++
 quic/state/TransportSettings.h          |  2 +
 3 files changed, 58 insertions(+)

diff --git a/quic/api/QuicTransportBase.cpp b/quic/api/QuicTransportBase.cpp
index c3a63b0cf..9b77aa9f6 100644
--- a/quic/api/QuicTransportBase.cpp
+++ b/quic/api/QuicTransportBase.cpp
@@ -765,6 +765,9 @@ void QuicTransportBase::invokeReadDataAndCallbacks() {
       readableStreams.begin(),
       readableStreams.end(),
       std::back_inserter(readableStreamsCopy));
+  if (self->conn_->transportSettings.orderedReadCallbacks) {
+    std::sort(readableStreamsCopy.begin(), readableStreamsCopy.end());
+  }
   for (StreamId streamId : readableStreamsCopy) {
     auto callback = self->readCallbacks_.find(streamId);
     if (callback == self->readCallbacks_.end()) {
diff --git a/quic/api/test/QuicTransportBaseTest.cpp b/quic/api/test/QuicTransportBaseTest.cpp
index 372f81494..cf01ada77 100644
--- a/quic/api/test/QuicTransportBaseTest.cpp
+++ b/quic/api/test/QuicTransportBaseTest.cpp
@@ -529,6 +529,59 @@ TEST_F(QuicTransportImplTest, ReadCallbackDataAvailable) {
   transport.reset();
 }
 
+TEST_F(QuicTransportImplTest, ReadCallbackDataAvailableOrdered) {
+  auto transportSettings = transport->getTransportSettings();
+  transportSettings.orderedReadCallbacks = true;
+  transport->setTransportSettings(transportSettings);
+
+  auto stream1 = transport->createBidirectionalStream().value();
+  auto stream2 = transport->createBidirectionalStream().value();
+  StreamId stream3 = 0x6;
+
+  InSequence s;
+  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));
+  transport->setReadCallback(stream3, &readCb3);
+
+  EXPECT_CALL(readCb1, readAvailable(stream1));
+  EXPECT_CALL(readCb3, readAvailable(stream3));
+  transport->driveReadCallbacks();
+
+  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, ReadCallbackChangeReadCallback) {
   auto stream1 = transport->createBidirectionalStream().value();
 
diff --git a/quic/state/TransportSettings.h b/quic/state/TransportSettings.h
index c56fd6af7..229a12f74 100644
--- a/quic/state/TransportSettings.h
+++ b/quic/state/TransportSettings.h
@@ -164,6 +164,8 @@ struct TransportSettings {
   // Whether or not we should stop writing a packet after writing a single
   // stream frame to it.
   bool streamFramePerPacket{false};
+  // Ensure read callbacks are ordered by Stream ID.
+  bool orderedReadCallbacks{false};
 };
 
 } // namespace quic