safa_buffer_pool/
lib.rs

1use std::sync::Arc;
2
3use tokio::sync::Mutex;
4
5///BufferPool for Multi-threading context
6#[derive(Clone, Default)]
7pub struct MultiThreadBufferPool {
8    inner_arc: Arc<Mutex<BufferPool>>,
9}
10
11impl MultiThreadBufferPool {
12    pub async fn get(&self) -> Option<BufferGuard> {
13        let mut pool = self.inner_arc.lock().await;
14        if let Some(new_buffer) = pool.all_available_buffer.pop() {
15            return Some(BufferGuard {
16                pool: self.clone(),
17                buffer: new_buffer,
18            });
19        }
20
21        if pool.total_nb_buffer == pool.max_nb_buffer {
22            return None;
23        }
24        pool.total_nb_buffer += 1;
25
26        let buffer_size = pool.buffer_size;
27
28        drop(pool); //Free lock
29
30        Some(BufferGuard {
31            pool: self.clone(),
32            buffer: vec![0u8; buffer_size].into_boxed_slice(),
33        })
34    }
35}
36
37///Unique BufferPool
38#[derive(Default)]
39struct BufferPool {
40    total_nb_buffer: usize,
41    max_nb_buffer: usize,
42    min_nb_buffer: usize,
43    buffer_size: usize,
44    all_available_buffer: Vec<Box<[u8]>>,
45}
46
47pub struct BufferPoolBuilder {
48    max_nb_buffer: usize,
49    min_nb_buffer: usize,
50    buffer_size: usize,
51}
52
53impl Default for BufferPoolBuilder {
54    fn default() -> Self {
55        Self {
56            max_nb_buffer: 1024,
57            min_nb_buffer: 1024,
58            buffer_size: 10240,
59        }
60    }
61}
62
63impl BufferPoolBuilder {
64    pub fn new() -> Self {
65        Self::default()
66    }
67
68    pub fn set_max_number_of_buffer(&mut self, number: usize) {
69        self.max_nb_buffer = number;
70    }
71    pub fn set_min_number_of_buffer(&mut self, number: usize) {
72        self.min_nb_buffer = number;
73    }
74    pub fn set_buffer_size(&mut self, number: usize) {
75        self.buffer_size = number;
76    }
77
78    pub fn build_multi_thread_buffer_pool(&self) -> MultiThreadBufferPool {
79        let mut all_buffer = Vec::with_capacity(self.max_nb_buffer);
80
81        for _ in 0..self.min_nb_buffer {
82            all_buffer.push(vec![0u8; self.buffer_size].into_boxed_slice());
83        }
84
85        MultiThreadBufferPool {
86            inner_arc: Arc::new(Mutex::new(BufferPool {
87                total_nb_buffer: self.min_nb_buffer,
88                max_nb_buffer: self.max_nb_buffer,
89                min_nb_buffer: self.min_nb_buffer,
90                buffer_size: self.buffer_size,
91                all_available_buffer: all_buffer,
92            })),
93        }
94    }
95}
96
97pub struct BufferGuard {
98    pool: MultiThreadBufferPool,
99    buffer: Box<[u8]>,
100}
101
102impl Drop for BufferGuard {
103    fn drop(&mut self) {
104        let pool = std::mem::take(&mut self.pool);
105        let buffer = std::mem::take(&mut self.buffer);
106
107        tokio::spawn(async move {
108            let mut pool = pool.inner_arc.lock().await;
109            pool.all_available_buffer.push(buffer);
110        });
111    }
112}
safa_buffer_pool/lib.rs

safa_buffer_pool/
lib.rs
