Implement actual zero-copy in EncapsulatedWebTransport. The existing implementation always copies data, which dates back to the time when WebTransport API did not accept memslices. Now that it does, we can avoid copying here. PiperOrigin-RevId: 931030987
diff --git a/quiche/web_transport/encapsulated/encapsulated_web_transport.cc b/quiche/web_transport/encapsulated/encapsulated_web_transport.cc index ccdf6bb..2ed07ff 100644 --- a/quiche/web_transport/encapsulated/encapsulated_web_transport.cc +++ b/quiche/web_transport/encapsulated/encapsulated_web_transport.cc
@@ -18,9 +18,11 @@ #include <utility> #include <vector> +#include "absl/container/inlined_vector.h" #include "absl/container/node_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/cord.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/time/time.h" @@ -32,9 +34,9 @@ #include "quiche/common/quiche_buffer_allocator.h" #include "quiche/common/quiche_callbacks.h" #include "quiche/common/quiche_circular_deque.h" +#include "quiche/common/quiche_cord_utils.h" #include "quiche/common/quiche_mem_slice.h" #include "quiche/common/quiche_status_utils.h" -#include "quiche/common/vectorized_io_utils.h" #include "quiche/web_transport/stream_helpers.h" #include "quiche/web_transport/web_transport.h" @@ -325,8 +327,9 @@ OnFatalError("Next scheduled stream is not in the map"); return; } - QUICHE_DCHECK(it->second.HasPendingWrite()); - it->second.FlushPendingWrite(); + auto& [id, stream] = *it; + QUICHE_DCHECK(stream.HasDataToWrite()); + stream.FlushWriteBuffer(/*notify_visitor=*/true); } } @@ -607,13 +610,6 @@ absl::Status EncapsulatedSession::InnerStream::Writev( const absl::Span<quiche::QuicheMemSlice> data, const StreamWriteOptions& options) { - // TODO: support zero copy. - std::vector<absl::string_view> views; - views.reserve(data.size()); - for (const quiche::QuicheMemSlice& slice : data) { - views.push_back(slice.AsStringView()); - } - if (write_side_closed_) { return absl::FailedPreconditionError( "Trying to write into an already-closed stream"); @@ -626,6 +622,10 @@ "Trying to write into a stream when CanWrite() = false"); } + const bool has_preexisting_buffered_data = !write_buffer_.empty(); + write_buffer_.Append(quiche::MemSliceSpanToCord(data)); + fin_buffered_ = options.send_fin(); + const absl::StatusOr<bool> should_yield = session_->scheduler_.ShouldYield(id_); if (!should_yield.ok()) { @@ -634,12 +634,8 @@ return absl::InternalError("Stream not registered with the scheduler"); } const bool write_blocked = !session_->underlying_->CanWrite() || - *should_yield || !pending_write_.empty(); + *should_yield || has_preexisting_buffered_data; if (write_blocked) { - fin_buffered_ = options.send_fin(); - for (absl::string_view chunk : views) { - absl::StrAppend(&pending_write_, chunk); - } absl::Status status = session_->scheduler_.Schedule(id_); if (!status.ok()) { QUICHE_BUG(WT_H2_Writev_CantSchedule) << status; @@ -649,72 +645,50 @@ return absl::OkStatus(); } - size_t bytes_written = WriteInner(views, options.send_fin()); - // TODO: handle partial writes when flow control requires those. - QUICHE_DCHECK(bytes_written == 0 || - bytes_written == quiche::TotalStringViewSpanSize(views)); - if (bytes_written == 0) { - for (absl::string_view chunk : views) { - absl::StrAppend(&pending_write_, chunk); - } - } - - if (options.send_fin()) { - CloseWriteSide(std::nullopt); - } + FlushWriteBuffer(/*notify_visitor=*/false); return absl::OkStatus(); } bool EncapsulatedSession::InnerStream::CanWrite() const { return session_->state_ != EncapsulatedSession::kSessionClosed && !write_side_closed_ && - (pending_write_.size() <= session_->max_stream_data_buffered_); + (write_buffer_.size() <= session_->max_stream_data_buffered_); } -void EncapsulatedSession::InnerStream::FlushPendingWrite() { +void EncapsulatedSession::InnerStream::FlushWriteBuffer(bool notify_visitor) { QUICHE_DCHECK(!write_side_closed_); QUICHE_DCHECK(session_->underlying_->CanWrite()); - QUICHE_DCHECK(!pending_write_.empty()); - absl::string_view to_write = pending_write_; - size_t bytes_written = - WriteInner(absl::MakeSpan(&to_write, 1), fin_buffered_); - if (bytes_written < to_write.size()) { - pending_write_ = pending_write_.substr(bytes_written); + const size_t bytes_to_write = write_buffer_.size(); + // TODO(vasilvv): adjust this value so that it does not exceed the flow + // control limit. + const bool fin = fin_buffered_ && (bytes_to_write == write_buffer_.size()); + + if (bytes_to_write == 0 && !fin) { + // Nothing to write. return; } - pending_write_.clear(); - if (fin_buffered_) { - CloseWriteSide(std::nullopt); - } - if (!write_side_closed_ && visitor_ != nullptr) { - visitor_->OnCanWrite(); - } -} -size_t EncapsulatedSession::InnerStream::WriteInner( - absl::Span<const absl::string_view> data, bool fin) { - size_t total_size = quiche::TotalStringViewSpanSize(data); - if (total_size == 0 && !fin) { - session_->OnFatalError("Attempted to make an empty write with fin=false"); - return 0; - } quiche::QuicheBuffer header = - quiche::SerializeWebTransportStreamCapsuleHeader(id_, fin, total_size, + quiche::SerializeWebTransportStreamCapsuleHeader(id_, fin, bytes_to_write, session_->allocator_); - std::vector<quiche::QuicheMemSlice> views_to_write; - views_to_write.reserve(data.size() + 1); - views_to_write.push_back(quiche::QuicheMemSlice(std::move(header))); - for (absl::string_view view : data) { - // TODO: support zero copy. - views_to_write.push_back(quiche::QuicheMemSlice::Copy(view)); - } + absl::InlinedVector<quiche::QuicheMemSlice, 8> slices; + slices.push_back(quiche::QuicheMemSlice(std::move(header))); + quiche::CordToMemSlicesTo(write_buffer_.Subcord(0, bytes_to_write), slices); absl::Status write_status = session_->underlying_->Writev( - absl::MakeSpan(views_to_write), kDefaultStreamWriteOptions); + absl::MakeSpan(slices), kDefaultStreamWriteOptions); if (!write_status.ok()) { session_->OnWriteError(write_status); - return 0; + return; } - return total_size; + write_buffer_.RemovePrefix(bytes_to_write); + + if (fin) { + CloseWriteSide(std::nullopt); + } + if (notify_visitor && !write_side_closed_ && visitor_ != nullptr && + CanWrite()) { + visitor_->OnCanWrite(); + } } void EncapsulatedSession::InnerStream::ResetWithUserCode( @@ -761,7 +735,7 @@ return; } write_side_closed_ = true; - pending_write_.clear(); + write_buffer_.Clear(); absl::Status status = session_->scheduler_.Unregister(id_); if (!status.ok()) { session_->OnFatalError("Failed to unregister closed stream");
diff --git a/quiche/web_transport/encapsulated/encapsulated_web_transport.h b/quiche/web_transport/encapsulated/encapsulated_web_transport.h index cb0d416..d711ad0 100644 --- a/quiche/web_transport/encapsulated/encapsulated_web_transport.h +++ b/quiche/web_transport/encapsulated/encapsulated_web_transport.h
@@ -16,6 +16,7 @@ #include "absl/base/attributes.h" #include "absl/container/node_hash_map.h" #include "absl/status/status.h" +#include "absl/strings/cord.h" #include "absl/strings/string_view.h" #include "absl/time/time.h" #include "absl/types/span.h" @@ -173,8 +174,10 @@ return read_side_closed_ && write_side_closed_; } - bool HasPendingWrite() const { return !pending_write_.empty(); } - void FlushPendingWrite(); + bool HasDataToWrite() const { + return !write_buffer_.empty() || fin_buffered_; + } + void FlushWriteBuffer(bool notify_visitor); void ProcessCapsule(const quiche::Capsule& capsule); @@ -195,15 +198,17 @@ size_t size() const { return data.size(); } }; - // Tries to send `data`; may send less if limited by flow control. - [[nodiscard]] size_t WriteInner(absl::Span<const absl::string_view> data, - bool fin); - EncapsulatedSession* session_; StreamId id_; std::unique_ptr<StreamVisitor> visitor_; quiche::QuicheCircularDeque<IncomingRead> incoming_reads_; - std::string pending_write_; + + // The write buffer for an encapsulated HTTP/2 stream is an absl::Cord + // wrapping multiple memslices around. Due to the flow control limitations, + // we cannot always pass memslices to the lower layer as-is, thus we need to + // split them by refcounting; absl::Cord does that automatically for us. + absl::Cord write_buffer_; + bool read_side_closed_; bool write_side_closed_; bool reset_frame_sent_ = false;
diff --git a/quiche/web_transport/encapsulated/encapsulated_web_transport_test.cc b/quiche/web_transport/encapsulated/encapsulated_web_transport_test.cc index f80d932..c8eddf0 100644 --- a/quiche/web_transport/encapsulated/encapsulated_web_transport_test.cc +++ b/quiche/web_transport/encapsulated/encapsulated_web_transport_test.cc
@@ -337,6 +337,27 @@ session->NotifySessionDraining(); } +TEST_F(EncapsulatedWebTransportTest, WriteErrorStream) { + std::unique_ptr<EncapsulatedSession> session = + CreateTransport(Perspective::kClient); + DefaultHandshakeForClient(*session); + Stream* stream = session->OpenOutgoingUnidirectionalStream(); + ASSERT_TRUE(stream != nullptr); + + // Let the CanWrite() check succeed, but the actual write fail. + EXPECT_CALL(underlying_, GetWriteStatus(/*is_write=*/false)) + .WillRepeatedly(Return(absl::OkStatus())); + EXPECT_CALL(underlying_, GetWriteStatus(/*is_write=*/true)) + .WillOnce(Return(absl::InternalError("Test write error"))); + EXPECT_CALL(fatal_error_callback_, Call) + .WillOnce([](absl::string_view error) { + EXPECT_THAT(error, HasSubstr("Test write error")); + }); + absl::Status status = WriteIntoStream(*stream, "test"); + QUICHE_EXPECT_OK(status); + EXPECT_EQ(session->state(), EncapsulatedSession::kSessionClosed); +} + TEST_F(EncapsulatedWebTransportTest, SimpleRead) { std::unique_ptr<EncapsulatedSession> session = CreateTransport(Perspective::kClient); @@ -573,6 +594,9 @@ EXPECT_EQ(capsule.web_transport_stream_data().data, "test"); return true; }); + auto visitor = std::make_unique<MockStreamVisitor>(); + EXPECT_CALL(*visitor, OnCanWrite).Times(0); + stream->SetVisitor(std::move(visitor)); absl::Status status = WriteIntoStream(*stream, "test"); EXPECT_THAT(status, StatusIs(absl::StatusCode::kOk)); } @@ -674,6 +698,9 @@ EXPECT_EQ(capsule.web_transport_stream_data().data, "abcdef"); return true; }); + auto visitor = std::make_unique<MockStreamVisitor>(); + EXPECT_CALL(*visitor, OnCanWrite); + stream->SetVisitor(std::move(visitor)); session_->OnCanWrite(); }
diff --git a/quiche/web_transport/test_tools/in_memory_stream.cc b/quiche/web_transport/test_tools/in_memory_stream.cc index 87f49a8..ad2236e 100644 --- a/quiche/web_transport/test_tools/in_memory_stream.cc +++ b/quiche/web_transport/test_tools/in_memory_stream.cc
@@ -60,7 +60,7 @@ absl::Status InMemoryStream::Writev( absl::Span<quiche::QuicheMemSlice> data, const webtransport::StreamWriteOptions& options) { - absl::Status status = GetWriteStatusWithExtraChecks(); + absl::Status status = GetWriteStatusWithExtraChecks(/*is_write=*/true); if (!status.ok()) { return status; } @@ -93,17 +93,18 @@ fin_received_ = false; } -absl::Status InMemoryStream::GetWriteStatus() const { +absl::Status InMemoryStream::GetWriteStatus(bool /*is_write*/) const { return absl::UnimplementedError( "Writing not implemented; use InMemoryStreamWithMockWrite"); } -absl::Status InMemoryStream::GetWriteStatusWithExtraChecks() const { +absl::Status InMemoryStream::GetWriteStatusWithExtraChecks( + bool is_write) const { if (fin_sent_) { return absl::FailedPreconditionError( "Can't write on a stream with FIN sent."); } - return GetWriteStatus(); + return GetWriteStatus(is_write); } InMemoryStreamWithMockWrite::InMemoryStreamWithMockWrite(StreamId id)
diff --git a/quiche/web_transport/test_tools/in_memory_stream.h b/quiche/web_transport/test_tools/in_memory_stream.h index 340fe7b..0cfe34e 100644 --- a/quiche/web_transport/test_tools/in_memory_stream.h +++ b/quiche/web_transport/test_tools/in_memory_stream.h
@@ -40,7 +40,7 @@ absl::Status Writev(absl::Span<quiche::QuicheMemSlice> data, const StreamWriteOptions& options) override; bool CanWrite() const override { - return GetWriteStatusWithExtraChecks().ok(); + return GetWriteStatusWithExtraChecks(/*is_write=*/false).ok(); } // webtransport::Stream implementation. @@ -81,11 +81,11 @@ protected: virtual void OnWrite(absl::string_view data) {} virtual void OnFin() {} - virtual absl::Status GetWriteStatus() const; + virtual absl::Status GetWriteStatus(bool is_write) const; private: void Terminate(); - absl::Status GetWriteStatusWithExtraChecks() const; + absl::Status GetWriteStatusWithExtraChecks(bool is_write) const; StreamId id_; std::unique_ptr<StreamVisitor> visitor_; @@ -105,7 +105,7 @@ MOCK_METHOD(void, OnWrite, (absl::string_view data), (override)); MOCK_METHOD(void, OnFin, (), (override)); - MOCK_METHOD(absl::Status, GetWriteStatus, (), (const, override)); + MOCK_METHOD(absl::Status, GetWriteStatus, (bool is_write), (const, override)); }; // An InMemoryStream where all writes are stored into a buffer. @@ -114,7 +114,9 @@ using InMemoryStream::InMemoryStream; void OnWrite(absl::string_view data) { write_buffer_.append(data); } - absl::Status GetWriteStatus() const { return absl::OkStatus(); } + absl::Status GetWriteStatus(bool is_write) const override { + return absl::OkStatus(); + } std::string& write_buffer() { return write_buffer_; }
diff --git a/quiche/web_transport/test_tools/in_memory_stream_test.cc b/quiche/web_transport/test_tools/in_memory_stream_test.cc index 8068dfb..e905eed 100644 --- a/quiche/web_transport/test_tools/in_memory_stream_test.cc +++ b/quiche/web_transport/test_tools/in_memory_stream_test.cc
@@ -86,10 +86,14 @@ TEST(InMemoryStreamTest, InMemoryStreamWithMockWrite) { InMemoryStreamWithMockWrite stream(0); + EXPECT_CALL(stream, GetWriteStatus(false)) + .WillOnce(testing::Return(absl::OkStatus())); EXPECT_TRUE(stream.CanWrite()); std::array write_vector = {quiche::QuicheMemSlice::Copy("test")}; EXPECT_CALL(stream, OnWrite("test")); + EXPECT_CALL(stream, GetWriteStatus(true)) + .WillOnce(testing::Return(absl::OkStatus())); StreamWriteOptions options; QUICHE_EXPECT_OK(stream.Writev(absl::MakeSpan(write_vector), options)); EXPECT_FALSE(stream.fin_sent()); @@ -102,6 +106,8 @@ EXPECT_CALL(stream, OnWrite("test2")); EXPECT_CALL(stream, OnFin()); } + EXPECT_CALL(stream, GetWriteStatus(true)) + .WillOnce(testing::Return(absl::OkStatus())); QUICHE_EXPECT_OK(stream.Writev(absl::MakeSpan(write_vector), options)); EXPECT_TRUE(stream.fin_sent()); EXPECT_FALSE(stream.CanWrite());