Ensure done_chan is Some, when refreshing a still receiving resource

3 files changed, 88 insertions, 19 deletions

diff --git a/components/net/http_cache.rs b/components/net/http_cache.rs
index e45c8577215..601a4f6cc65 100644
--- a/components/net/http_cache.rs
+++ b/components/net/http_cache.rs
@@ -624,6 +624,23 @@ impl HttpCache {
         let entry_key = CacheKey::new(request.clone());
         if let Some(cached_resources) = self.entries.get_mut(&entry_key) {
             for cached_resource in cached_resources.iter_mut() {
+                // done_chan will have been set to Some(..) by http_network_fetch.
+                // If the body is not receiving data, set the done_chan back to None.
+                // Otherwise, create a new dedicated channel to update the consumer.
+                // The response constructed here will replace the 304 one from the network.
+                let in_progress_channel = match *cached_resource.body.lock().unwrap() {
+                    ResponseBody::Receiving(..) => {
+                        Some(channel())
+                    },
+                    ResponseBody::Empty | ResponseBody::Done(..) => None
+                };
+                match in_progress_channel {
+                    Some((done_sender, done_receiver)) => {
+                        *done_chan = Some((done_sender.clone(), done_receiver));
+                        cached_resource.awaiting_body.lock().unwrap().push(done_sender);
+                    },
+                    None => *done_chan = None
+                }
                 // Received a response with 304 status code, in response to a request that matches a cached resource.
                 // 1. update the headers of the cached resource.
                 // 2. return a response, constructed from the cached resource.
@@ -635,9 +652,6 @@ impl HttpCache {
                 constructed_response.referrer_policy = request.referrer_policy.clone();
                 constructed_response.raw_status = cached_resource.data.raw_status.clone();
                 constructed_response.url_list = cached_resource.data.url_list.clone();
-                // done_chan will have been set to Some by http_network_fetch,
-                // set it back to None since the response returned here replaces the 304 one from the network.
-                *done_chan = None;
                 cached_resource.data.expires = get_response_expiry(&constructed_response);
                 let mut stored_headers = cached_resource.data.metadata.headers.lock().unwrap();
                 stored_headers.extend(response.headers.iter());
diff --git a/components/net/http_loader.rs b/components/net/http_loader.rs
index a8efbc4d297..3f4eb381c93 100644
--- a/components/net/http_loader.rs
+++ b/components/net/http_loader.rs
@@ -905,26 +905,30 @@ fn http_network_or_cache_fetch(request: &mut Request,
         }
     }
 
-    if let Some(ref ch) = *done_chan {
-        // The cache constructed a response with a body of ResponseBody::Receiving.
-        // We wait for the response in the cache to "finish",
-        // with a body of either Done or Cancelled.
-        loop {
-            match ch.1.recv()
-                    .expect("HTTP cache should always send Done or Cancelled") {
-                Data::Payload(_) => {},
-                Data::Done => break, // Return the full response as if it was initially cached as such.
-                Data::Cancelled => {
-                    // The response was cancelled while the fetch was ongoing.
-                    // Set response to None, which will trigger a network fetch below.
-                    response = None;
-                    break;
+    fn wait_for_cached_response(done_chan: &mut DoneChannel, response: &mut Option<Response>) {
+        if let Some(ref ch) = *done_chan {
+            // The cache constructed a response with a body of ResponseBody::Receiving.
+            // We wait for the response in the cache to "finish",
+            // with a body of either Done or Cancelled.
+            loop {
+                match ch.1.recv()
+                        .expect("HTTP cache should always send Done or Cancelled") {
+                    Data::Payload(_) => {},
+                    Data::Done => break, // Return the full response as if it was initially cached as such.
+                    Data::Cancelled => {
+                        // The response was cancelled while the fetch was ongoing.
+                        // Set response to None, which will trigger a network fetch below.
+                        *response = None;
+                        break;
+                    }
                 }
             }
         }
+        // Set done_chan back to None, it's cache-related usefulness ends here.
+        *done_chan = None;
     }
-    // Set done_chan back to None, it's cache-related usefulness ends here.
-    *done_chan = None;
+
+    wait_for_cached_response(done_chan, &mut response);
 
     // Step 22
     if response.is_none() {
@@ -951,6 +955,7 @@ fn http_network_or_cache_fetch(request: &mut Request,
         if revalidating_flag && forward_response.status.map_or(false, |s| s == StatusCode::NotModified) {
             if let Ok(mut http_cache) = context.state.http_cache.write() {
                 response = http_cache.refresh(&http_request, forward_response.clone(), done_chan);
+                wait_for_cached_response(done_chan, &mut response);
             }
         }
 
diff --git a/components/net/tests/http_cache.rs b/components/net/tests/http_cache.rs
new file mode 100644
index 00000000000..c63ba8131c7
--- /dev/null
+++ b/components/net/tests/http_cache.rs
@@ -0,0 +1,50 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+use hyper::header::{Expires, HttpDate};
+use hyper::method::Method;
+use hyper::status::StatusCode;
+use msg::constellation_msg::TEST_PIPELINE_ID;
+use net::http_cache::HttpCache;
+use net_traits::request::{Destination, Request, RequestInit};
+use net_traits::response::{Response, ResponseBody};
+use servo_url::ServoUrl;
+use std::sync::mpsc::channel;
+use time;
+
+
+#[test]
+fn test_refreshing_resource_sets_done_chan_the_appropriate_value() {
+    let response_bodies = vec![ResponseBody::Receiving(vec![]),
+                               ResponseBody::Empty,
+                               ResponseBody::Done(vec![])];
+    let url = ServoUrl::parse("https://servo.org").unwrap();
+    let request = Request::from_init(RequestInit {
+        url: url.clone(),
+        method: Method::Get,
+        destination: Destination::Document,
+        origin: url.clone().origin(),
+        pipeline_id: Some(TEST_PIPELINE_ID),
+        .. RequestInit::default()
+    });
+    let mut response = Response::new(url.clone());
+    // Expires header makes the response cacheable.
+    response.headers.set(Expires(HttpDate(time::now())));
+    response_bodies.iter().for_each(|body| {
+        let mut cache = HttpCache::new();
+        *response.body.lock().unwrap() = body;
+        // First, store the 'normal' response.
+        cache.store(&request, &response);
+        // Second, mutate the response into a 304 response, and refresh the stored one.
+        response.status = Some(StatusCode::NotModified);
+        let mut done_chan = Some(channel());
+        let refreshed_response = cache.refresh(&request, response, &mut done_chan);
+        // Ensure a resource was found, and refreshed.
+        assert!(refreshed_response.is_some());
+        match body {
+            ResponseBody::Receiving(_) => assert!(done_chan.is_some()),
+            ResponseBody::Empty | ResponseBody::Done(_) => assert!(done_chan.is_none())
+        }
+    })
+}