use std::collections::{HashMap, VecDeque}; use std::fmt; use std::ops::{Deref, DerefMut}; use std::sync::{Arc, Mutex, Weak}; use std::time::{Duration, Instant}; use futures::{Future, Async, Poll, Stream}; use futures::sync::oneshot; use tokio::reactor::{Handle, Interval}; pub struct Pool { inner: Arc>>, } // Before using a pooled connection, make sure the sender is not dead. // // This is a trait to allow the `client::pool::tests` to work for `i32`. // // See https://github.com/hyperium/hyper/issues/1429 pub trait Closed { fn is_closed(&self) -> bool; } struct PoolInner { enabled: bool, // These are internal Conns sitting in the event loop in the KeepAlive // state, waiting to receive a new Request to send on the socket. idle: HashMap, Vec>>, // These are outstanding Checkouts that are waiting for a socket to be // able to send a Request one. This is used when "racing" for a new // connection. // // The Client starts 2 tasks, 1 to connect a new socket, and 1 to wait // for the Pool to receive an idle Conn. When a Conn becomes idle, // this list is checked for any parked Checkouts, and tries to notify // them that the Conn could be used instead of waiting for a brand new // connection. parked: HashMap, VecDeque>>, timeout: Option, // Used to prevent multiple intervals from being spawned to clear // expired connections. // // TODO(0.12): Remove the need for this when Client::schedule_pool_timer // can be done in Client::new. expired_timer_spawned: bool, } impl Pool { pub fn new(enabled: bool, timeout: Option) -> Pool { Pool { inner: Arc::new(Mutex::new(PoolInner { enabled: enabled, idle: HashMap::new(), parked: HashMap::new(), timeout: timeout, expired_timer_spawned: false, })), } } } impl Pool { pub fn checkout(&self, key: &str) -> Checkout { Checkout { key: Arc::new(key.to_owned()), pool: self.clone(), parked: None, } } fn take(&self, key: &Arc) -> Option> { let entry = { let mut inner = self.inner.lock().unwrap(); let expiration = Expiration::new(inner.timeout); let mut should_remove = false; let entry = inner.idle.get_mut(key).and_then(|list| { trace!("take; url = {:?}, expiration = {:?}", key, expiration.0); while let Some(entry) = list.pop() { if !expiration.expires(entry.idle_at) { if !entry.value.is_closed() { should_remove = list.is_empty(); return Some(entry); } } trace!("removing unacceptable pooled {:?}", key); // every other case the Idle should just be dropped // 1. Idle but expired // 2. Busy (something else somehow took it?) // 3. Disabled don't reuse of course } should_remove = true; None }); if should_remove { inner.idle.remove(key); } entry }; entry.map(|e| self.reuse(key, e.value)) } pub fn pooled(&self, key: Arc, value: T) -> Pooled { Pooled { is_reused: false, key: key, pool: Arc::downgrade(&self.inner), value: Some(value) } } fn reuse(&self, key: &Arc, value: T) -> Pooled { debug!("reuse idle connection for {:?}", key); Pooled { is_reused: true, key: key.clone(), pool: Arc::downgrade(&self.inner), value: Some(value), } } fn park(&mut self, key: Arc, tx: oneshot::Sender) { trace!("park; waiting for idle connection: {:?}", key); self.inner.lock().unwrap() .parked.entry(key) .or_insert(VecDeque::new()) .push_back(tx); } } impl PoolInner { fn put(&mut self, key: Arc, value: T) { if !self.enabled { return; } trace!("Pool::put {:?}", key); let mut remove_parked = false; let mut value = Some(value); if let Some(parked) = self.parked.get_mut(&key) { while let Some(tx) = parked.pop_front() { if !tx.is_canceled() { match tx.send(value.take().unwrap()) { Ok(()) => break, Err(e) => { value = Some(e); } } } trace!("Pool::put removing canceled parked {:?}", key); } remove_parked = parked.is_empty(); } if remove_parked { self.parked.remove(&key); } match value { Some(value) => { debug!("pooling idle connection for {:?}", key); self.idle.entry(key) .or_insert(Vec::new()) .push(Idle { value: value, idle_at: Instant::now(), }); } None => trace!("Pool::put found parked {:?}", key), } } } impl PoolInner { /// Any `FutureResponse`s that were created will have made a `Checkout`, /// and possibly inserted into the pool that it is waiting for an idle /// connection. If a user ever dropped that future, we need to clean out /// those parked senders. fn clean_parked(&mut self, key: &Arc) { let mut remove_parked = false; if let Some(parked) = self.parked.get_mut(key) { parked.retain(|tx| { !tx.is_canceled() }); remove_parked = parked.is_empty(); } if remove_parked { self.parked.remove(key); } } } impl PoolInner { fn clear_expired(&mut self) { let dur = if let Some(dur) = self.timeout { dur } else { return }; let now = Instant::now(); //self.last_idle_check_at = now; self.idle.retain(|_key, values| { values.retain(|entry| { if entry.value.is_closed() { return false; } now - entry.idle_at < dur }); // returning false evicts this key/val !values.is_empty() }); } } impl Pool { pub(super) fn spawn_expired_interval(&self, handle: &Handle) { let dur = { let mut inner = self.inner.lock().unwrap(); if !inner.enabled { return; } if inner.expired_timer_spawned { return; } inner.expired_timer_spawned = true; if let Some(dur) = inner.timeout { dur } else { return } }; let interval = Interval::new(dur, handle) .expect("reactor is gone"); handle.spawn(IdleInterval { interval: interval, pool: Arc::downgrade(&self.inner), }); } } impl Clone for Pool { fn clone(&self) -> Pool { Pool { inner: self.inner.clone(), } } } pub struct Pooled { value: Option, is_reused: bool, key: Arc, pool: Weak>>, } impl Pooled { pub fn is_reused(&self) -> bool { self.is_reused } fn as_ref(&self) -> &T { self.value.as_ref().expect("not dropped") } fn as_mut(&mut self) -> &mut T { self.value.as_mut().expect("not dropped") } } impl Deref for Pooled { type Target = T; fn deref(&self) -> &T { self.as_ref() } } impl DerefMut for Pooled { fn deref_mut(&mut self) -> &mut T { self.as_mut() } } impl Drop for Pooled { fn drop(&mut self) { if let Some(value) = self.value.take() { if let Some(inner) = self.pool.upgrade() { if let Ok(mut inner) = inner.lock() { inner.put(self.key.clone(), value); } } else { trace!("pool dropped, dropping pooled ({:?})", self.key); } } } } impl fmt::Debug for Pooled { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_struct("Pooled") .field("key", &self.key) .finish() } } struct Idle { idle_at: Instant, value: T, } pub struct Checkout { key: Arc, pool: Pool, parked: Option>, } struct NotParked; impl Checkout { fn poll_parked(&mut self) -> Poll, NotParked> { let mut drop_parked = false; if let Some(ref mut rx) = self.parked { match rx.poll() { Ok(Async::Ready(value)) => { if !value.is_closed() { return Ok(Async::Ready(self.pool.reuse(&self.key, value))); } drop_parked = true; }, Ok(Async::NotReady) => return Ok(Async::NotReady), Err(_canceled) => drop_parked = true, } } if drop_parked { self.parked.take(); } Err(NotParked) } fn park(&mut self) { if self.parked.is_none() { let (tx, mut rx) = oneshot::channel(); let _ = rx.poll(); // park this task self.pool.park(self.key.clone(), tx); self.parked = Some(rx); } } } impl Future for Checkout { type Item = Pooled; type Error = ::Error; fn poll(&mut self) -> Poll { match self.poll_parked() { Ok(async) => return Ok(async), Err(_not_parked) => (), } let entry = self.pool.take(&self.key); if let Some(pooled) = entry { Ok(Async::Ready(pooled)) } else { self.park(); Ok(Async::NotReady) } } } impl Drop for Checkout { fn drop(&mut self) { self.parked.take(); if let Ok(mut inner) = self.pool.inner.lock() { inner.clean_parked(&self.key); } } } struct Expiration(Option); impl Expiration { fn new(dur: Option) -> Expiration { Expiration(dur) } fn expires(&self, instant: Instant) -> bool { match self.0 { Some(timeout) => instant.elapsed() > timeout, None => false, } } } struct IdleInterval { interval: Interval, pool: Weak>>, } impl Future for IdleInterval { type Item = (); type Error = (); fn poll(&mut self) -> Poll { loop { try_ready!(self.interval.poll().map_err(|_| unreachable!("interval cannot error"))); if let Some(inner) = self.pool.upgrade() { if let Ok(mut inner) = inner.lock() { inner.clear_expired(); continue; } } return Ok(Async::Ready(())); } } } #[cfg(test)] mod tests { use std::sync::Arc; use std::time::Duration; use futures::{Async, Future}; use futures::future; use super::{Closed, Pool}; impl Closed for i32 { fn is_closed(&self) -> bool { false } } #[test] fn test_pool_checkout_smoke() { let pool = Pool::new(true, Some(Duration::from_secs(5))); let key = Arc::new("foo".to_string()); let pooled = pool.pooled(key.clone(), 41); drop(pooled); match pool.checkout(&key).poll().unwrap() { Async::Ready(pooled) => assert_eq!(*pooled, 41), _ => panic!("not ready"), } } #[test] fn test_pool_checkout_returns_none_if_expired() { future::lazy(|| { let pool = Pool::new(true, Some(Duration::from_millis(100))); let key = Arc::new("foo".to_string()); let pooled = pool.pooled(key.clone(), 41); drop(pooled); ::std::thread::sleep(pool.inner.lock().unwrap().timeout.unwrap()); assert!(pool.checkout(&key).poll().unwrap().is_not_ready()); ::futures::future::ok::<(), ()>(()) }).wait().unwrap(); } #[test] fn test_pool_checkout_removes_expired() { future::lazy(|| { let pool = Pool::new(true, Some(Duration::from_millis(100))); let key = Arc::new("foo".to_string()); pool.pooled(key.clone(), 41); pool.pooled(key.clone(), 5); pool.pooled(key.clone(), 99); assert_eq!(pool.inner.lock().unwrap().idle.get(&key).map(|entries| entries.len()), Some(3)); ::std::thread::sleep(pool.inner.lock().unwrap().timeout.unwrap()); // checkout.poll() should clean out the expired pool.checkout(&key).poll().unwrap(); assert!(pool.inner.lock().unwrap().idle.get(&key).is_none()); Ok::<(), ()>(()) }).wait().unwrap(); } #[test] fn test_pool_timer_removes_expired() { let mut core = ::tokio::reactor::Core::new().unwrap(); let pool = Pool::new(true, Some(Duration::from_millis(100))); pool.spawn_expired_interval(&core.handle()); let key = Arc::new("foo".to_string()); pool.pooled(key.clone(), 41); pool.pooled(key.clone(), 5); pool.pooled(key.clone(), 99); assert_eq!(pool.inner.lock().unwrap().idle.get(&key).map(|entries| entries.len()), Some(3)); let timeout = ::tokio::reactor::Timeout::new( Duration::from_millis(400), // allow for too-good resolution &core.handle() ).unwrap(); core.run(timeout).unwrap(); assert!(pool.inner.lock().unwrap().idle.get(&key).is_none()); } #[test] fn test_pool_checkout_task_unparked() { let pool = Pool::new(true, Some(Duration::from_secs(10))); let key = Arc::new("foo".to_string()); let pooled = pool.pooled(key.clone(), 41); let checkout = pool.checkout(&key).join(future::lazy(move || { // the checkout future will park first, // and then this lazy future will be polled, which will insert // the pooled back into the pool // // this test makes sure that doing so will unpark the checkout drop(pooled); Ok(()) })).map(|(entry, _)| entry); assert_eq!(*checkout.wait().unwrap(), 41); } #[test] fn test_pool_checkout_drop_cleans_up_parked() { future::lazy(|| { let pool = Pool::::new(true, Some(Duration::from_secs(10))); let key = Arc::new("localhost:12345".to_string()); let mut checkout1 = pool.checkout(&key); let mut checkout2 = pool.checkout(&key); // first poll needed to get into Pool's parked checkout1.poll().unwrap(); assert_eq!(pool.inner.lock().unwrap().parked.get(&key).unwrap().len(), 1); checkout2.poll().unwrap(); assert_eq!(pool.inner.lock().unwrap().parked.get(&key).unwrap().len(), 2); // on drop, clean up Pool drop(checkout1); assert_eq!(pool.inner.lock().unwrap().parked.get(&key).unwrap().len(), 1); drop(checkout2); assert!(pool.inner.lock().unwrap().parked.get(&key).is_none()); ::futures::future::ok::<(), ()>(()) }).wait().unwrap(); } }