Repository: algesten/ureq
Pull request #83 (2020-06-22 16:59:24)

1.3.0
To be merged after #81.

lolgesten: @jsha let's semantic version it and call it 1.3.0. it's new features! :D 

Patch:
diff --git a/Cargo.toml b/Cargo.toml
index ed99436..21729d5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ureq"
-version = "1.2.0"
+version = "1.3.0"
 authors = ["Martin Algesten <martin@algesten.se>"]
 description = "Minimal HTTP request library"
 license = "MIT/Apache-2.0"
@@ -41,3 +41,9 @@ native-tls = { version = "0.2", optional = true }
 
 [dev-dependencies]
 serde = { version = "1", features = ["derive"] }
+rayon = "1.3.0"
+rayon-core = "1.7.0"
+chrono = "0.4.11"
+
+[[example]]
+name = "smoke-test"
diff --git a/src/agent.rs b/src/agent.rs
index ad8266e..4eaf6d1 100644
--- a/src/agent.rs
+++ b/src/agent.rs
@@ -262,6 +262,7 @@ pub(crate) fn basic_auth(user: &str, pass: &str) -> String {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::io::Read;
 
     ///////////////////// AGENT TESTS //////////////////////////////
 
@@ -273,6 +274,35 @@ mod tests {
         });
     }
 
+    #[test]
+    #[cfg(any(feature = "tls", feature = "native-tls"))]
+    fn agent_pool() {
+        let agent = crate::agent();
+        let url = "https://ureq.s3.eu-central-1.amazonaws.com/sherlock.txt";
+        // req 1
+        let resp = agent.get(url).call();
+        assert!(resp.ok());
+        let mut reader = resp.into_reader();
+        let mut buf = vec![];
+        // reading the entire content will return the connection to the pool
+        reader.read_to_end(&mut buf).unwrap();
+
+        fn poolsize(agent: &Agent) -> usize {
+            let mut lock = agent.state.lock().unwrap();
+            let state = lock.as_mut().unwrap();
+            state.pool().len()
+        }
+        assert_eq!(poolsize(&agent), 1);
+
+        // req 2 should be done with a reused connection
+        let resp = agent.get(url).call();
+        assert!(resp.ok());
+        assert_eq!(poolsize(&agent), 0);
+        let mut reader = resp.into_reader();
+        let mut buf = vec![];
+        reader.read_to_end(&mut buf).unwrap();
+    }
+
     //////////////////// REQUEST TESTS /////////////////////////////
 
     #[test]
diff --git a/src/pool.rs b/src/pool.rs
index 0fe8a56..07e28c4 100644
--- a/src/pool.rs
+++ b/src/pool.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, VecDeque};
 use std::io::{Read, Result as IoResult};
 
 use crate::stream::Stream;
@@ -7,14 +7,24 @@ use crate::unit::Unit;
 use url::Url;
 
 pub const DEFAULT_HOST: &str = "localhost";
+const MAX_IDLE_CONNECTIONS: usize = 100;
 
 /// Holder of recycled connections.
 ///
+/// Invariant: The length of recycle and lru are the same.
+/// Invariant: Each PoolKey exists as a key in recycle, and vice versa.
+/// Invariant: Each PoolKey exists in recycle at most once and lru at most once.
+///
 /// *Internal API*
 #[derive(Default, Debug)]
 pub(crate) struct ConnectionPool {
     // the actual pooled connection. however only one per hostname:port.
     recycle: HashMap<PoolKey, Stream>,
+    // This is used to keep track of which streams to expire when the
+    // pool reaches MAX_IDLE_CONNECTIONS. The corresponding PoolKeys for
+    // recently used Streams are added to the back of the queue;
+    // old streams are removed from the front.
+    lru: VecDeque<PoolKey>,
 }
 
 impl ConnectionPool {
@@ -26,50 +36,135 @@ impl ConnectionPool {
 
     /// How the unit::connect tries to get a pooled connection.
     pub fn try_get_connection(&mut self, url: &Url) -> Option<Stream> {
-        self.recycle.remove(&PoolKey::new(url))
+        let key = PoolKey::new(url);
+        self.remove(&key)
+    }
+
+    fn remove(&mut self, key: &PoolKey) -> Option<Stream> {
+        if !self.recycle.contains_key(&key) {
+            return None;
+        }
+        let index = self.lru.iter().position(|k| k == key);
+        assert!(
+            index.is_some(),
+            "invariant failed: key existed in recycle but not lru"
+        );
+        self.lru.remove(index.unwrap());
+        self.recycle.remove(&key)
+    }
+
+    fn add(&mut self, key: PoolKey, stream: Stream) {
+        // If an entry with the same key already exists, remove it.
+        // The more recently used stream is likely to live longer.
+        self.remove(&key);
+        if self.recycle.len() + 1 > MAX_IDLE_CONNECTIONS {
+            self.remove_oldest();
+        }
+        self.lru.push_back(key.clone());
+        self.recycle.insert(key, stream);
+    }
+
+    fn remove_oldest(&mut self) {
+        if let Some(key) = self.lru.pop_front() {
+            let removed = self.recycle.remove(&key);
+            assert!(
+                removed.is_some(),
+                "invariant failed: key existed in lru but not in recycle"
+            );
+        } else {
+            panic!("tried to remove oldest but no entries found!");
+        }
     }
 
     #[cfg(test)]
     pub fn len(&self) -> usize {
         self.recycle.len()
     }
-
-    #[cfg(all(test, any(feature = "tls", feature = "native-tls")))]
-    pub fn get(&self, hostname: &str, port: u16) -> Option<&Stream> {
-        let key = PoolKey {
-            hostname: hostname.into(),
-            port,
-        };
-        self.recycle.get(&key)
-    }
 }
 
-#[derive(Debug, PartialEq, Clone, Eq, Hash)]
+#[derive(PartialEq, Clone, Eq, Hash)]
 struct PoolKey {
+    scheme: String,
     hostname: String,
-    port: u16,
+    port: Option<u16>,
+}
+
+use std::fmt;
+
+impl fmt::Debug for PoolKey {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_fmt(format_args!(
+            "{}|{}|{}",
+            self.scheme,
+            self.hostname,
+            self.port.unwrap_or(0)
+        ))
+    }
 }
 
 impl PoolKey {
     fn new(url: &Url) -> Self {
-        let port = if cfg!(test) {
-            if let Some(p) = url.port_or_known_default() {
-                Some(p)
-            } else if url.scheme() == "test" {
-                Some(42)
-            } else {
-                None
-            }
-        } else {
-            url.port_or_known_default()
-        };
+        let port = url.port_or_known_default();
         PoolKey {
-            hostname: url.host_str().unwrap_or(DEFAULT_HOST).into(),
-            port: port.expect("Failed to get port for pool key"),
+            scheme: url.scheme().to_string(),
+            hostname: url.host_str().unwrap_or("").to_string(),
+            port,
         }
     }
 }
 
+#[test]
+fn poolkey_new() {
+    // Test that PoolKey::new() does not panic on unrecognized schemes.
+    PoolKey::new(&Url::parse("zzz:///example.com").unwrap());
+}
+
+#[test]
+fn pool_size_limit() {
+    assert_eq!(MAX_IDLE_CONNECTIONS, 100);
+    let mut pool = ConnectionPool::new();
+    let hostnames = (0..200).map(|i| format!("{}.example", i));
+    let poolkeys = hostnames.map(|hostname| PoolKey {
+        scheme: "https".to_string(),
+        hostname,
+        port: Some(999),
+    });
+    for key in poolkeys.clone() {
+        pool.add(key, Stream::Cursor(std::io::Cursor::new(vec![])));
+    }
+    assert_eq!(pool.len(), 100);
+
+    for key in poolkeys.skip(100) {
+        let result = pool.remove(&key);
+        assert!(result.is_some(), "expected key was not in pool");
+    }
+}
+
+#[test]
+fn pool_duplicates_limit() {
+    // Test inserting duplicates into the pool, and subsequently
+    // filling and draining it. The duplicates should evict earlier
+    // entries with the same key.
+    assert_eq!(MAX_IDLE_CONNECTIONS, 100);
+    let mut pool = ConnectionPool::new();
+    let hostnames = (0..100).map(|i| format!("{}.example", i));
+    let poolkeys = hostnames.map(|hostname| PoolKey {
+        scheme: "https".to_string(),
+        hostname,
+        port: Some(999),
+    });
+    for key in poolkeys.clone() {
+        pool.add(key.clone(), Stream::Cursor(std::io::Cursor::new(vec![])));
+        pool.add(key, Stream::Cursor(std::io::Cursor::new(vec![])));
+    }
+    assert_eq!(pool.len(), 100);
+
+    for key in poolkeys {
+        let result = pool.remove(&key);
+        assert!(result.is_some(), "expected key was not in pool");
+    }
+}
+
 /// Read wrapper that returns the stream to the pool once the
 /// read is exhausted (reached a 0).
 ///
@@ -102,7 +197,7 @@ impl<R: Read + Sized + Into<Stream>> PoolReturnRead<R> {
                 }
                 // insert back into pool
                 let key = PoolKey::new(&unit.url);
-                agent.pool().recycle.insert(key, stream);
+                agent.pool().add(key, stream);
             }
         }
     }
diff --git a/src/stream.rs b/src/stream.rs
index f7b4088..4cc0d71 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -81,12 +81,12 @@ impl Read for DeadlineStream {
 // then. Otherwise return a TimedOut error.
 fn time_until_deadline(deadline: Instant) -> IoResult<Duration> {
     let now = Instant::now();
-    match now.checked_duration_since(deadline) {
-        Some(_) => Err(IoError::new(
+    match deadline.checked_duration_since(now) {
+        None => Err(IoError::new(
             ErrorKind::TimedOut,
             "timed out reading response",
         )),
-        None => Ok(deadline - now),
+        Some(duration) => Ok(duration),
     }
 }
 
@@ -394,7 +394,7 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
     // Setting it to None will disable the native system timeout
     if let Some(deadline) = deadline {
         stream
-            .set_read_timeout(Some(deadline - Instant::now()))
+            .set_read_timeout(Some(time_until_deadline(deadline)?))
             .ok();
     } else if unit.timeout_read > 0 {
         stream
@@ -406,7 +406,7 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
 
     if let Some(deadline) = deadline {
         stream
-            .set_write_timeout(Some(deadline - Instant::now()))
+            .set_write_timeout(Some(time_until_deadline(deadline)?))
             .ok();
     } else if unit.timeout_write > 0 {
         stream
@@ -467,7 +467,7 @@ fn socks5_local_nslookup(hostname: &str, port: u16) -> Result<TargetAddr, std::i
 #[cfg(feature = "socks-proxy")]
 fn connect_socks5(
     proxy: Proxy,
-    deadline: Option<time::Instant>,
+    deadline: Option<Instant>,
     proxy_addr: SocketAddr,
     host: &str,
     port: u16,
@@ -523,7 +523,9 @@ fn connect_socks5(
         let (lock, cvar) = &*master_signal;
         let done = lock.lock().unwrap();
 
-        let done_result = cvar.wait_timeout(done, deadline - Instant::now()).unwrap();
+        let done_result = cvar
+            .wait_timeout(done, time_until_deadline(deadline)?)
+            .unwrap();
         let done = done_result.0;
         if *done {
             rx.recv().unwrap()?
diff --git a/src/unit.rs b/src/unit.rs
index cea0e05..b6fd18b 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -291,6 +291,10 @@ pub(crate) fn combine_query(url: &Url, query: &QString, mix_queries: bool) -> St
 
 /// Connect the socket, either by using the pool or grab a new one.
 fn connect_socket(unit: &Unit, use_pooled: bool) -> Result<(Stream, bool), Error> {
+    match unit.url.scheme() {
+        "http" | "https" | "test" => (),
+        _ => return Err(Error::UnknownScheme(unit.url.scheme().to_string())),
+    };
     if use_pooled {
         let state = &mut unit.agent.lock().unwrap();
         if let Some(agent) = state.as_mut() {


Test patch:
diff --git a/examples/smoke-test/main.rs b/examples/smoke-test/main.rs
new file mode 100644
index 0000000..39770a7
--- /dev/null
+++ b/examples/smoke-test/main.rs
@@ -0,0 +1,107 @@
+use chrono::Local;
+use rayon::prelude::*;
+use rayon_core;
+
+use std::io::{self, BufRead, BufReader, Read};
+use std::iter::Iterator;
+use std::time::Duration;
+use std::{env, error, fmt, result};
+
+use ureq;
+
+#[derive(Debug)]
+struct Oops(String);
+
+impl From<io::Error> for Oops {
+    fn from(e: io::Error) -> Oops {
+        Oops(e.to_string())
+    }
+}
+
+impl From<&ureq::Error> for Oops {
+    fn from(e: &ureq::Error) -> Oops {
+        Oops(e.to_string())
+    }
+}
+
+impl From<rayon_core::ThreadPoolBuildError> for Oops {
+    fn from(e: rayon_core::ThreadPoolBuildError) -> Oops {
+        Oops(e.to_string())
+    }
+}
+
+impl error::Error for Oops {}
+
+impl fmt::Display for Oops {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+type Result<T> = result::Result<T, Oops>;
+
+fn get(agent: &ureq::Agent, url: &String) -> Result<Vec<u8>> {
+    let response = agent
+        .get(url)
+        .timeout_connect(5_000)
+        .timeout(Duration::from_secs(20))
+        .call();
+    if let Some(err) = response.synthetic_error() {
+        return Err(err.into());
+    }
+    let mut reader = response.into_reader();
+    let mut bytes = vec![];
+    reader.read_to_end(&mut bytes)?;
+    Ok(bytes)
+}
+
+fn get_and_write(agent: &ureq::Agent, url: &String) -> Result<()> {
+    println!("🕷️ {} {}", Local::now(), url);
+    match get(agent, url) {
+        Ok(_) => println!("✔️ {} {}", Local::now(), url),
+        Err(e) => println!("⚠️ {} {} {}", Local::now(), url, e),
+    }
+    Ok(())
+}
+
+fn get_many(urls: Vec<String>, simultaneous_fetches: usize) -> Result<()> {
+    let agent = ureq::Agent::default().build();
+    let pool = rayon::ThreadPoolBuilder::new()
+        .num_threads(simultaneous_fetches)
+        .build()?;
+    pool.scope(|_| {
+        urls.par_iter().map(|u| get_and_write(&agent, u)).count();
+    });
+    Ok(())
+}
+
+fn main() -> Result<()> {
+    let args = env::args();
+    if args.len() == 1 {
+        println!(
+            r##"Usage: {:#?} top-1m.csv
+        
+Where top-1m.csv is a simple, unquoted CSV containing two fields, a rank and
+a domain name. For instance you can get such a list from https://tranco-list.eu/.
+
+For each domain, this program will attempt to GET four URLs: The domain name
+name with HTTP and HTTPS, and with and without a www prefix. It will fetch
+using 50 threads concurrently.
+"##,
+            env::current_exe()?
+        );
+        return Ok(());
+    }
+    let file = std::fs::File::open(args.skip(1).next().unwrap())?;
+    let bufreader = BufReader::new(file);
+    let mut urls = vec![];
+    for line in bufreader.lines() {
+        let domain = line?.rsplit(",").next().unwrap().to_string();
+        urls.push(format!("http://{}/", domain));
+        urls.push(format!("https://{}/", domain));
+        urls.push(format!("http://www.{}/", domain));
+        urls.push(format!("https://www.{}/", domain));
+    }
+    get_many(urls, 50)?;
+    Ok(())
+}
diff --git a/src/test/range.rs b/src/test/range.rs
index 9c4adbc..cdef429 100644
--- a/src/test/range.rs
+++ b/src/test/range.rs
@@ -20,41 +20,3 @@ fn read_range() {
         [83, 99, 111, 116, 116, 34, 10, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32]
     )
 }
-
-#[test]
-#[cfg(any(feature = "tls", feature = "native-tls"))]
-fn agent_pool() {
-    let agent = agent();
-
-    // req 1
-    let resp = agent
-        .get("https://ureq.s3.eu-central-1.amazonaws.com/sherlock.txt")
-        .set("Range", "bytes=1000-1999")
-        .call();
-    assert_eq!(resp.status(), 206);
-    let mut reader = resp.into_reader();
-    let mut buf = vec![];
-    // reading the entire content will return the connection to the pool
-    let len = reader.read_to_end(&mut buf).unwrap();
-    assert_eq!(len, 1000);
-
-    {
-        let mut lock = agent.state.lock().unwrap();
-        let state = lock.as_mut().unwrap();
-        let pool = state.pool();
-        assert_eq!(pool.len(), 1);
-        let f = format!("{:?}", pool.get("ureq.s3.eu-central-1.amazonaws.com", 443));
-        assert_eq!(f, "Some(Stream[https])"); // not a great way of testing.
-    }
-
-    // req 2 should be done with a reused connection
-    let resp = agent
-        .get("https://ureq.s3.eu-central-1.amazonaws.com/sherlock.txt")
-        .set("Range", "bytes=5000-6999")
-        .call();
-    assert_eq!(resp.status(), 206);
-    let mut reader = resp.into_reader();
-    let mut buf = vec![];
-    let len = reader.read_to_end(&mut buf).unwrap();
-    assert_eq!(len, 2000);
-}