compio_pool/
guard.rs

1//! The checkout guard and the cancellation-safety primitive that goes with it.
2
3use std::{
4    cell::Cell,
5    ops::{Deref, DerefMut},
6    rc::Rc,
7};
8
9use crate::{
10    exchange::Exchange,
11    manage::Manage,
12    pool::Pool,
13    shard::Shard,
14    slot::{Slot, SlotMeta},
15};
16
17/// A connection checked out of the pool.
18///
19/// Derefs to the underlying connection. Returns itself to the pool on drop —
20/// unless it has been poisoned, in which case it is closed instead.
21///
22/// # Cancellation safety
23///
24/// This is the part that differs most from a readiness-based pool. With
25/// completion-based IO, dropping a future with an operation in flight does not
26/// unwind the operation: the kernel still owns the buffer and the peer's
27/// response is still coming. `compio` keeps the *buffer* sound, but your
28/// *protocol* is now out of step — the next reader on this connection will see
29/// the tail of someone else's response.
30///
31/// A connection cancelled mid-operation must therefore be destroyed, never
32/// reused. Wrap each operation in [`begin_op`](Pooled::begin_op):
33///
34/// ```ignore
35/// let op = conn.begin_op();
36/// let n = conn.read(&mut buf).await?;   // if this await is cancelled…
37/// op.complete_op();                     // …this never runs, and the conn is dropped
38/// ```
39///
40/// If you never cancel (no `select!`, no timeouts, no early return between
41/// submit and completion) you can skip it. Everyone believes that about their
42/// code right up until they add a timeout.
43pub struct Pooled<M: Manage, X: Exchange<M>> {
44    slot: Option<Slot<M::Connection>>,
45    shard: Rc<Shard<M>>,
46    pool: Pool<M, X>,
47    poison: Rc<Cell<bool>>,
48}
49
50impl<M: Manage, X: Exchange<M>> Pooled<M, X> {
51    pub(crate) fn new(slot: Slot<M::Connection>, shard: Rc<Shard<M>>, pool: Pool<M, X>) -> Self {
52        Self {
53            slot: Some(slot),
54            shard,
55            pool,
56            poison: Rc::new(Cell::new(false)),
57        }
58    }
59
60    /// Lifecycle metadata: age, idle time, checkout count.
61    pub fn meta(&self) -> &SlotMeta {
62        &self.slot.as_ref().expect("slot present until drop").meta
63    }
64
65    /// Marks this connection unusable. It is closed on drop instead of pooled.
66    ///
67    /// Call this whenever you learn the connection is in an unknown state: a
68    /// protocol error, a partial write, a response you could not parse.
69    pub fn poison(&self) {
70        self.poison.set(true);
71    }
72
73    /// True if [`poison`](Pooled::poison) has been called, directly or by a
74    /// dropped [`OpGuard`].
75    pub fn is_poisoned(&self) -> bool {
76        self.poison.get()
77    }
78
79    /// Arms cancellation protection for one operation.
80    ///
81    /// The returned guard poisons this connection when dropped, unless
82    /// [`OpGuard::complete`] runs first. It borrows nothing from `self`, so the
83    /// connection stays fully usable while the guard is alive.
84    pub fn begin_op(&self) -> OpGuard {
85        OpGuard {
86            poison: self.poison.clone(),
87            armed: true,
88        }
89    }
90
91    /// Removes the connection from the pool's control and hands it to you.
92    ///
93    /// The shard's budget is freed immediately, so a replacement may be opened.
94    pub fn take(mut self) -> M::Connection {
95        let slot = self.slot.take().expect("slot present until drop");
96        self.pool.forget(&self.shard);
97        slot.conn
98    }
99}
100
101impl<M: Manage, X: Exchange<M>> Deref for Pooled<M, X> {
102    type Target = M::Connection;
103
104    fn deref(&self) -> &Self::Target {
105        &self.slot.as_ref().expect("slot present until drop").conn
106    }
107}
108
109impl<M: Manage, X: Exchange<M>> DerefMut for Pooled<M, X> {
110    fn deref_mut(&mut self) -> &mut Self::Target {
111        &mut self.slot.as_mut().expect("slot present until drop").conn
112    }
113}
114
115impl<M: Manage, X: Exchange<M>> Drop for Pooled<M, X> {
116    fn drop(&mut self) {
117        // `Drop` cannot await, so the connection goes back synchronously and
118        // any async validation happens on the next `acquire`, in
119        // `Manage::recycle`. That is why `recycle` is an acquire-side hook.
120        if let Some(slot) = self.slot.take() {
121            self.pool.release(&self.shard, slot, self.poison.get());
122        }
123    }
124}
125
126impl<M: Manage, X: Exchange<M>> std::fmt::Debug for Pooled<M, X> {
127    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128        f.debug_struct("Pooled")
129            .field("poisoned", &self.poison.get())
130            .field("meta", self.meta())
131            .finish_non_exhaustive()
132    }
133}
134
135/// Poisons its connection unless [`complete`](OpGuard::complete) is called.
136///
137/// Created by [`Pooled::begin_op`]. See that type's docs for why this exists.
138#[must_use = "an OpGuard that is dropped immediately poisons the connection"]
139pub struct OpGuard {
140    poison: Rc<Cell<bool>>,
141    armed: bool,
142}
143
144impl OpGuard {
145    /// Disarms the guard: the operation finished, the connection is intact.
146    pub fn complete_op(mut self) {
147        self.armed = false;
148    }
149}
150
151impl Drop for OpGuard {
152    fn drop(&mut self) {
153        if self.armed {
154            self.poison.set(true);
155        }
156    }
157}
158
159impl std::fmt::Debug for OpGuard {
160    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
161        f.debug_struct("OpGuard")
162            .field("armed", &self.armed)
163            .finish()
164    }
165}
compio_pool/guard.rs

compio_pool/
guard.rs
