minipool/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::ops::{Deref, DerefMut};
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::Arc;
6
7use tokio::sync::{mpsc, Mutex};
8
9#[cfg(test)]
10mod tests;
11
12/// Lightweight, generic tokio-based pool implementation for Rust.
13///
14/// This pool is save to be put into an [`Arc`].
15pub struct Pool<T: PoolEntry> {
16  len: Arc<AtomicUsize>,
17  tx: mpsc::UnboundedSender<T>,
18  rx: Mutex<mpsc::UnboundedReceiver<T>>,
19}
20
21/// Repsesents a entity that can be a member of a [`Pool`].
22pub trait PoolEntry {
23  /// Determinates if the PoolEntry should be reinserted into the pool and thereby is able to be
24  /// reused in future invocations.
25  ///
26  /// # Examples
27  /// - A closed TCP connection should probably be not inserted.
28  /// - A simple buffer can always be reinserted.
29  fn is_closed(&self) -> bool;
30}
31
32/// A wrapper around an occupied entry of a [`Pool`].
33///
34/// When this structure is dropped (falls out of
35/// scope) and the entry is not yet closed ([`is_closed`]), the entry will be reinserted into the pool.
36///
37/// The data protected by the mutex can be accessed through this guard via its
38/// [`Deref`] and [`DerefMut`] implementations.
39///
40/// This structure is created by the [`acquire`] method on [`Pool`].
41///
42/// [`is_closed`]: PoolEntry::is_closed
43/// [`acquire`]: Pool::acquire
44#[must_use = "if unused the PoolEntry will immediately reinserted"]
45pub struct PoolGuard<T: PoolEntry> {
46  inner: Option<T>,
47  back: mpsc::UnboundedSender<T>,
48  pool_len: Arc<AtomicUsize>,
49}
50
51unsafe impl<T: Send + PoolEntry> Send for PoolGuard<T> {}
52unsafe impl<T: Sync + PoolEntry> Sync for PoolGuard<T> {}
53
54impl<T: PoolEntry> Drop for PoolGuard<T> {
55  fn drop(&mut self) {
56    if let Some(inner) = self.inner.take() {
57      if inner.is_closed() {
58        self.pool_len.fetch_sub(1, Ordering::Relaxed);
59      } else if self.back.send(inner).is_err() {
60        // NOOP: channel is closed, probably because the program is shutting down
61        //       so it's not bad if the entry get's dropped
62      }
63    }
64  }
65}
66
67impl<T: PoolEntry> Deref for PoolGuard<T> {
68  type Target = T;
69  fn deref(&self) -> &Self::Target {
70    match self.inner.as_ref() {
71      Some(inner) => inner,
72      None => unreachable!("PoolGuard can only be None after itself got dropped."),
73    }
74  }
75}
76
77impl<T: PoolEntry> DerefMut for PoolGuard<T> {
78  fn deref_mut(&mut self) -> &mut Self::Target {
79    match self.inner.as_mut() {
80      Some(inner) => inner,
81      None => unreachable!("PoolGuard can only be None after itself got dropped."),
82    }
83  }
84}
85
86impl<T: PoolEntry> Default for Pool<T> {
87  fn default() -> Self {
88    let (tx, rx) = mpsc::unbounded_channel();
89    Self {
90      len: Arc::new(AtomicUsize::new(0)),
91      tx,
92      rx: Mutex::new(rx),
93    }
94  }
95}
96
97impl<T: PoolEntry> Pool<T> {
98  /// Tries to acquire an entry from the pool. May blocks untail an entry becomes
99  /// available again. Entries are handed out on a first-come, first-served basis.
100  ///
101  /// Returns [`None`] if the pool is empty.
102  pub async fn acquire(&self) -> Option<PoolGuard<T>> {
103    if self.is_empty() {
104      return None;
105    }
106
107    let mut rx = self.rx.lock().await;
108
109    let inner = match rx.recv().await {
110      Some(inner) => inner,
111      None => unreachable!("Channel is closed"),
112    };
113
114    Some(PoolGuard {
115      inner: Some(inner),
116      back: self.tx.clone(),
117      pool_len: self.len.clone(),
118    })
119  }
120}
121
122impl<T: PoolEntry> Pool<T> {
123  /// Inserts an new [`PoolEntry`] into the pool.
124  pub fn push(&self, inner: T) {
125    self.len.fetch_add(1, Ordering::Relaxed);
126    if self.tx.send(inner).is_err() {
127      unreachable!("Cannel is closed");
128    }
129  }
130
131  /// Returns the number of entries in the pool.
132  pub fn len(&self) -> usize {
133    self.len.load(Ordering::SeqCst)
134  }
135
136  /// Returns [`true`] if the pool has a length of 0.
137  pub fn is_empty(&self) -> bool {
138    self.len() == 0
139  }
140}
minipool/lib.rs

minipool/
lib.rs
