Skip to main content

dma_api/
pool.rs

1use alloc::{
2    collections::VecDeque,
3    sync::{Arc, Weak},
4};
5use core::ops::{Deref, DerefMut};
6
7use ax_kspin::SpinNoIrq as Mutex;
8
9use crate::{ContiguousArray, DeviceDma, DmaDirection, DmaError};
10
11#[derive(Clone, Debug)]
12pub(crate) struct ContiguousBufferConfig {
13    pub size: usize,
14    pub align: usize,
15    pub direction: DmaDirection,
16}
17
18#[derive(Clone)]
19pub struct ContiguousBufferPool {
20    inner: Arc<Mutex<Inner>>,
21}
22
23pub struct ContiguousBuffer {
24    data: Option<ContiguousArray<u8>>,
25    pool: Weak<Mutex<Inner>>,
26}
27
28unsafe impl Send for ContiguousBuffer {}
29
30impl Deref for ContiguousBuffer {
31    type Target = ContiguousArray<u8>;
32
33    fn deref(&self) -> &Self::Target {
34        self.data.as_ref().unwrap()
35    }
36}
37
38impl DerefMut for ContiguousBuffer {
39    fn deref_mut(&mut self) -> &mut Self::Target {
40        self.data.as_mut().unwrap()
41    }
42}
43
44impl Drop for ContiguousBuffer {
45    fn drop(&mut self) {
46        if let Some(data) = self.data.take()
47            && let Some(pool) = self.pool.upgrade()
48        {
49            let mut inner = pool.lock();
50            inner.dealloc(data);
51        }
52    }
53}
54
55struct Inner {
56    dev: DeviceDma,
57    config: ContiguousBufferConfig,
58    pool: VecDeque<ContiguousArray<u8>>,
59}
60
61impl Inner {
62    fn alloc(&mut self) -> Option<ContiguousArray<u8>> {
63        self.pool.pop_front()
64    }
65
66    fn dealloc(&mut self, data: ContiguousArray<u8>) {
67        self.pool.push_back(data);
68    }
69}
70
71impl ContiguousBufferPool {
72    pub(crate) fn with_capacity(
73        dev: DeviceDma,
74        config: ContiguousBufferConfig,
75        cap: usize,
76    ) -> ContiguousBufferPool {
77        let mut pool = VecDeque::with_capacity(cap);
78        for _ in 0..cap {
79            if let Ok(data) = ContiguousArray::new_zero_with_align(
80                &dev,
81                config.size,
82                config.align,
83                config.direction,
84            ) {
85                pool.push_back(data);
86            }
87        }
88
89        ContiguousBufferPool {
90            inner: Arc::new(Mutex::new(Inner { dev, pool, config })),
91        }
92    }
93
94    pub fn alloc(&self) -> Result<ContiguousBuffer, DmaError> {
95        let config;
96        let dev;
97        {
98            let mut inner = self.inner.lock();
99            if let Some(data) = inner.alloc() {
100                return Ok(ContiguousBuffer {
101                    data: Some(data),
102                    pool: Arc::downgrade(&self.inner),
103                });
104            } else {
105                config = inner.config.clone();
106                dev = inner.dev.clone();
107            }
108        };
109
110        let data = ContiguousArray::new_zero_with_align(
111            &dev,
112            config.size,
113            config.align,
114            config.direction,
115        )?;
116        Ok(ContiguousBuffer {
117            data: Some(data),
118            pool: Arc::downgrade(&self.inner),
119        })
120    }
121}