refactor: Move pool into collections module (#1057)

benches/compression.rs

@@ -2,8 +2,8 @@ mod shared;
 
 use divan::Bencher;
 use quilkin::{
+    collections::BufferPool,
     filters::compress::{Compressor, Mode},
-    pool::BufferPool,
 };
 use shared::*;
 use std::sync::Arc;

benches/token_router.rs

@@ -23,7 +23,7 @@ fn token_router(b: Bencher, token_kind: &str) {
     }
 
     let total_token_size: usize = tokens.iter().map(|t| t.len()).sum();
-    let pool = std::sync::Arc::new(quilkin::pool::BufferPool::new(1, 1));
+    let pool = std::sync::Arc::new(quilkin::collections::BufferPool::new(1, 1));
 
     let mut rand = rand::rngs::SmallRng::seed_from_u64(42);
 

crates/test/src/lib.rs

@@ -1,4 +1,5 @@
 use quilkin::{
+    collections::{BufferPool, PoolBuffer},
     components::{self, RunArgs},
     config::Providers,
     net::TcpListener,
@@ -9,11 +10,11 @@ pub use serde_json::json;
 use std::{net::SocketAddr, num::NonZeroUsize, path::PathBuf, sync::Arc};
 use tokio::sync::mpsc;
 
-pub static BUFFER_POOL: once_cell::sync::Lazy<Arc<quilkin::pool::BufferPool>> =
-    once_cell::sync::Lazy::new(|| Arc::new(quilkin::pool::BufferPool::default()));
+pub static BUFFER_POOL: once_cell::sync::Lazy<Arc<BufferPool>> =
+    once_cell::sync::Lazy::new(|| Arc::new(BufferPool::default()));
 
 #[inline]
-pub fn alloc_buffer(data: impl AsRef<[u8]>) -> quilkin::pool::PoolBuffer {
+pub fn alloc_buffer(data: impl AsRef<[u8]>) -> PoolBuffer {
     BUFFER_POOL.clone().alloc_slice(data.as_ref())
 }
 

src/collections.rs

@@ -16,4 +16,10 @@
 
 //! Collection types designed for use with Quilkin.
 
+pub mod pool;
 pub mod ttl;
+
+pub use self::{
+    pool::{BufferPool, FrozenPoolBuffer, PoolBuffer},
+    ttl::TtlMap,
+};

src/components/proxy.rs

@@ -107,14 +107,14 @@ pub struct SendPacket {
     /// The destination address of the packet
     pub destination: socket2::SockAddr,
     /// The packet data being sent
-    pub data: crate::pool::FrozenPoolBuffer,
+    pub data: crate::collections::FrozenPoolBuffer,
     /// The asn info for the sender, used for metrics
     pub asn_info: Option<crate::net::maxmind_db::MetricsIpNetEntry>,
 }
 
 pub struct RecvPacket {
     pub source: SocketAddr,
-    pub data: crate::pool::PoolBuffer,
+    pub data: crate::collections::PoolBuffer,
 }
 
 #[derive(Clone, Debug)]
@@ -356,7 +356,7 @@ impl Proxy {
         }
 
         let num_workers = self.num_workers.get();
-        let buffer_pool = Arc::new(crate::pool::BufferPool::new(num_workers, 2 * 1024));
+        let buffer_pool = Arc::new(crate::collections::BufferPool::new(num_workers, 2 * 1024));
 
         let mut worker_sends = Vec::with_capacity(num_workers);
         let mut session_sends = Vec::with_capacity(num_workers);

src/components/proxy/io_uring_shared.rs

@@ -21,9 +21,9 @@
 //! enough that it doesn't make sense to share the same code
 
 use crate::{
+    collections::PoolBuffer,
     components::proxy::{self, PendingSends, PipelineError, SendPacket},
     metrics,
-    pool::PoolBuffer,
     time::UtcTimestamp,
 };
 use io_uring::{squeue::Entry, types::Fd};
@@ -295,7 +295,7 @@ impl<'uring> LoopCtx<'uring> {
 
     /// Enqueues a recv_from on the socket
     #[inline]
-    fn enqueue_recv(&mut self, buffer: crate::pool::PoolBuffer) {
+    fn enqueue_recv(&mut self, buffer: crate::collections::PoolBuffer) {
         let packet = LoopPacketInner::Recv(RecvPacket {
             buffer,
             source: empty_net_addr(),
@@ -423,7 +423,7 @@ impl IoUringLoop {
         thread_name: String,
         mut ctx: PacketProcessorCtx,
         pending_sends: (PendingSends, EventFd),
-        buffer_pool: Arc<crate::pool::BufferPool>,
+        buffer_pool: Arc<crate::collections::BufferPool>,
     ) -> Result<(), PipelineError> {
         let dispatcher = tracing::dispatcher::get_default(|d| d.clone());
 

src/components/proxy/packet_router.rs

@@ -150,7 +150,7 @@ pub struct DownstreamReceiveWorkerConfig {
     pub config: Arc<Config>,
     pub sessions: Arc<SessionPool>,
     pub error_sender: super::error::ErrorSender,
-    pub buffer_pool: Arc<crate::pool::BufferPool>,
+    pub buffer_pool: Arc<crate::collections::BufferPool>,
 }
 
 /// Spawns a background task that sits in a loop, receiving packets from the passed in socket.
@@ -163,7 +163,7 @@ pub async fn spawn_receivers(
     socket: socket2::Socket,
     worker_sends: Vec<(super::PendingSends, super::PacketSendReceiver)>,
     sessions: &Arc<SessionPool>,
-    buffer_pool: Arc<crate::pool::BufferPool>,
+    buffer_pool: Arc<crate::collections::BufferPool>,
 ) -> crate::Result<()> {
     let (error_sender, mut error_receiver) = mpsc::channel(128);
 

src/components/proxy/sessions.rs

@@ -25,12 +25,12 @@ use std::{
 use tokio::time::Instant;
 
 use crate::{
+    collections::{BufferPool, FrozenPoolBuffer, PoolBuffer},
     components::proxy::SendPacket,
     config::Config,
     filters::Filter,
     metrics,
     net::maxmind_db::{IpNetEntry, MetricsIpNetEntry},
-    pool::{BufferPool, FrozenPoolBuffer, PoolBuffer},
     time::UtcTimestamp,
     Loggable,
 };

src/filters/token_router.rs

@@ -309,7 +309,7 @@ mod tests {
 
     fn with_ctx(
         dest: &mut Vec<crate::net::EndpointAddress>,
-        test: impl FnOnce(ReadContext<'_, crate::pool::PoolBuffer>),
+        test: impl FnOnce(ReadContext<'_, crate::collections::PoolBuffer>),
     ) {
         let endpoint1 = Endpoint::with_metadata(
             "127.0.0.1:80".parse().unwrap(),
@@ -324,7 +324,7 @@ mod tests {
             },
         );
 
-        let pool = std::sync::Arc::new(crate::pool::BufferPool::new(1, 5));
+        let pool = std::sync::Arc::new(crate::collections::BufferPool::new(1, 5));
 
         let endpoints = crate::net::cluster::ClusterMap::default();
         endpoints.insert_default([endpoint1, endpoint2].into());

src/lib.rs

@@ -17,9 +17,8 @@
 #![deny(unused_must_use)]
 
 pub mod alloc;
-pub(crate) mod collections;
+pub mod collections;
 pub(crate) mod metrics;
-pub mod pool;
 pub mod time;
 
 // Above other modules for thr `uring_spawn` macro.

src/test.rs

@@ -22,12 +22,12 @@ use tokio::sync::{mpsc, oneshot};
 use tracing_subscriber::EnvFilter;
 
 use crate::{
+    collections::BufferPool,
     config::Config,
     filters::{prelude::*, FilterRegistry},
     net::endpoint::metadata::Value,
     net::endpoint::{Endpoint, EndpointAddress},
     net::DualStackEpollSocket as DualStackLocalSocket,
-    pool::BufferPool,
     ShutdownKind, ShutdownRx, ShutdownTx,
 };
 
@@ -354,7 +354,7 @@ pub static BUFFER_POOL: once_cell::sync::Lazy<Arc<BufferPool>> =
     once_cell::sync::Lazy::new(|| Arc::new(BufferPool::default()));
 
 #[inline]
-pub fn alloc_buffer(data: impl AsRef<[u8]>) -> crate::pool::PoolBuffer {
+pub fn alloc_buffer(data: impl AsRef<[u8]>) -> crate::collections::PoolBuffer {
     BUFFER_POOL.clone().alloc_slice(data.as_ref())
 }