tree.rs

  1// SPDX-License-Identifier: GPL-2.0
  2
  3// Copyright (C) 2025 Google LLC.
  4
  5use kernel::{
  6    page::PAGE_SIZE,
  7    prelude::*,
  8    rbtree::{RBTree, RBTreeNode, RBTreeNodeReservation},
  9    seq_file::SeqFile,
 10    seq_print,
 11    task::Pid,
 12};
 13
 14use crate::range_alloc::{DescriptorState, FreedRange, Range};
 15
 16/// Keeps track of allocations in a process' mmap.
 17///
 18/// Each process has an mmap where the data for incoming transactions will be placed. This struct
 19/// keeps track of allocations made in the mmap. For each allocation, we store a descriptor that
 20/// has metadata related to the allocation. We also keep track of available free space.
 21pub(super) struct TreeRangeAllocator<T> {
 22    /// This collection contains descriptors for *both* ranges containing an allocation, *and* free
 23    /// ranges between allocations. The free ranges get merged, so there are never two free ranges
 24    /// next to each other.
 25    tree: RBTree<usize, Descriptor<T>>,
 26    /// Contains an entry for every free range in `self.tree`. This tree sorts the ranges by size,
 27    /// letting us look up the smallest range whose size is at least some lower bound.
 28    free_tree: RBTree<FreeKey, ()>,
 29    size: usize,
 30    free_oneway_space: usize,
 31}
 32
 33impl<T> TreeRangeAllocator<T> {
 34    pub(crate) fn from_array(
 35        size: usize,
 36        ranges: &mut KVec<Range<T>>,
 37        alloc: &mut FromArrayAllocs<T>,
 38    ) -> Self {
 39        let mut tree = TreeRangeAllocator {
 40            tree: RBTree::new(),
 41            free_tree: RBTree::new(),
 42            size,
 43            free_oneway_space: size / 2,
 44        };
 45
 46        let mut free_offset = 0;
 47        for range in ranges.drain_all() {
 48            let free_size = range.offset - free_offset;
 49            if free_size > 0 {
 50                let free_node = alloc.free_tree.pop().unwrap();
 51                tree.free_tree
 52                    .insert(free_node.into_node((free_size, free_offset), ()));
 53                let tree_node = alloc.tree.pop().unwrap();
 54                tree.tree.insert(
 55                    tree_node.into_node(free_offset, Descriptor::new(free_offset, free_size)),
 56                );
 57            }
 58            free_offset = range.endpoint();
 59
 60            if range.state.is_oneway() {
 61                tree.free_oneway_space = tree.free_oneway_space.saturating_sub(range.size);
 62            }
 63
 64            let free_res = alloc.free_tree.pop().unwrap();
 65            let tree_node = alloc.tree.pop().unwrap();
 66            let mut desc = Descriptor::new(range.offset, range.size);
 67            desc.state = Some((range.state, free_res));
 68            tree.tree.insert(tree_node.into_node(range.offset, desc));
 69        }
 70
 71        // After the last range, we may need a free range.
 72        if free_offset < size {
 73            let free_size = size - free_offset;
 74            let free_node = alloc.free_tree.pop().unwrap();
 75            tree.free_tree
 76                .insert(free_node.into_node((free_size, free_offset), ()));
 77            let tree_node = alloc.tree.pop().unwrap();
 78            tree.tree
 79                .insert(tree_node.into_node(free_offset, Descriptor::new(free_offset, free_size)));
 80        }
 81
 82        tree
 83    }
 84
 85    pub(crate) fn is_empty(&self) -> bool {
 86        let mut tree_iter = self.tree.values();
 87        // There's always at least one range, because index zero is either the start of a free or
 88        // allocated range.
 89        let first_value = tree_iter.next().unwrap();
 90        if tree_iter.next().is_some() {
 91            // There are never two free ranges next to each other, so if there is more than one
 92            // descriptor, then at least one of them must hold an allocated range.
 93            return false;
 94        }
 95        // There is only one descriptor. Return true if it is for a free range.
 96        first_value.state.is_none()
 97    }
 98
 99    pub(crate) fn total_size(&self) -> usize {
100        self.size
101    }
102
103    pub(crate) fn free_oneway_space(&self) -> usize {
104        self.free_oneway_space
105    }
106
107    pub(crate) fn count_buffers(&self) -> usize {
108        self.tree
109            .values()
110            .filter(|desc| desc.state.is_some())
111            .count()
112    }
113
114    pub(crate) fn debug_print(&self, m: &SeqFile) -> Result<()> {
115        for desc in self.tree.values() {
116            let state = match &desc.state {
117                Some(state) => &state.0,
118                None => continue,
119            };
120            seq_print!(
121                m,
122                "  buffer: {} size {} pid {}",
123                desc.offset,
124                desc.size,
125                state.pid(),
126            );
127            if state.is_oneway() {
128                seq_print!(m, " oneway");
129            }
130            match state {
131                DescriptorState::Reserved(_res) => {
132                    seq_print!(m, " reserved\n");
133                }
134                DescriptorState::Allocated(_alloc) => {
135                    seq_print!(m, " allocated\n");
136                }
137            }
138        }
139        Ok(())
140    }
141
142    fn find_best_match(&mut self, size: usize) -> Option<&mut Descriptor<T>> {
143        let free_cursor = self.free_tree.cursor_lower_bound(&(size, 0))?;
144        let ((_, offset), ()) = free_cursor.current();
145        self.tree.get_mut(offset)
146    }
147
148    /// Try to reserve a new buffer, using the provided allocation if necessary.
149    pub(crate) fn reserve_new(
150        &mut self,
151        debug_id: usize,
152        size: usize,
153        is_oneway: bool,
154        pid: Pid,
155        alloc: ReserveNewTreeAlloc<T>,
156    ) -> Result<(usize, bool)> {
157        // Compute new value of free_oneway_space, which is set only on success.
158        let new_oneway_space = if is_oneway {
159            match self.free_oneway_space.checked_sub(size) {
160                Some(new_oneway_space) => new_oneway_space,
161                None => return Err(ENOSPC),
162            }
163        } else {
164            self.free_oneway_space
165        };
166
167        let (found_size, found_off, tree_node, free_tree_node) = match self.find_best_match(size) {
168            None => {
169                pr_warn!("ENOSPC from range_alloc.reserve_new - size: {}", size);
170                return Err(ENOSPC);
171            }
172            Some(desc) => {
173                let found_size = desc.size;
174                let found_offset = desc.offset;
175
176                // In case we need to break up the descriptor
177                let new_desc = Descriptor::new(found_offset + size, found_size - size);
178                let (tree_node, free_tree_node, desc_node_res) = alloc.initialize(new_desc);
179
180                desc.state = Some((
181                    DescriptorState::new(is_oneway, debug_id, pid),
182                    desc_node_res,
183                ));
184                desc.size = size;
185
186                (found_size, found_offset, tree_node, free_tree_node)
187            }
188        };
189        self.free_oneway_space = new_oneway_space;
190        self.free_tree.remove(&(found_size, found_off));
191
192        if found_size != size {
193            self.tree.insert(tree_node);
194            self.free_tree.insert(free_tree_node);
195        }
196
197        // Start detecting spammers once we have less than 20%
198        // of async space left (which is less than 10% of total
199        // buffer size).
200        //
201        // (This will short-circuit, so `low_oneway_space` is
202        // only called when necessary.)
203        let oneway_spam_detected =
204            is_oneway && new_oneway_space < self.size / 10 && self.low_oneway_space(pid);
205
206        Ok((found_off, oneway_spam_detected))
207    }
208
209    pub(crate) fn reservation_abort(&mut self, offset: usize) -> Result<FreedRange> {
210        let mut cursor = self.tree.cursor_lower_bound_mut(&offset).ok_or_else(|| {
211            pr_warn!(
212                "EINVAL from range_alloc.reservation_abort - offset: {}",
213                offset
214            );
215            EINVAL
216        })?;
217
218        let (_, desc) = cursor.current_mut();
219
220        if desc.offset != offset {
221            pr_warn!(
222                "EINVAL from range_alloc.reservation_abort - offset: {}",
223                offset
224            );
225            return Err(EINVAL);
226        }
227
228        let (reservation, free_node_res) = desc.try_change_state(|state| match state {
229            Some((DescriptorState::Reserved(reservation), free_node_res)) => {
230                (None, Ok((reservation, free_node_res)))
231            }
232            None => {
233                pr_warn!(
234                    "EINVAL from range_alloc.reservation_abort - offset: {}",
235                    offset
236                );
237                (None, Err(EINVAL))
238            }
239            allocated => {
240                pr_warn!(
241                    "EPERM from range_alloc.reservation_abort - offset: {}",
242                    offset
243                );
244                (allocated, Err(EPERM))
245            }
246        })?;
247
248        let mut size = desc.size;
249        let mut offset = desc.offset;
250        let free_oneway_space_add = if reservation.is_oneway { size } else { 0 };
251
252        self.free_oneway_space += free_oneway_space_add;
253
254        let mut freed_range = FreedRange::interior_pages(offset, size);
255        // Compute how large the next free region needs to be to include one more page in
256        // the newly freed range.
257        let add_next_page_needed = match (offset + size) % PAGE_SIZE {
258            0 => usize::MAX,
259            unalign => PAGE_SIZE - unalign,
260        };
261        // Compute how large the previous free region needs to be to include one more page
262        // in the newly freed range.
263        let add_prev_page_needed = match offset % PAGE_SIZE {
264            0 => usize::MAX,
265            unalign => unalign,
266        };
267
268        // Merge next into current if next is free
269        let remove_next = match cursor.peek_next() {
270            Some((_, next)) if next.state.is_none() => {
271                if next.size >= add_next_page_needed {
272                    freed_range.end_page_idx += 1;
273                }
274                self.free_tree.remove(&(next.size, next.offset));
275                size += next.size;
276                true
277            }
278            _ => false,
279        };
280
281        if remove_next {
282            let (_, desc) = cursor.current_mut();
283            desc.size = size;
284            cursor.remove_next();
285        }
286
287        // Merge current into prev if prev is free
288        match cursor.peek_prev_mut() {
289            Some((_, prev)) if prev.state.is_none() => {
290                if prev.size >= add_prev_page_needed {
291                    freed_range.start_page_idx -= 1;
292                }
293                // merge previous with current, remove current
294                self.free_tree.remove(&(prev.size, prev.offset));
295                offset = prev.offset;
296                size += prev.size;
297                prev.size = size;
298                cursor.remove_current();
299            }
300            _ => {}
301        };
302
303        self.free_tree
304            .insert(free_node_res.into_node((size, offset), ()));
305
306        Ok(freed_range)
307    }
308
309    pub(crate) fn reservation_commit(&mut self, offset: usize, data: &mut Option<T>) -> Result {
310        let desc = self.tree.get_mut(&offset).ok_or(ENOENT)?;
311
312        desc.try_change_state(|state| match state {
313            Some((DescriptorState::Reserved(reservation), free_node_res)) => (
314                Some((
315                    DescriptorState::Allocated(reservation.allocate(data.take())),
316                    free_node_res,
317                )),
318                Ok(()),
319            ),
320            other => (other, Err(ENOENT)),
321        })
322    }
323
324    /// Takes an entry at the given offset from [`DescriptorState::Allocated`] to
325    /// [`DescriptorState::Reserved`].
326    ///
327    /// Returns the size of the existing entry and the data associated with it.
328    pub(crate) fn reserve_existing(&mut self, offset: usize) -> Result<(usize, usize, Option<T>)> {
329        let desc = self.tree.get_mut(&offset).ok_or_else(|| {
330            pr_warn!(
331                "ENOENT from range_alloc.reserve_existing - offset: {}",
332                offset
333            );
334            ENOENT
335        })?;
336
337        let (debug_id, data) = desc.try_change_state(|state| match state {
338            Some((DescriptorState::Allocated(allocation), free_node_res)) => {
339                let (reservation, data) = allocation.deallocate();
340                let debug_id = reservation.debug_id;
341                (
342                    Some((DescriptorState::Reserved(reservation), free_node_res)),
343                    Ok((debug_id, data)),
344                )
345            }
346            other => {
347                pr_warn!(
348                    "ENOENT from range_alloc.reserve_existing - offset: {}",
349                    offset
350                );
351                (other, Err(ENOENT))
352            }
353        })?;
354
355        Ok((desc.size, debug_id, data))
356    }
357
358    /// Call the provided callback at every allocated region.
359    ///
360    /// This destroys the range allocator. Used only during shutdown.
361    pub(crate) fn take_for_each<F: Fn(usize, usize, usize, Option<T>)>(&mut self, callback: F) {
362        for (_, desc) in self.tree.iter_mut() {
363            if let Some((DescriptorState::Allocated(allocation), _)) = &mut desc.state {
364                callback(
365                    desc.offset,
366                    desc.size,
367                    allocation.debug_id(),
368                    allocation.take(),
369                );
370            }
371        }
372    }
373
374    /// Find the amount and size of buffers allocated by the current caller.
375    ///
376    /// The idea is that once we cross the threshold, whoever is responsible
377    /// for the low async space is likely to try to send another async transaction,
378    /// and at some point we'll catch them in the act.  This is more efficient
379    /// than keeping a map per pid.
380    fn low_oneway_space(&self, calling_pid: Pid) -> bool {
381        let mut total_alloc_size = 0;
382        let mut num_buffers = 0;
383        for (_, desc) in self.tree.iter() {
384            if let Some((state, _)) = &desc.state {
385                if state.is_oneway() && state.pid() == calling_pid {
386                    total_alloc_size += desc.size;
387                    num_buffers += 1;
388                }
389            }
390        }
391
392        // Warn if this pid has more than 50 transactions, or more than 50% of
393        // async space (which is 25% of total buffer size). Oneway spam is only
394        // detected when the threshold is exceeded.
395        num_buffers > 50 || total_alloc_size > self.size / 4
396    }
397}
398
399type TreeDescriptorState<T> = (DescriptorState<T>, FreeNodeRes);
400struct Descriptor<T> {
401    size: usize,
402    offset: usize,
403    state: Option<TreeDescriptorState<T>>,
404}
405
406impl<T> Descriptor<T> {
407    fn new(offset: usize, size: usize) -> Self {
408        Self {
409            size,
410            offset,
411            state: None,
412        }
413    }
414
415    fn try_change_state<F, Data>(&mut self, f: F) -> Result<Data>
416    where
417        F: FnOnce(Option<TreeDescriptorState<T>>) -> (Option<TreeDescriptorState<T>>, Result<Data>),
418    {
419        let (new_state, result) = f(self.state.take());
420        self.state = new_state;
421        result
422    }
423}
424
425// (Descriptor.size, Descriptor.offset)
426type FreeKey = (usize, usize);
427type FreeNodeRes = RBTreeNodeReservation<FreeKey, ()>;
428
429/// An allocation for use by `reserve_new`.
430pub(crate) struct ReserveNewTreeAlloc<T> {
431    tree_node_res: RBTreeNodeReservation<usize, Descriptor<T>>,
432    free_tree_node_res: FreeNodeRes,
433    desc_node_res: FreeNodeRes,
434}
435
436impl<T> ReserveNewTreeAlloc<T> {
437    pub(crate) fn try_new() -> Result<Self> {
438        let tree_node_res = RBTreeNodeReservation::new(GFP_KERNEL)?;
439        let free_tree_node_res = RBTreeNodeReservation::new(GFP_KERNEL)?;
440        let desc_node_res = RBTreeNodeReservation::new(GFP_KERNEL)?;
441        Ok(Self {
442            tree_node_res,
443            free_tree_node_res,
444            desc_node_res,
445        })
446    }
447
448    fn initialize(
449        self,
450        desc: Descriptor<T>,
451    ) -> (
452        RBTreeNode<usize, Descriptor<T>>,
453        RBTreeNode<FreeKey, ()>,
454        FreeNodeRes,
455    ) {
456        let size = desc.size;
457        let offset = desc.offset;
458        (
459            self.tree_node_res.into_node(offset, desc),
460            self.free_tree_node_res.into_node((size, offset), ()),
461            self.desc_node_res,
462        )
463    }
464}
465
466/// An allocation for creating a tree from an `ArrayRangeAllocator`.
467pub(crate) struct FromArrayAllocs<T> {
468    tree: KVec<RBTreeNodeReservation<usize, Descriptor<T>>>,
469    free_tree: KVec<RBTreeNodeReservation<FreeKey, ()>>,
470}
471
472impl<T> FromArrayAllocs<T> {
473    pub(crate) fn try_new(len: usize) -> Result<Self> {
474        let num_descriptors = 2 * len + 1;
475
476        let mut tree = KVec::with_capacity(num_descriptors, GFP_KERNEL)?;
477        for _ in 0..num_descriptors {
478            tree.push(RBTreeNodeReservation::new(GFP_KERNEL)?, GFP_KERNEL)?;
479        }
480
481        let mut free_tree = KVec::with_capacity(num_descriptors, GFP_KERNEL)?;
482        for _ in 0..num_descriptors {
483            free_tree.push(RBTreeNodeReservation::new(GFP_KERNEL)?, GFP_KERNEL)?;
484        }
485
486        Ok(Self { tree, free_tree })
487    }
488}