Repository: algesten/ureq
Pull request #87 (2020-06-24 03:20:03)

Remove some panics in deadline checking.
time_until_deadline had a time of check to time of use problem - the
deadline could pass between a call to checked_duration_since and the
evaluation of `deadline - now` (which panics if the result would be
negative). Resolve that by flipping the order of
checked_duration_since's arguments and using the result rather than
ignoring it.

Also there were three places that called deadline - now(), which could
panic. Replace those with time_until_deadline().


Patch:
diff --git a/Cargo.toml b/Cargo.toml
index ed99436..4071b67 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -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/pool.rs b/src/pool.rs
index f5db413..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,7 +36,44 @@ 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)]
@@ -35,13 +82,26 @@ impl ConnectionPool {
     }
 }
 
-#[derive(Debug, PartialEq, Clone, Eq, Hash)]
+#[derive(PartialEq, Clone, Eq, Hash)]
 struct PoolKey {
     scheme: String,
     hostname: String,
     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 = url.port_or_known_default();
@@ -59,6 +119,52 @@ fn poolkey_new() {
     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).
 ///
@@ -91,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..3716bb3 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
@@ -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()?


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(())
+}