Reland "Search Prefetch Streaming requests can be cancelled after headers"

This is a reland of 583f52a3

See https://chromium-review.googlesource.com/c/chromium/src/+/2585991/2
for underlying MSAN issue unrelated to this CL (just caught by this
CL's tests)

Original change's description:
> Search Prefetch Streaming requests can be cancelled after headers
>
> The goal of this CL is to allow unneeded requests to be cancelled
> when they are unlikely to be navigated to. This is achieved through
> a few avenues of change:
>
> -- When the user navigates to something in omnibox, the URL is reported
> to Search Prefetch Service to be able to mark it as likely to be served
> (using kCanBeServedAndUserClicked when the request is kCanBeServed).
> This occurs before the omnibox is closed and all results are wiped.
> -- kCanBeServed are wiped when omnibox closes.
> -- Streaming requests will no longer pause the mojo channel, so they can
> witness the Complete event to be able to make the request as complete.
> -- As part of above, we store the response body in memory instead of
> just handing off the data pipe. Data pipes have a 500k data limit, which
> search requests typically go over. Net won't send the complete event
> until all data is in the pipe. Therefore we store the data and create a
> pipe to navigation if needed.
> -- Full body implementation now reports kComplete where it used to
> report kCanBeServed.
>
> Bug: 1156325
> Change-Id: Ied70a3eafd896801e263f631cb4d80556e762180
> Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2583112
> Reviewed-by: Tarun Bansal <tbansal@chromium.org>
> Reviewed-by: Robert Ogden <robertogden@chromium.org>
> Commit-Queue: Ryan Sturm <ryansturm@chromium.org>
> Cr-Commit-Position: refs/heads/master@{#835897}

Bug: 1156325
Change-Id: Ifeeed421e9124f99a0f6aa27ca08734a1b130878
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2587263Reviewed-by: Robert Ogden <robertogden@chromium.org>
Reviewed-by: Tarun Bansal <tbansal@chromium.org>
Commit-Queue: Ryan Sturm <ryansturm@chromium.org>
Cr-Commit-Position: refs/heads/master@{#836212}

chrome/browser/prefetch/search_prefetch/base_search_prefetch_request.cc

@@ -206,17 +206,17 @@ bool BaseSearchPrefetchRequest::StartPrefetchRequest(Profile* profile) {
 }
 
 void BaseSearchPrefetchRequest::CancelPrefetch() {
-  DCHECK(current_status_ == SearchPrefetchStatus::kInFlight);
+  DCHECK(current_status_ == SearchPrefetchStatus::kInFlight ||
+         current_status_ == SearchPrefetchStatus::kCanBeServed);
   current_status_ = SearchPrefetchStatus::kRequestCancelled;
   StopPrefetch();
 }
 
 void BaseSearchPrefetchRequest::ErrorEncountered() {
   DCHECK(!report_error_callback_.is_null());
-  // A streaming response can still encounter an error after the headers, so
-  // both these states are possible.
   DCHECK(current_status_ == SearchPrefetchStatus::kInFlight ||
-         current_status_ == SearchPrefetchStatus::kCanBeServed);
+         current_status_ == SearchPrefetchStatus::kCanBeServed ||
+         current_status_ == SearchPrefetchStatus::kCanBeServedAndUserClicked);
   current_status_ = SearchPrefetchStatus::kRequestFailed;
   std::move(report_error_callback_).Run();
   StopPrefetch();
@@ -227,6 +227,18 @@ void BaseSearchPrefetchRequest::MarkPrefetchAsServable() {
   current_status_ = SearchPrefetchStatus::kCanBeServed;
 }
 
+void BaseSearchPrefetchRequest::MarkPrefetchAsComplete() {
+  DCHECK(current_status_ == SearchPrefetchStatus::kInFlight ||
+         current_status_ == SearchPrefetchStatus::kCanBeServed ||
+         current_status_ == SearchPrefetchStatus::kCanBeServedAndUserClicked);
+  current_status_ = SearchPrefetchStatus::kComplete;
+}
+
+void BaseSearchPrefetchRequest::MarkPrefetchAsClicked() {
+  DCHECK(current_status_ == SearchPrefetchStatus::kCanBeServed);
+  current_status_ = SearchPrefetchStatus::kCanBeServedAndUserClicked;
+}
+
 bool BaseSearchPrefetchRequest::CanServePrefetchRequest(
     const scoped_refptr<net::HttpResponseHeaders> headers) {
   if (!headers)

chrome/browser/prefetch/search_prefetch/base_search_prefetch_request.h

@@ -27,12 +27,22 @@ enum class SearchPrefetchStatus {
   // The request is on the network and may move to any other state.
   kInFlight = 1,
   // The request can be served to the navigation stack, but may still encounter
-  // errors and move to |kRequestFailed|.
+  // errors and move to |kRequestFailed| or it may complete and move to
+  // |kComplete|. It may also move to |kCanBeServedAndUserClicked| when the user
+  // navigates to the result in omnibox or |kRequestCancelled| if the user
+  // closes omnibox.
   kCanBeServed = 2,
+  // The request can be served to the navigation stack, and is marked as being
+  // clicked by the user. At this point, it may move to |kComplete| or
+  // |kRequestFailed|.
+  kCanBeServedAndUserClicked = 3,
+  // The request can be served to the navigation stack, and has fully streamed
+  // the response with no errors. This is a terminal state.
+  kComplete = 4,
   // The request hit an error and cannot be served. This is a terminal state.
-  kRequestFailed = 3,
+  kRequestFailed = 5,
   // The request was cancelled before completion. This is terminal state.
-  kRequestCancelled = 4,
+  kRequestCancelled = 6,
 };
 
 // A class representing a prefetch used by the Search Prefetch Service.
@@ -63,6 +73,12 @@ class BaseSearchPrefetchRequest {
   // Update the status when the request is serveable.
   void MarkPrefetchAsServable();
 
+  // Update the status when the request is complete.
+  void MarkPrefetchAsComplete();
+
+  // Update the status when the relevant search item is clicked in omnibox.
+  void MarkPrefetchAsClicked();
+
   // Whether the prefetch should be served based on |headers|.
   bool CanServePrefetchRequest(
       const scoped_refptr<net::HttpResponseHeaders> headers);

chrome/browser/prefetch/search_prefetch/full_body_search_prefetch_request.cc

@@ -67,7 +67,7 @@ void FullBodySearchPrefetchRequest::LoadDone(
     return;
   }
 
-  MarkPrefetchAsServable();
+  MarkPrefetchAsComplete();
   prefetch_response_container_ = std::make_unique<PrefetchedResponseContainer>(
       simple_loader_->ResponseInfo()->Clone(), std::move(response_body));
   simple_loader_.reset();

chrome/browser/prefetch/search_prefetch/search_prefetch_service.cc

@@ -21,6 +21,8 @@
 #include "components/content_settings/core/common/content_settings.h"
 #include "components/omnibox/browser/autocomplete_controller.h"
 #include "components/omnibox/browser/base_search_provider.h"
+#include "components/omnibox/browser/omnibox_event_global_tracker.h"
+#include "components/omnibox/browser/omnibox_log.h"
 #include "components/prefs/pref_service.h"
 #include "components/search_engines/template_url_service.h"
 #include "url/origin.h"
@@ -89,6 +91,11 @@ bool SearchPrefetchService::MaybePrefetchURL(const GURL& url) {
     if (template_url) {
       template_url_service_data_ = template_url->data();
     }
+
+    omnibox_subscription_ =
+        OmniboxEventGlobalTracker::GetInstance()->RegisterCallback(
+            base::BindRepeating(&SearchPrefetchService::OnURLOpenedFromOmnibox,
+                                base::Unretained(this)));
   }
 
   base::string16 search_terms;
@@ -140,6 +147,33 @@ bool SearchPrefetchService::MaybePrefetchURL(const GURL& url) {
   return true;
 }
 
+void SearchPrefetchService::OnURLOpenedFromOmnibox(OmniboxLog* log) {
+  if (!log)
+    return;
+  const AutocompleteMatch& match = log->result.match_at(log->selected_index);
+  const GURL& opened_url = match.destination_url;
+
+  auto* template_url_service =
+      TemplateURLServiceFactory::GetForProfile(profile_);
+  DCHECK(template_url_service);
+  auto* default_search = template_url_service->GetDefaultSearchProvider();
+  if (!default_search)
+    return;
+
+  base::string16 match_search_terms;
+
+  default_search->ExtractSearchTermsFromURL(
+      opened_url, template_url_service->search_terms_data(),
+      &match_search_terms);
+
+  if (prefetches_.find(match_search_terms) == prefetches_.end() ||
+      prefetches_[match_search_terms]->current_status() !=
+          SearchPrefetchStatus::kCanBeServed) {
+    return;
+  }
+  prefetches_[match_search_terms]->MarkPrefetchAsClicked();
+}
+
 base::Optional<SearchPrefetchStatus>
 SearchPrefetchService::GetSearchPrefetchStatusForTesting(
     base::string16 search_terms) {
@@ -193,7 +227,9 @@ SearchPrefetchService::TakePrefetchResponse(const GURL& url) {
     return nullptr;
   }
 
-  if (iter->second->current_status() != SearchPrefetchStatus::kCanBeServed) {
+  if (iter->second->current_status() != SearchPrefetchStatus::kComplete &&
+      iter->second->current_status() !=
+          SearchPrefetchStatus::kCanBeServedAndUserClicked) {
     return nullptr;
   }
 
@@ -247,7 +283,9 @@ void SearchPrefetchService::OnResultChanged(
       const auto& search_terms = kv_pair.first;
       auto& prefetch_request = kv_pair.second;
       if (prefetch_request->current_status() !=
-          SearchPrefetchStatus::kInFlight) {
+              SearchPrefetchStatus::kInFlight &&
+          prefetch_request->current_status() !=
+              SearchPrefetchStatus::kCanBeServed) {
         continue;
       }
       bool should_cancel_request = true;

chrome/browser/prefetch/search_prefetch/search_prefetch_service.h

@@ -8,6 +8,7 @@
 #include <map>
 
 #include "base/callback.h"
+#include "base/callback_list.h"
 #include "base/optional.h"
 #include "base/scoped_observation.h"
 #include "base/strings/string16.h"
@@ -19,10 +20,11 @@
 #include "components/search_engines/template_url_service_observer.h"
 #include "url/gurl.h"
 
+class AutocompleteController;
+struct OmniboxLog;
 class Profile;
 class SearchPrefetchURLLoader;
 
-class AutocompleteController;
 
 class SearchPrefetchService : public KeyedService,
                               public TemplateURLServiceObserver {
@@ -64,6 +66,11 @@ class SearchPrefetchService : public KeyedService,
   // Records the current time to prevent prefetches for a set duration.
   void ReportError();
 
+  // If the navigation URL matches with a prefetch that can be served, this
+  // function marks that prefetch as clicked to prevent deletion when omnibox
+  // closes.
+  void OnURLOpenedFromOmnibox(OmniboxLog* log);
+
   // Prefetches that are started are stored using search terms as a key. Only
   // one prefetch should be started for a given search term until the old
   // prefetch expires.
@@ -80,6 +87,10 @@ class SearchPrefetchService : public KeyedService,
   // The current state of the DSE.
   base::Optional<TemplateURLData> template_url_service_data_;
 
+  // A subscription to the omnibox log service to track when a navigation is
+  // about to happen.
+  base::CallbackListSubscription omnibox_subscription_;
+
   base::ScopedObservation<TemplateURLService, TemplateURLServiceObserver>
       observer_{this};
 

chrome/browser/prefetch/search_prefetch/streaming_search_prefetch_url_loader.cc

@@ -16,9 +16,11 @@
 #include "base/time/time.h"
 #include "chrome/browser/profiles/profile.h"
 #include "content/public/browser/storage_partition.h"
+#include "mojo/public/c/system/data_pipe.h"
 #include "net/http/http_response_headers.h"
 #include "net/http/http_util.h"
 #include "net/traffic_annotation/network_traffic_annotation.h"
+#include "services/network/public/cpp/constants.h"
 #include "services/network/public/cpp/resource_request.h"
 #include "services/network/public/cpp/shared_url_loader_factory.h"
 #include "url/gurl.h"
@@ -42,9 +44,9 @@ StreamingSearchPrefetchURLLoader::StreamingSearchPrefetchURLLoader(
       url_loader_receiver_.BindNewPipeAndPassRemote(
           base::ThreadTaskRunnerHandle::Get()),
       net::MutableNetworkTrafficAnnotationTag(network_traffic_annotation));
-  url_loader_receiver_.set_disconnect_handler(
-      base::BindOnce(&StreamingSearchPrefetchURLLoader::OnMojoDisconnect,
-                     base::Unretained(this)));
+  url_loader_receiver_.set_disconnect_handler(base::BindOnce(
+      &StreamingSearchPrefetchURLLoader::OnURLLoaderMojoDisconnect,
+      base::Unretained(this)));
 }
 
 StreamingSearchPrefetchURLLoader::~StreamingSearchPrefetchURLLoader() = default;
@@ -66,17 +68,16 @@ void StreamingSearchPrefetchURLLoader::SetUpForwardingClient(
   DCHECK(!streaming_prefetch_request_);
   // Bind to the content/ navigation code.
   DCHECK(!receiver_.is_bound());
-
-  network_url_loader_->SetPriority(resource_request.priority, -1);
+  if (network_url_loader_)
+    network_url_loader_->SetPriority(resource_request.priority, -1);
 
   // At this point, we are bound to the mojo receiver, so we can release
   // |loader|, which points to |this|.
   receiver_.Bind(std::move(receiver));
-  receiver_.set_disconnect_handler(
-      base::BindOnce(&StreamingSearchPrefetchURLLoader::OnMojoDisconnect,
-                     weak_factory_.GetWeakPtr()));
   loader.release();
-
+  receiver_.set_disconnect_handler(base::BindOnce(
+      &StreamingSearchPrefetchURLLoader::OnURLLoaderClientMojoDisconnect,
+      weak_factory_.GetWeakPtr()));
   forwarding_client_.Bind(std::move(forwarding_client));
 
   if (!resource_request.report_raw_headers) {
@@ -84,9 +85,7 @@ void StreamingSearchPrefetchURLLoader::SetUpForwardingClient(
   }
 
   forwarding_client_->OnReceiveResponse(std::move(resource_response_));
-
-  // Resume previously paused network service URLLoader.
-  url_loader_receiver_.Resume();
+  RunEventQueue();
 }
 
 void StreamingSearchPrefetchURLLoader::OnReceiveResponse(
@@ -96,6 +95,11 @@ void StreamingSearchPrefetchURLLoader::OnReceiveResponse(
 
   // Store head and pause new messages until the forwarding client is set up.
   resource_response_ = std::move(head);
+  estimated_length_ = resource_response_->content_length < 0
+                          ? 0
+                          : resource_response_->content_length;
+  if (estimated_length_ > 0)
+    body_content_.reserve(estimated_length_);
 
   if (!streaming_prefetch_request_->CanServePrefetchRequest(
           resource_response_->headers)) {
@@ -105,7 +109,6 @@ void StreamingSearchPrefetchURLLoader::OnReceiveResponse(
   }
 
   streaming_prefetch_request_->MarkPrefetchAsServable();
-  url_loader_receiver_.Pause();
 }
 
 void StreamingSearchPrefetchURLLoader::OnReceiveRedirect(
@@ -133,26 +136,153 @@ void StreamingSearchPrefetchURLLoader::OnReceiveCachedMetadata(
 
 void StreamingSearchPrefetchURLLoader::OnTransferSizeUpdated(
     int32_t transfer_size_diff) {
-  DCHECK(forwarding_client_);
-  forwarding_client_->OnTransferSizeUpdated(transfer_size_diff);
+  if (forwarding_client_) {
+    DCHECK(forwarding_client_);
+    forwarding_client_->OnTransferSizeUpdated(transfer_size_diff);
+    return;
+  }
+  estimated_length_ += transfer_size_diff;
+  if (estimated_length_ > 0)
+    body_content_.reserve(estimated_length_);
+  event_queue_.push_back(
+      base::BindOnce(&StreamingSearchPrefetchURLLoader::OnTransferSizeUpdated,
+                     base::Unretained(this), transfer_size_diff));
 }
 
 void StreamingSearchPrefetchURLLoader::OnStartLoadingResponseBody(
     mojo::ScopedDataPipeConsumerHandle body) {
-  DCHECK(forwarding_client_);
-  DCHECK(!streaming_prefetch_request_);
-  forwarding_client_->OnStartLoadingResponseBody(std::move(body));
+  if (forwarding_client_) {
+    DCHECK(forwarding_client_);
+    DCHECK(!streaming_prefetch_request_);
+    forwarding_client_->OnStartLoadingResponseBody(std::move(body));
+    return;
+  }
+
+  serving_from_data_ = true;
+
+  pipe_drainer_ =
+      std::make_unique<mojo::DataPipeDrainer>(this, std::move(body));
+
+  event_queue_.push_back(base::BindOnce(
+      &StreamingSearchPrefetchURLLoader::OnStartLoadingResponseBodyFromData,
+      base::Unretained(this)));
+}
+
+void StreamingSearchPrefetchURLLoader::OnDataAvailable(const void* data,
+                                                       size_t num_bytes) {
+  body_content_.append(std::string(static_cast<const char*>(data), num_bytes));
+  bytes_of_raw_data_to_transfer_ += num_bytes;
+
+  if (forwarding_client_)
+    PushData();
+}
+
+void StreamingSearchPrefetchURLLoader::OnDataComplete() {
+  drain_complete_ = true;
+}
+
+void StreamingSearchPrefetchURLLoader::OnStartLoadingResponseBodyFromData() {
+  mojo::ScopedDataPipeConsumerHandle consumer_handle;
+
+  MojoCreateDataPipeOptions options;
+  options.struct_size = sizeof(MojoCreateDataPipeOptions);
+  options.flags = MOJO_CREATE_DATA_PIPE_FLAG_NONE;
+  options.element_num_bytes = 1;
+  options.capacity_num_bytes = network::kDataPipeDefaultAllocationSize;
+
+  MojoResult rv =
+      mojo::CreateDataPipe(&options, &producer_handle_, &consumer_handle);
+
+  if (rv != MOJO_RESULT_OK) {
+    delete this;
+    return;
+  }
+
+  handle_watcher_ = std::make_unique<mojo::SimpleWatcher>(
+      FROM_HERE, mojo::SimpleWatcher::ArmingPolicy::MANUAL,
+      base::SequencedTaskRunnerHandle::Get());
+  handle_watcher_->Watch(
+      producer_handle_.get(), MOJO_HANDLE_SIGNAL_WRITABLE,
+      MOJO_WATCH_CONDITION_SATISFIED,
+      base::BindRepeating(&StreamingSearchPrefetchURLLoader::OnHandleReady,
+                          weak_factory_.GetWeakPtr()));
+
+  forwarding_client_->OnStartLoadingResponseBody(std::move(consumer_handle));
+
+  PushData();
+}
+
+void StreamingSearchPrefetchURLLoader::OnHandleReady(
+    MojoResult result,
+    const mojo::HandleSignalsState& state) {
+  if (result != MOJO_RESULT_OK) {
+    delete this;
+    return;
+  }
+  PushData();
+}
+
+void StreamingSearchPrefetchURLLoader::PushData() {
+  while (true) {
+    DCHECK_GE(bytes_of_raw_data_to_transfer_, write_position_);
+    uint32_t write_size =
+        static_cast<uint32_t>(bytes_of_raw_data_to_transfer_ - write_position_);
+    if (write_size == 0) {
+      if (drain_complete_)
+        Finish();
+      return;
+    }
+
+    MojoResult result =
+        producer_handle_->WriteData(body_content_.data() + write_position_,
+                                    &write_size, MOJO_WRITE_DATA_FLAG_NONE);
+    if (result == MOJO_RESULT_SHOULD_WAIT) {
+      handle_watcher_->ArmOrNotify();
+      return;
+    }
+
+    if (result != MOJO_RESULT_OK) {
+      delete this;
+      return;
+    }
+
+    // |write_position_| should only be updated when the mojo pipe has
+    // successfully been written to.
+    write_position_ += write_size;
+  }
+}
+
+void StreamingSearchPrefetchURLLoader::Finish() {
+  serving_from_data_ = false;
+  handle_watcher_.reset();
+  producer_handle_.reset();
+  if (status_) {
+    forwarding_client_->OnComplete(status_.value());
+  }
 }
 
 void StreamingSearchPrefetchURLLoader::OnComplete(
     const network::URLLoaderCompletionStatus& status) {
-  DCHECK(!streaming_prefetch_request_);
-  if (forwarding_client_) {
+  network_url_loader_.reset();
+  if (forwarding_client_ && !serving_from_data_) {
+    DCHECK(!streaming_prefetch_request_);
     forwarding_client_->OnComplete(status);
     return;
   }
 
-  NOTREACHED();
+  if (!forwarding_client_) {
+    DCHECK(streaming_prefetch_request_);
+    streaming_prefetch_request_->MarkPrefetchAsComplete();
+  }
+
+  status_ = status;
+}
+
+void StreamingSearchPrefetchURLLoader::RunEventQueue() {
+  for (auto& event : event_queue_) {
+    std::move(event).Run();
+  }
+  event_queue_.clear();
 }
 
 void StreamingSearchPrefetchURLLoader::FollowRedirect(
@@ -168,27 +298,42 @@ void StreamingSearchPrefetchURLLoader::SetPriority(
     net::RequestPriority priority,
     int32_t intra_priority_value) {
   // Pass through.
-  network_url_loader_->SetPriority(priority, intra_priority_value);
+  if (network_url_loader_)
+    network_url_loader_->SetPriority(priority, intra_priority_value);
 }
 
 void StreamingSearchPrefetchURLLoader::PauseReadingBodyFromNet() {
   // Pass through.
-  network_url_loader_->PauseReadingBodyFromNet();
+  if (network_url_loader_)
+    network_url_loader_->PauseReadingBodyFromNet();
 }
 
 void StreamingSearchPrefetchURLLoader::ResumeReadingBodyFromNet() {
   // Pass through.
-  network_url_loader_->ResumeReadingBodyFromNet();
+  if (network_url_loader_)
+    network_url_loader_->ResumeReadingBodyFromNet();
 }
 
-void StreamingSearchPrefetchURLLoader::OnMojoDisconnect() {
-  if (streaming_prefetch_request_) {
+void StreamingSearchPrefetchURLLoader::OnURLLoaderMojoDisconnect() {
+  if (!network_url_loader_) {
+    // The connection should close after complete.
+    return;
+  }
+
+  if (!forwarding_client_) {
+    DCHECK(streaming_prefetch_request_);
     streaming_prefetch_request_->ErrorEncountered();
   } else {
     delete this;
   }
 }
 
+void StreamingSearchPrefetchURLLoader::OnURLLoaderClientMojoDisconnect() {
+  DCHECK(forwarding_client_);
+  DCHECK(!streaming_prefetch_request_);
+  delete this;
+}
+
 void StreamingSearchPrefetchURLLoader::ClearOwnerPointer() {
   streaming_prefetch_request_ = nullptr;
 }

chrome/browser/prefetch/search_prefetch/streaming_search_prefetch_url_loader.h

@@ -19,15 +19,19 @@
 #include "mojo/public/cpp/bindings/receiver.h"
 #include "mojo/public/cpp/bindings/remote.h"
 #include "mojo/public/cpp/system/data_pipe.h"
+#include "mojo/public/cpp/system/data_pipe_drainer.h"
 #include "services/network/public/mojom/url_loader.mojom-forward.h"
 #include "services/network/public/mojom/url_response_head.mojom-forward.h"
 
 // This class starts a search prefetch and is able to serve it once headers are
 // received. This allows streaming the response from memory as the response
-// finishes from the network.
+// finishes from the network. The class drains the network request URL Loader,
+// and creates a data pipe to handoff, so it may close the network URL Loader
+// after the read from the network is done.
 class StreamingSearchPrefetchURLLoader : public network::mojom::URLLoader,
                                          public network::mojom::URLLoaderClient,
-                                         public SearchPrefetchURLLoader {
+                                         public SearchPrefetchURLLoader,
+                                         public mojo::DataPipeDrainer::Client {
  public:
   // Creates a network service URLLoader, binds to the URL Loader, and starts
   // the request.
@@ -45,6 +49,10 @@ class StreamingSearchPrefetchURLLoader : public network::mojom::URLLoader,
   void ClearOwnerPointer();
 
  private:
+  // mojo::DataPipeDrainer::Client:
+  void OnDataAvailable(const void* data, size_t num_bytes) override;
+  void OnDataComplete() override;
+
   // SearchPrefetchURLLoader:
   SearchPrefetchURLLoader::RequestHandler ServingResponseHandler(
       std::unique_ptr<SearchPrefetchURLLoader> loader) override;
@@ -73,9 +81,28 @@ class StreamingSearchPrefetchURLLoader : public network::mojom::URLLoader,
       mojo::ScopedDataPipeConsumerHandle body) override;
   void OnComplete(const network::URLLoaderCompletionStatus& status) override;
 
-  // When a disconnection occurs in either mojo pipe, this object's lifetime
-  // needs to be managed and the connections need to be closed.
-  void OnMojoDisconnect();
+  // When a disconnection occurs in the network URLLoader mojo pipe, this
+  // object's lifetime needs to be managed and the connections need to be closed
+  // unless complete has happened.
+  void OnURLLoaderMojoDisconnect();
+
+  // When a disconnection occurs in the navigation client mojo pipe, this
+  // object's lifetime needs to be managed and the connections need to be
+  // closed.
+  void OnURLLoaderClientMojoDisconnect();
+
+  // Start serving the response from |producer_handle_|, which serves
+  // |body_content_|.
+  void OnStartLoadingResponseBodyFromData();
+
+  // Called when more data can be sent into |producer_handle_|.
+  void OnHandleReady(MojoResult result, const mojo::HandleSignalsState& state);
+
+  // Push data into |producer_handle_|.
+  void PushData();
+
+  // Clears |producer_handle_| and |handle_watcher_|.
+  void Finish();
 
   // Sets up mojo forwarding to the navigation path. Resumes
   // |network_url_loader_| calls. Serves the start of the response to the
@@ -88,6 +115,9 @@ class StreamingSearchPrefetchURLLoader : public network::mojom::URLLoader,
       mojo::PendingReceiver<network::mojom::URLLoader> receiver,
       mojo::PendingRemote<network::mojom::URLLoaderClient> forwarding_client);
 
+  // Forwards all queued events to |forwarding_client_|.
+  void RunEventQueue();
+
   // The network URLLoader that fetches the prefetch URL and its receiver.
   mojo::Remote<network::mojom::URLLoader> network_url_loader_;
   mojo::Receiver<network::mojom::URLLoaderClient> url_loader_receiver_{this};
@@ -104,10 +134,35 @@ class StreamingSearchPrefetchURLLoader : public network::mojom::URLLoader,
   // the navigation stack.
   StreamingSearchPrefetchRequest* streaming_prefetch_request_;
 
+  // Whether we are serving from |bdoy_content_|.
+  bool serving_from_data_ = false;
+
+  // The status returned from |network_url_loader_|.
+  base::Optional<network::URLLoaderCompletionStatus> status_;
+
+  // Total amount of bytes to transfer.
+  int bytes_of_raw_data_to_transfer_ = 0;
+  // Bytes sent to |producer_handle_| already.
+  int write_position_ = 0;
+  // The request body.
+  std::string body_content_;
+  int estimated_length_ = 0;
+  // Whether the body has fully been drained from |network_url_loader_|.
+  bool drain_complete_ = false;
+  // Drainer for the content in |network_url_loader_|.
+  std::unique_ptr<mojo::DataPipeDrainer> pipe_drainer_;
+
+  // URL Loader Events that occur before serving to the navigation stack should
+  // be queued internally until the request is being served.
+  std::vector<base::OnceClosure> event_queue_;
+
   // Forwarding client receiver.
   mojo::Receiver<network::mojom::URLLoader> receiver_{this};
   mojo::Remote<network::mojom::URLLoaderClient> forwarding_client_;
 
+  mojo::ScopedDataPipeProducerHandle producer_handle_;
+  std::unique_ptr<mojo::SimpleWatcher> handle_watcher_;
+
   base::WeakPtrFactory<StreamingSearchPrefetchURLLoader> weak_factory_{this};
 };