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

quic/quic_transport/quic_transport_stream.cc

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 {

quic/quic_transport/quic_transport_stream_test.cc

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