From 95ff9198ddb0d19ef3ac5c4f757887f7b347fc70 Mon Sep 17 00:00:00 2001 From: Luca Niccolini Date: Thu, 16 Dec 2021 11:43:58 -0800 Subject: [PATCH] Fix Datagram Write in case there is no space in the QUIC packet Summary: When there is not enough space in the QUIC packet to write a datagram frame, we should not dequeue the datagram Reviewed By: mjoras Differential Revision: D33109603 fbshipit-source-id: 937a5a0ffe55c7f88e39faf224e7ad06ca599708 --- quic/api/QuicPacketScheduler.cpp | 15 ++++++++++++--- quic/api/test/QuicPacketSchedulerTest.cpp | 16 ++++++++++++---- .../fizz/client/test/QuicClientTransportTest.cpp | 11 ++++++++--- quic/server/test/QuicServerTransportTest.cpp | 3 ++- 4 files changed, 34 insertions(+), 11 deletions(-) diff --git a/quic/api/QuicPacketScheduler.cpp b/quic/api/QuicPacketScheduler.cpp index 5fce23eb1..2177f1092 100644 --- a/quic/api/QuicPacketScheduler.cpp +++ b/quic/api/QuicPacketScheduler.cpp @@ -630,9 +630,18 @@ bool DatagramFrameScheduler::writeDatagramFrames( for (size_t i = 0; i <= conn_.datagramState.writeBuffer.size(); ++i) { auto& payload = conn_.datagramState.writeBuffer.front(); auto len = payload.chainLength(); - auto datagramFrame = DatagramFrame(len, payload.move()); - if (writeFrame(datagramFrame, builder) > 0) { - QUIC_STATS(conn_.statsCallback, onDatagramWrite, payload.chainLength()); + uint64_t spaceLeft = builder.remainingSpaceInPkt(); + QuicInteger frameTypeQuicInt(static_cast(FrameType::DATAGRAM_LEN)); + QuicInteger datagramLenInt(len); + auto datagramFrameLength = + frameTypeQuicInt.getSize() + len + datagramLenInt.getSize(); + if (folly::to(datagramFrameLength) <= spaceLeft) { + auto datagramFrame = DatagramFrame(len, payload.move()); + auto res = writeFrame(datagramFrame, builder); + // Must always succeed since we have already checked that there is enough + // space to write the frame + CHECK_GT(res, 0); + QUIC_STATS(conn_.statsCallback, onDatagramWrite, len); conn_.datagramState.writeBuffer.pop_front(); sent = true; } diff --git a/quic/api/test/QuicPacketSchedulerTest.cpp b/quic/api/test/QuicPacketSchedulerTest.cpp index dc68a6795..bebd21dea 100644 --- a/quic/api/test/QuicPacketSchedulerTest.cpp +++ b/quic/api/test/QuicPacketSchedulerTest.cpp @@ -1965,7 +1965,9 @@ TEST_F(QuicPacketSchedulerTest, DatagramFrameSchedulerMultipleFramesPerPacket) { })); NiceMock quicStats; conn.statsCallback = &quicStats; - EXPECT_CALL(quicStats, onDatagramWrite(_)).Times(2); + EXPECT_CALL(quicStats, onDatagramWrite(_)) + .Times(2) + .WillRepeatedly(Invoke([](uint64_t bytes) { EXPECT_GT(bytes, 0); })); // Call scheduler auto& frames = builder.frames_; scheduler.writeDatagramFrames(builder); @@ -1993,10 +1995,14 @@ TEST_F(QuicPacketSchedulerTest, DatagramFrameSchedulerOneFramePerPacket) { conn.statsCallback = &quicStats; // Call scheduler auto& frames = builder.frames_; - EXPECT_CALL(quicStats, onDatagramWrite(_)).Times(1); + EXPECT_CALL(quicStats, onDatagramWrite(_)) + .Times(1) + .WillRepeatedly(Invoke([](uint64_t bytes) { EXPECT_GT(bytes, 0); })); scheduler.writeDatagramFrames(builder); ASSERT_EQ(frames.size(), 1); - EXPECT_CALL(quicStats, onDatagramWrite(_)).Times(1); + EXPECT_CALL(quicStats, onDatagramWrite(_)) + .Times(1) + .WillRepeatedly(Invoke([](uint64_t bytes) { EXPECT_GT(bytes, 0); })); scheduler.writeDatagramFrames(builder); ASSERT_EQ(frames.size(), 2); } @@ -2025,7 +2031,9 @@ TEST_F(QuicPacketSchedulerTest, DatagramFrameWriteWhenRoomAvailable) { ASSERT_EQ(frames.size(), 0); EXPECT_CALL(builder, remainingSpaceInPkt()) .WillRepeatedly(Return(conn.udpSendPacketLen / 2)); - EXPECT_CALL(quicStats, onDatagramWrite(_)).Times(1); + EXPECT_CALL(quicStats, onDatagramWrite(_)) + .Times(1) + .WillRepeatedly(Invoke([](uint64_t bytes) { EXPECT_GT(bytes, 0); })); scheduler.writeDatagramFrames(builder); ASSERT_EQ(frames.size(), 1); } diff --git a/quic/fizz/client/test/QuicClientTransportTest.cpp b/quic/fizz/client/test/QuicClientTransportTest.cpp index d57adf75b..083348a90 100644 --- a/quic/fizz/client/test/QuicClientTransportTest.cpp +++ b/quic/fizz/client/test/QuicClientTransportTest.cpp @@ -4629,7 +4629,8 @@ TEST_F(QuicClientTransportAfterStartTest, ReceiveDatagramFrameAndStore) { conn.datagramState.maxReadBufferSize = 10; EXPECT_CALL(*quicStats_, onDatagramRead(_)) - .Times(conn.datagramState.maxReadBufferSize); + .Times(conn.datagramState.maxReadBufferSize) + .WillRepeatedly(Invoke([](uint64_t bytes) { EXPECT_GT(bytes, 0); })); EXPECT_CALL(*quicStats_, onDatagramDroppedOnRead()) .Times(conn.datagramState.maxReadBufferSize); for (uint64_t i = 0; i < conn.datagramState.maxReadBufferSize * 2; i++) { @@ -4679,7 +4680,9 @@ TEST_F( datagramPayload1.size(), IOBuf::copyBuffer(datagramPayload1)); writeFrame(datagramFrame1, builder1); auto packet1 = packetToBuf(std::move(builder1).buildPacket()); - EXPECT_CALL(*quicStats_, onDatagramRead(_)).Times(1); + EXPECT_CALL(*quicStats_, onDatagramRead(_)) + .Times(1) + .WillRepeatedly(Invoke([](uint64_t bytes) { EXPECT_GT(bytes, 0); })); deliverData(packet1->coalesce()); ASSERT_EQ( client->getConn().datagramState.readBuffer.size(), @@ -4698,7 +4701,9 @@ TEST_F( writeFrame(datagramFrame2, builder2); auto packet2 = packetToBuf(std::move(builder2).buildPacket()); EXPECT_CALL(*quicStats_, onDatagramDroppedOnRead()).Times(1); - EXPECT_CALL(*quicStats_, onDatagramRead(_)).Times(1); + EXPECT_CALL(*quicStats_, onDatagramRead(_)) + .Times(1) + .WillRepeatedly(Invoke([](uint64_t bytes) { EXPECT_GT(bytes, 0); })); deliverData(packet2->coalesce()); ASSERT_EQ( client->getConn().datagramState.readBuffer.size(), diff --git a/quic/server/test/QuicServerTransportTest.cpp b/quic/server/test/QuicServerTransportTest.cpp index c0edcd433..cac55788a 100644 --- a/quic/server/test/QuicServerTransportTest.cpp +++ b/quic/server/test/QuicServerTransportTest.cpp @@ -2750,7 +2750,8 @@ TEST_F(QuicServerTransportTest, ReceiveDatagramFrameAndStore) { conn.datagramState.maxReadBufferSize = 10; EXPECT_CALL(*quicStats_, onDatagramRead(_)) - .Times(conn.datagramState.maxReadBufferSize); + .Times(conn.datagramState.maxReadBufferSize) + .WillRepeatedly(Invoke([](uint64_t bytes) { EXPECT_GT(bytes, 0); })); EXPECT_CALL(*quicStats_, onDatagramDroppedOnRead()) .Times(conn.datagramState.maxReadBufferSize); for (uint64_t i = 0; i < conn.datagramState.maxReadBufferSize * 2; i++) {