Skip to content

Commit

Permalink
refactor: Move pool into collections module
Browse files Browse the repository at this point in the history
  • Loading branch information
XAMPPRocky committed Jan 8, 2025
1 parent c81df40 commit 8254e02
Show file tree
Hide file tree
Showing 10 changed files with 19 additions and 14 deletions.
2 changes: 1 addition & 1 deletion benches/compression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
2 changes: 1 addition & 1 deletion benches/token_router.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
6 changes: 6 additions & 0 deletions src/collections.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
File renamed without changes.
6 changes: 3 additions & 3 deletions src/components/proxy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -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);
Expand Down
4 changes: 2 additions & 2 deletions src/components/proxy/packet_router.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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);

Expand Down
2 changes: 1 addition & 1 deletion src/components/proxy/sessions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand Down
4 changes: 2 additions & 2 deletions src/filters/token_router.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand All @@ -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());
Expand Down
3 changes: 1 addition & 2 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
4 changes: 2 additions & 2 deletions src/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down Expand Up @@ -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())
}

Expand Down

0 comments on commit 8254e02

Please sign in to comment.