Skip to content

Commit

Permalink
feat(client): add poison to Connected (#3145)
Browse files Browse the repository at this point in the history
Add `poison` method to `Connected`. This allows callers to mark a connection as poisoned which prevents the pool from reusing it on subsequent requests. `is_open` will consider poisoning prior to returning a connection to the pool.
  • Loading branch information
rcoh committed Feb 20, 2023
1 parent 40c01df commit 37ed5a2
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 0 deletions.
4 changes: 4 additions & 0 deletions src/client/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -689,6 +689,10 @@ where
B: Send + 'static,
{
fn is_open(&self) -> bool {
if self.conn_info.poisoned.poisoned() {
trace!("marking {:?} as closed because it was poisoned", self.conn_info);
return false;
}
match self.tx {
PoolTx::Http1(ref tx) => tx.is_ready(),
#[cfg(feature = "http2")]
Expand Down
43 changes: 43 additions & 0 deletions src/client/connect/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,9 @@
//! [`AsyncWrite`]: tokio::io::AsyncWrite
//! [`Connection`]: Connection
use std::fmt;
use std::fmt::{Debug, Formatter};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use ::http::Extensions;

Expand Down Expand Up @@ -113,6 +116,34 @@ pub struct Connected {
pub(super) alpn: Alpn,
pub(super) is_proxied: bool,
pub(super) extra: Option<Extra>,
pub(super) poisoned: PoisonPill,
}

#[derive(Clone)]
pub(crate) struct PoisonPill {
poisoned: Arc<AtomicBool>,
}

impl Debug for PoisonPill {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
// print the address of the pill—this makes debugging issues much easier
write!(f, "PoisonPill@{:p} {{ poisoned: {} }}", self.poisoned, self.poisoned.load(Ordering::Relaxed))
}
}

impl PoisonPill {
pub(crate) fn healthy() -> Self {
Self {
poisoned: Arc::new(AtomicBool::new(false)),
}
}
pub(crate) fn poison(&self) {
self.poisoned.store(true, Ordering::Relaxed)
}

pub(crate) fn poisoned(&self) -> bool {
self.poisoned.load(Ordering::Relaxed)
}
}

pub(super) struct Extra(Box<dyn ExtraInner>);
Expand All @@ -130,6 +161,7 @@ impl Connected {
alpn: Alpn::None,
is_proxied: false,
extra: None,
poisoned: PoisonPill::healthy(),
}
}

Expand Down Expand Up @@ -189,6 +221,16 @@ impl Connected {
self.alpn == Alpn::H2
}

/// Poison this connection
///
/// A poisoned connection will not be reused for subsequent requests by the pool
pub fn poison(&self) {
self.poisoned.poison();
tracing::debug!(
poison_pill = ?self.poisoned, "connection was poisoned"
);
}

// Don't public expose that `Connected` is `Clone`, unsure if we want to
// keep that contract...
#[cfg(feature = "http2")]
Expand All @@ -197,6 +239,7 @@ impl Connected {
alpn: self.alpn.clone(),
is_proxied: self.is_proxied,
extra: self.extra.clone(),
poisoned: self.poisoned.clone(),
}
}
}
Expand Down

0 comments on commit 37ed5a2

Please sign in to comment.