connection pooling

src/unit.rs

@@ -1,9 +1,11 @@
 use body::{send_body, Payload, SizedReader};
-use std::io::Write;
-use stream::{connect_http, connect_https, connect_test};
+use std::io::{Result as IoResult, Write};
+use stream::{connect_http, connect_https, connect_test, Stream};
 use url::Url;
 //
 
+/// It's a "unit of work". Maybe a bad name for it?
+#[derive(Debug)]
 pub struct Unit {
     pub agent: Arc<Mutex<Option<AgentState>>>,
     pub url: Url,
@@ -97,64 +99,35 @@ impl Unit {
 }
 
 pub fn connect(
-    unit: Unit,
-    url: Url,
+    mut unit: Unit,
     method: &str,
+    use_pooled: bool,
     redirects: u32,
     body: SizedReader,
 ) -> Result<Response, Error> {
     //
 
     // open socket
-    let mut stream = match url.scheme() {
-        "http" => connect_http(&unit),
-        "https" => connect_https(&unit),
-        "test" => connect_test(&unit),
-        _ => Err(Error::UnknownScheme(url.scheme().to_string())),
-    }?;
+    let (mut stream, is_recycled) = connect_socket(&unit, use_pooled)?;
 
-    // send the request start + headers
-    let mut prelude: Vec<u8> = vec![];
-    write!(
-        prelude,
-        "{} {}{} HTTP/1.1\r\n",
-        method,
-        url.path(),
-        &unit.query_string
-    )?;
-    if !has_header(&unit.headers, "host") {
-        write!(prelude, "Host: {}\r\n", url.host().unwrap())?;
-    }
-    for header in &unit.headers {
-        write!(prelude, "{}: {}\r\n", header.name(), header.value())?;
-    }
-    write!(prelude, "\r\n")?;
+    let send_result = send_prelude(&unit, method, &mut stream);
 
-    stream.write_all(&mut prelude[..])?;
+    if send_result.is_err() {
+        if is_recycled {
+            // we try open a new connection, this time there will be
+            // no connection in the pool. don't use it.
+            return connect(unit, method, false, redirects, body);
+        } else {
+            // not a pooled connection, propagate the error.
+            return Err(send_result.unwrap_err().into());
+        }
+    }
 
     // start reading the response to process cookies and redirects.
     let mut resp = Response::from_read(&mut stream);
 
     // squirrel away cookies
-    {
-        let state = &mut unit.agent.lock().unwrap();
-        if let Some(add_jar) = state.as_mut().map(|state| &mut state.jar) {
-            for raw_cookie in resp.all("set-cookie").iter() {
-                let to_parse = if raw_cookie.to_lowercase().contains("domain=") {
-                    raw_cookie.to_string()
-                } else {
-                    format!("{}; Domain={}", raw_cookie, &unit.hostname)
-                };
-                match Cookie::parse_encoded(&to_parse[..]) {
-                    Err(_) => (), // ignore unparseable cookies
-                    Ok(mut cookie) => {
-                        let cookie = cookie.into_owned();
-                        add_jar.add(cookie)
-                    }
-                }
-            }
-        }
-    }
+    save_cookies(&unit, &resp);
 
     // handle redirects
     if resp.redirect() {
@@ -166,18 +139,22 @@ pub fn connect(
         let location = resp.header("location");
         if let Some(location) = location {
             // join location header to current url in case it it relative
-            let new_url = url
+            let new_url = unit
+                .url
                 .join(location)
                 .map_err(|_| Error::BadUrl(format!("Bad redirection: {}", location)))?;
 
+            // change this for every redirect since it is used when connection pooling.
+            unit.url = new_url;
+
             // perform the redirect differently depending on 3xx code.
             return match resp.status() {
                 301 | 302 | 303 => {
                     send_body(body, unit.is_chunked, &mut stream)?;
                     let empty = Payload::Empty.into_read();
-                    connect(unit, new_url, "GET", redirects - 1, empty)
+                    connect(unit, "GET", use_pooled, redirects - 1, empty)
                 }
-                307 | 308 | _ => connect(unit, new_url, method, redirects - 1, body),
+                307 | 308 | _ => connect(unit, method, use_pooled, redirects - 1, body),
             };
         }
     }
@@ -233,3 +210,71 @@ fn combine_query(url: &Url, query: &QString) -> String {
         (None, false) => "".to_string(),
     }
 }
+
+fn connect_socket(unit: &Unit, use_pooled: bool) -> Result<(Stream, bool), Error> {
+    if use_pooled {
+        let state = &mut unit.agent.lock().unwrap();
+        if let Some(agent) = state.as_mut() {
+            if let Some(stream) = agent.pool.try_get_connection(&unit.url) {
+                return Ok((stream, true));
+            }
+        }
+    }
+    let stream = match unit.url.scheme() {
+        "http" => connect_http(&unit),
+        "https" => connect_https(&unit),
+        "test" => connect_test(&unit),
+        _ => Err(Error::UnknownScheme(unit.url.scheme().to_string())),
+    };
+    Ok((stream?, false))
+}
+
+fn send_prelude(unit: &Unit, method: &str, stream: &mut Stream) -> IoResult<()> {
+    // send the request start + headers
+    let mut prelude: Vec<u8> = vec![];
+    write!(
+        prelude,
+        "{} {}{} HTTP/1.1\r\n",
+        method,
+        unit.url.path(),
+        &unit.query_string
+    )?;
+    if !has_header(&unit.headers, "host") {
+        write!(prelude, "Host: {}\r\n", unit.url.host().unwrap())?;
+    }
+    for header in &unit.headers {
+        write!(prelude, "{}: {}\r\n", header.name(), header.value())?;
+    }
+    write!(prelude, "\r\n")?;
+
+    stream.write_all(&mut prelude[..])?;
+
+    Ok(())
+}
+
+fn save_cookies(unit: &Unit, resp: &Response) {
+    //
+
+    let cookies = resp.all("set-cookie");
+    if cookies.is_empty() {
+        return;
+    }
+
+    let state = &mut unit.agent.lock().unwrap();
+    if let Some(add_jar) = state.as_mut().map(|state| &mut state.jar) {
+        for raw_cookie in cookies.iter() {
+            let to_parse = if raw_cookie.to_lowercase().contains("domain=") {
+                raw_cookie.to_string()
+            } else {
+                format!("{}; Domain={}", raw_cookie, &unit.hostname)
+            };
+            match Cookie::parse_encoded(&to_parse[..]) {
+                Err(_) => (), // ignore unparseable cookies
+                Ok(mut cookie) => {
+                    let cookie = cookie.into_owned();
+                    add_jar.add(cookie)
+                }
+            }
+        }
+    }
+}