Use WriteMemSlices() instead of WriteOrBufferData() in QuicTransportStream.
Originally, this was supposed to fix the fact that WriteOrBufferData() has unlimited buffering. However, as I wrote a test for that, I discovered that I've already added a check for this by accident when I originally wrote this code. I'm still switching this code to use WriteMemSlices() as it actually returns whether it has succeeded.
gfe-relnote: n/a (not used in production)
PiperOrigin-RevId: 288496889
Change-Id: I801c2f870dd04b9fa478e12452834ec38e20a870
diff --git a/quic/quic_transport/quic_transport_stream.cc b/quic/quic_transport/quic_transport_stream.cc
index ac54fa3..db90597 100644
--- a/quic/quic_transport/quic_transport_stream.cc
+++ b/quic/quic_transport/quic_transport_stream.cc
@@ -6,6 +6,8 @@
#include <sys/types.h>
+#include "net/third_party/quiche/src/quic/core/quic_buffer_allocator.h"
+#include "net/third_party/quiche/src/quic/core/quic_error_codes.h"
#include "net/third_party/quiche/src/quic/core/quic_types.h"
#include "net/third_party/quiche/src/quic/core/quic_utils.h"
#include "net/third_party/quiche/src/common/platform/api/quiche_string_piece.h"
@@ -54,9 +56,32 @@
return false;
}
- // TODO(vasilvv): use WriteMemSlices()
- WriteOrBufferData(data, /*fin=*/false, nullptr);
- return true;
+ QuicUniqueBufferPtr buffer = MakeUniqueBuffer(
+ session()->connection()->helper()->GetStreamSendBufferAllocator(),
+ data.size());
+ memcpy(buffer.get(), data.data(), data.size());
+ QuicMemSlice memslice(std::move(buffer), data.size());
+ QuicConsumedData consumed =
+ WriteMemSlices(QuicMemSliceSpan(&memslice), /*fin=*/false);
+
+ if (consumed.bytes_consumed == data.size()) {
+ return true;
+ }
+ if (consumed.bytes_consumed == 0) {
+ return false;
+ }
+ // QuicTransportStream::Write() is an all-or-nothing write API. To achieve
+ // that property, it relies on WriteMemSlices() being an all-or-nothing API.
+ // If WriteMemSlices() fails to provide that guarantee, we have no way to
+ // communicate a partial write to the caller, and thus it's safer to just
+ // close the connection.
+ QUIC_BUG << "WriteMemSlices() unexpectedly partially consumed the input "
+ "data, provided: "
+ << data.size() << ", written: " << consumed.bytes_consumed;
+ CloseConnectionWithDetails(
+ QUIC_INTERNAL_ERROR,
+ "WriteMemSlices() unexpectedly partially consumed the input data");
+ return false;
}
bool QuicTransportStream::SendFin() {
@@ -64,8 +89,11 @@
return false;
}
- WriteOrBufferData(quiche::QuicheStringPiece(), /*fin=*/true, nullptr);
- return true;
+ QuicMemSlice empty;
+ QuicConsumedData consumed =
+ WriteMemSlices(QuicMemSliceSpan(&empty), /*fin=*/true);
+ DCHECK_EQ(consumed.bytes_consumed, 0u);
+ return consumed.fin_consumed;
}
bool QuicTransportStream::CanWrite() const {
diff --git a/quic/quic_transport/quic_transport_stream_test.cc b/quic/quic_transport/quic_transport_stream_test.cc
index eda6af6..f4311ce 100644
--- a/quic/quic_transport/quic_transport_stream_test.cc
+++ b/quic/quic_transport/quic_transport_stream_test.cc
@@ -120,6 +120,22 @@
ASSERT_EQ(stream_->Read(&buffer), 4u);
}
+TEST_F(QuicTransportStreamTest, WritingTooMuchData) {
+ EXPECT_CALL(interface_, IsSessionReady()).WillRepeatedly(Return(true));
+ ASSERT_TRUE(stream_->CanWrite());
+
+ std::string a_little_bit_of_data(128, 'A');
+ std::string a_lot_of_data(GetQuicFlag(FLAGS_quic_buffered_data_threshold) * 2,
+ 'a');
+
+ EXPECT_TRUE(stream_->Write(a_little_bit_of_data));
+ EXPECT_TRUE(stream_->Write(a_little_bit_of_data));
+ EXPECT_TRUE(stream_->Write(a_little_bit_of_data));
+
+ EXPECT_TRUE(stream_->Write(a_lot_of_data));
+ EXPECT_FALSE(stream_->Write(a_lot_of_data));
+}
+
} // namespace
} // namespace test
} // namespace quic
