mod.rs

  1// SPDX-License-Identifier: GPL-2.0
  2
  3// Copyright (C) 2025 Google LLC.
  4
  5use kernel::{page::PAGE_SIZE, prelude::*, seq_file::SeqFile, task::Pid};
  6
  7mod tree;
  8use self::tree::{FromArrayAllocs, ReserveNewTreeAlloc, TreeRangeAllocator};
  9
 10mod array;
 11use self::array::{ArrayRangeAllocator, EmptyArrayAlloc};
 12
 13enum DescriptorState<T> {
 14    Reserved(Reservation),
 15    Allocated(Allocation<T>),
 16}
 17
 18impl<T> DescriptorState<T> {
 19    fn new(is_oneway: bool, debug_id: usize, pid: Pid) -> Self {
 20        DescriptorState::Reserved(Reservation {
 21            debug_id,
 22            is_oneway,
 23            pid,
 24        })
 25    }
 26
 27    fn pid(&self) -> Pid {
 28        match self {
 29            DescriptorState::Reserved(inner) => inner.pid,
 30            DescriptorState::Allocated(inner) => inner.reservation.pid,
 31        }
 32    }
 33
 34    fn is_oneway(&self) -> bool {
 35        match self {
 36            DescriptorState::Reserved(inner) => inner.is_oneway,
 37            DescriptorState::Allocated(inner) => inner.reservation.is_oneway,
 38        }
 39    }
 40}
 41
 42#[derive(Clone)]
 43struct Reservation {
 44    debug_id: usize,
 45    is_oneway: bool,
 46    pid: Pid,
 47}
 48
 49impl Reservation {
 50    fn allocate<T>(self, data: Option<T>) -> Allocation<T> {
 51        Allocation {
 52            data,
 53            reservation: self,
 54        }
 55    }
 56}
 57
 58struct Allocation<T> {
 59    reservation: Reservation,
 60    data: Option<T>,
 61}
 62
 63impl<T> Allocation<T> {
 64    fn deallocate(self) -> (Reservation, Option<T>) {
 65        (self.reservation, self.data)
 66    }
 67
 68    fn debug_id(&self) -> usize {
 69        self.reservation.debug_id
 70    }
 71
 72    fn take(&mut self) -> Option<T> {
 73        self.data.take()
 74    }
 75}
 76
 77/// The array implementation must switch to the tree if it wants to go beyond this number of
 78/// ranges.
 79const TREE_THRESHOLD: usize = 8;
 80
 81/// Represents a range of pages that have just become completely free.
 82#[derive(Copy, Clone)]
 83pub(crate) struct FreedRange {
 84    pub(crate) start_page_idx: usize,
 85    pub(crate) end_page_idx: usize,
 86}
 87
 88impl FreedRange {
 89    fn interior_pages(offset: usize, size: usize) -> FreedRange {
 90        FreedRange {
 91            // Divide round up
 92            start_page_idx: offset.div_ceil(PAGE_SIZE),
 93            // Divide round down
 94            end_page_idx: (offset + size) / PAGE_SIZE,
 95        }
 96    }
 97}
 98
 99struct Range<T> {
100    offset: usize,
101    size: usize,
102    state: DescriptorState<T>,
103}
104
105impl<T> Range<T> {
106    fn endpoint(&self) -> usize {
107        self.offset + self.size
108    }
109}
110
111pub(crate) struct RangeAllocator<T> {
112    inner: Impl<T>,
113}
114
115enum Impl<T> {
116    Empty(usize),
117    Array(ArrayRangeAllocator<T>),
118    Tree(TreeRangeAllocator<T>),
119}
120
121impl<T> RangeAllocator<T> {
122    pub(crate) fn new(size: usize) -> Self {
123        Self {
124            inner: Impl::Empty(size),
125        }
126    }
127
128    pub(crate) fn free_oneway_space(&self) -> usize {
129        match &self.inner {
130            Impl::Empty(size) => size / 2,
131            Impl::Array(array) => array.free_oneway_space(),
132            Impl::Tree(tree) => tree.free_oneway_space(),
133        }
134    }
135
136    pub(crate) fn count_buffers(&self) -> usize {
137        match &self.inner {
138            Impl::Empty(_size) => 0,
139            Impl::Array(array) => array.count_buffers(),
140            Impl::Tree(tree) => tree.count_buffers(),
141        }
142    }
143
144    pub(crate) fn debug_print(&self, m: &SeqFile) -> Result<()> {
145        match &self.inner {
146            Impl::Empty(_size) => Ok(()),
147            Impl::Array(array) => array.debug_print(m),
148            Impl::Tree(tree) => tree.debug_print(m),
149        }
150    }
151
152    /// Try to reserve a new buffer, using the provided allocation if necessary.
153    pub(crate) fn reserve_new(&mut self, mut args: ReserveNewArgs<T>) -> Result<ReserveNew<T>> {
154        match &mut self.inner {
155            Impl::Empty(size) => {
156                let empty_array = match args.empty_array_alloc.take() {
157                    Some(empty_array) => ArrayRangeAllocator::new(*size, empty_array),
158                    None => {
159                        return Ok(ReserveNew::NeedAlloc(ReserveNewNeedAlloc {
160                            args,
161                            need_empty_array_alloc: true,
162                            need_new_tree_alloc: false,
163                            need_tree_alloc: false,
164                        }))
165                    }
166                };
167
168                self.inner = Impl::Array(empty_array);
169                self.reserve_new(args)
170            }
171            Impl::Array(array) if array.is_full() => {
172                let allocs = match args.new_tree_alloc {
173                    Some(ref mut allocs) => allocs,
174                    None => {
175                        return Ok(ReserveNew::NeedAlloc(ReserveNewNeedAlloc {
176                            args,
177                            need_empty_array_alloc: false,
178                            need_new_tree_alloc: true,
179                            need_tree_alloc: true,
180                        }))
181                    }
182                };
183
184                let new_tree =
185                    TreeRangeAllocator::from_array(array.total_size(), &mut array.ranges, allocs);
186
187                self.inner = Impl::Tree(new_tree);
188                self.reserve_new(args)
189            }
190            Impl::Array(array) => {
191                let (offset, oneway_spam_detected) =
192                    array.reserve_new(args.debug_id, args.size, args.is_oneway, args.pid)?;
193                Ok(ReserveNew::Success(ReserveNewSuccess {
194                    offset,
195                    oneway_spam_detected,
196                    _empty_array_alloc: args.empty_array_alloc,
197                    _new_tree_alloc: args.new_tree_alloc,
198                    _tree_alloc: args.tree_alloc,
199                }))
200            }
201            Impl::Tree(tree) => {
202                let alloc = match args.tree_alloc {
203                    Some(alloc) => alloc,
204                    None => {
205                        return Ok(ReserveNew::NeedAlloc(ReserveNewNeedAlloc {
206                            args,
207                            need_empty_array_alloc: false,
208                            need_new_tree_alloc: false,
209                            need_tree_alloc: true,
210                        }));
211                    }
212                };
213                let (offset, oneway_spam_detected) =
214                    tree.reserve_new(args.debug_id, args.size, args.is_oneway, args.pid, alloc)?;
215                Ok(ReserveNew::Success(ReserveNewSuccess {
216                    offset,
217                    oneway_spam_detected,
218                    _empty_array_alloc: args.empty_array_alloc,
219                    _new_tree_alloc: args.new_tree_alloc,
220                    _tree_alloc: None,
221                }))
222            }
223        }
224    }
225
226    /// Deletes the allocations at `offset`.
227    pub(crate) fn reservation_abort(&mut self, offset: usize) -> Result<FreedRange> {
228        match &mut self.inner {
229            Impl::Empty(_size) => Err(EINVAL),
230            Impl::Array(array) => array.reservation_abort(offset),
231            Impl::Tree(tree) => {
232                let freed_range = tree.reservation_abort(offset)?;
233                if tree.is_empty() {
234                    self.inner = Impl::Empty(tree.total_size());
235                }
236                Ok(freed_range)
237            }
238        }
239    }
240
241    /// Called when an allocation is no longer in use by the kernel.
242    ///
243    /// The value in `data` will be stored, if any. A mutable reference is used to avoid dropping
244    /// the `T` when an error is returned.
245    pub(crate) fn reservation_commit(&mut self, offset: usize, data: &mut Option<T>) -> Result {
246        match &mut self.inner {
247            Impl::Empty(_size) => Err(EINVAL),
248            Impl::Array(array) => array.reservation_commit(offset, data),
249            Impl::Tree(tree) => tree.reservation_commit(offset, data),
250        }
251    }
252
253    /// Called when the kernel starts using an allocation.
254    ///
255    /// Returns the size of the existing entry and the data associated with it.
256    pub(crate) fn reserve_existing(&mut self, offset: usize) -> Result<(usize, usize, Option<T>)> {
257        match &mut self.inner {
258            Impl::Empty(_size) => Err(EINVAL),
259            Impl::Array(array) => array.reserve_existing(offset),
260            Impl::Tree(tree) => tree.reserve_existing(offset),
261        }
262    }
263
264    /// Call the provided callback at every allocated region.
265    ///
266    /// This destroys the range allocator. Used only during shutdown.
267    pub(crate) fn take_for_each<F: Fn(usize, usize, usize, Option<T>)>(&mut self, callback: F) {
268        match &mut self.inner {
269            Impl::Empty(_size) => {}
270            Impl::Array(array) => array.take_for_each(callback),
271            Impl::Tree(tree) => tree.take_for_each(callback),
272        }
273    }
274}
275
276/// The arguments for `reserve_new`.
277#[derive(Default)]
278pub(crate) struct ReserveNewArgs<T> {
279    pub(crate) size: usize,
280    pub(crate) is_oneway: bool,
281    pub(crate) debug_id: usize,
282    pub(crate) pid: Pid,
283    pub(crate) empty_array_alloc: Option<EmptyArrayAlloc<T>>,
284    pub(crate) new_tree_alloc: Option<FromArrayAllocs<T>>,
285    pub(crate) tree_alloc: Option<ReserveNewTreeAlloc<T>>,
286}
287
288/// The return type of `ReserveNew`.
289pub(crate) enum ReserveNew<T> {
290    Success(ReserveNewSuccess<T>),
291    NeedAlloc(ReserveNewNeedAlloc<T>),
292}
293
294/// Returned by `reserve_new` when the reservation was successul.
295pub(crate) struct ReserveNewSuccess<T> {
296    pub(crate) offset: usize,
297    pub(crate) oneway_spam_detected: bool,
298
299    // If the user supplied an allocation that we did not end up using, then we return it here.
300    // The caller will kfree it outside of the lock.
301    _empty_array_alloc: Option<EmptyArrayAlloc<T>>,
302    _new_tree_alloc: Option<FromArrayAllocs<T>>,
303    _tree_alloc: Option<ReserveNewTreeAlloc<T>>,
304}
305
306/// Returned by `reserve_new` to request the caller to make an allocation before calling the method
307/// again.
308pub(crate) struct ReserveNewNeedAlloc<T> {
309    args: ReserveNewArgs<T>,
310    need_empty_array_alloc: bool,
311    need_new_tree_alloc: bool,
312    need_tree_alloc: bool,
313}
314
315impl<T> ReserveNewNeedAlloc<T> {
316    /// Make the necessary allocations for another call to `reserve_new`.
317    pub(crate) fn make_alloc(mut self) -> Result<ReserveNewArgs<T>> {
318        if self.need_empty_array_alloc && self.args.empty_array_alloc.is_none() {
319            self.args.empty_array_alloc = Some(EmptyArrayAlloc::try_new(TREE_THRESHOLD)?);
320        }
321        if self.need_new_tree_alloc && self.args.new_tree_alloc.is_none() {
322            self.args.new_tree_alloc = Some(FromArrayAllocs::try_new(TREE_THRESHOLD)?);
323        }
324        if self.need_tree_alloc && self.args.tree_alloc.is_none() {
325            self.args.tree_alloc = Some(ReserveNewTreeAlloc::try_new()?);
326        }
327        Ok(self.args)
328    }
329}