vm_allocator/
id_allocator.rs

1// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// Copyright © 2019 Intel Corporation. All Rights Reserved.
3// SPDX-License-Identifier: Apache-2.0 OR BSD-3-Clause
4
5//! Provides allocation and releasing strategy for IDs.
6//!
7//! This module implements an allocation strategies for all resource types
8//! that can be abstracted to an integer.
9
10use crate::{Error, Result};
11use alloc::collections::BTreeSet;
12
13/// An unique ID allocator that allows management of IDs in a given interval.
14// Internal representation of IdAllocator. Contains the ends of the interval
15// that is managed, a field that points at the next available ID, and a
16// BTreeSet used to store the released IDs. The reason we chose a
17// BTreeSet is that the average complexity for deletion and insertion is
18// O(logN) compared to Vec for example, another benefit is that the entries
19// are sorted so we will always use the first available ID.
20#[derive(Debug, Clone)]
21#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
22pub struct IdAllocator {
23    // Beginning of the range of IDs that we want to manage.
24    range_base: u32,
25    // First available id that was never allocated.
26    next_id: Option<u32>,
27    // End of the range of IDs that we want to manage.
28    range_end: u32,
29    // Set of all freed ids that can be reused at subsequent allocations.
30    freed_ids: BTreeSet<u32>,
31}
32
33impl IdAllocator {
34    /// Creates a new instance of IdAllocator that will be used to manage the
35    /// allocation and release of ids from the interval specified by
36    /// `range_base` and `range_end`
37    pub fn new(range_base: u32, range_end: u32) -> Result<Self> {
38        if range_end < range_base {
39            return Err(Error::InvalidRange(range_base.into(), range_end.into()));
40        }
41        Ok(IdAllocator {
42            range_base,
43            next_id: Some(range_base),
44            range_end,
45            freed_ids: BTreeSet::<u32>::new(),
46        })
47    }
48
49    fn id_in_range(&self, id: u32) -> bool {
50        // Check for out of range.
51        self.range_base <= id && id <= self.range_end
52    }
53
54    /// Allocate an ID from the managed range.
55    /// We first try to reuse one of the IDs released before. If there is no
56    /// ID to reuse we return the next available one from the managed range.
57    pub fn allocate_id(&mut self) -> Result<u32> {
58        // If the set containing all freed ids is not empty we extract the
59        // first entry from that set and return it.
60        if !self.freed_ids.is_empty() {
61            let ret_value = *self.freed_ids.iter().next().unwrap();
62            self.freed_ids.remove(&ret_value);
63            return Ok(ret_value);
64        }
65        // If no id was freed before we return the next available id.
66        if let Some(next_id) = self.next_id {
67            if next_id > self.range_end {
68                return Err(Error::ResourceNotAvailable);
69            }
70            // Prepare the next available id. If the addition overflows we
71            // set the next_id field to None and return Overflow at the next
72            // allocation request.
73            self.next_id = next_id.checked_add(1);
74            return Ok(next_id);
75        }
76        Err(Error::Overflow)
77    }
78
79    /// Returns `true` if `id` is currently allocated from the managed range.
80    ///
81    /// An id is considered allocated if it has been handed out by
82    /// [`allocate_id`](Self::allocate_id) and not since released by
83    /// [`free_id`](Self::free_id). Ids outside the managed range are never
84    /// allocated.
85    pub fn is_allocated(&self, id: u32) -> bool {
86        if !self.id_in_range(id) || self.freed_ids.contains(&id) {
87            return false;
88        }
89        // An id is allocated iff it has already been handed out. `next_id` is
90        // the next id to hand out; `None` means the whole range has been
91        // handed out.
92        match self.next_id {
93            Some(next) => id < next,
94            None => true,
95        }
96    }
97
98    /// Frees an id from the managed range.
99    pub fn free_id(&mut self, id: u32) -> Result<u32> {
100        // Check if the id belongs to the managed range and if it was not
101        // released before. Return error if any of the conditions is not met.
102        if !self.id_in_range(id) {
103            return Err(Error::OutOfRange(id));
104        }
105        if let Some(next_id) = self.next_id {
106            if next_id <= id {
107                return Err(Error::NeverAllocated(id));
108            }
109        }
110
111        // Insert the released id in the set of released id to avoid releasing
112        // it in next iterations.
113        self.freed_ids
114            .insert(id)
115            .then_some(id)
116            .ok_or(Error::AlreadyReleased(id))
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use crate::id_allocator::IdAllocator;
123    use crate::Error;
124
125    #[test]
126    fn test_slot_id_allocation() {
127        let faulty_allocator = IdAllocator::new(23, 5);
128        assert_eq!(faulty_allocator.unwrap_err(), Error::InvalidRange(23, 5));
129        let mut legacy_irq_allocator = IdAllocator::new(5, 23).unwrap();
130        assert_eq!(legacy_irq_allocator.range_base, 5);
131        assert_eq!(legacy_irq_allocator.range_end, 23);
132
133        let id = legacy_irq_allocator.allocate_id().unwrap();
134        assert_eq!(id, 5);
135        assert_eq!(legacy_irq_allocator.next_id.unwrap(), 6);
136
137        for _ in 1..19 {
138            assert!(legacy_irq_allocator.allocate_id().is_ok());
139        }
140
141        assert_eq!(
142            legacy_irq_allocator.allocate_id().unwrap_err(),
143            Error::ResourceNotAvailable
144        );
145    }
146
147    #[test]
148    fn test_u32_overflow() {
149        let mut allocator = IdAllocator::new(u32::MAX - 1, u32::MAX).unwrap();
150        assert_eq!(allocator.allocate_id().unwrap(), u32::MAX - 1);
151        assert_eq!(allocator.allocate_id().unwrap(), u32::MAX);
152        let res = allocator.allocate_id();
153        assert!(res.is_err());
154        assert_eq!(res.unwrap_err(), Error::Overflow);
155    }
156
157    #[test]
158    fn test_slot_id_free() {
159        let mut legacy_irq_allocator = IdAllocator::new(5, 23).unwrap();
160        assert_eq!(
161            legacy_irq_allocator.free_id(3).unwrap_err(),
162            Error::OutOfRange(3)
163        );
164        assert_eq!(legacy_irq_allocator.freed_ids.len(), 0);
165
166        for _ in 1..10 {
167            let _id = legacy_irq_allocator.allocate_id().unwrap();
168        }
169
170        let irq = 10;
171        legacy_irq_allocator.free_id(irq).unwrap();
172        assert!(legacy_irq_allocator.freed_ids.contains(&irq));
173        assert_eq!(
174            legacy_irq_allocator.free_id(10).unwrap_err(),
175            Error::AlreadyReleased(10)
176        );
177        let irq = 9;
178        legacy_irq_allocator.free_id(irq).unwrap();
179        assert_eq!(legacy_irq_allocator.freed_ids.len(), 2);
180        assert_eq!(*legacy_irq_allocator.freed_ids.iter().next().unwrap(), 9);
181
182        let irq = legacy_irq_allocator.allocate_id().unwrap();
183        assert_eq!(irq, 9);
184        assert!(!legacy_irq_allocator.freed_ids.contains(&irq));
185        assert_eq!(legacy_irq_allocator.freed_ids.len(), 1);
186        assert_eq!(
187            legacy_irq_allocator.free_id(21).unwrap_err(),
188            Error::NeverAllocated(21)
189        );
190    }
191
192    #[test]
193    fn test_free_id_never_allocated_boundary() {
194        let mut allocator = IdAllocator::new(5, 23).unwrap();
195
196        assert_eq!(allocator.next_id.unwrap(), 5);
197        assert_eq!(allocator.free_id(5).unwrap_err(), Error::NeverAllocated(5));
198        assert_eq!(allocator.freed_ids.len(), 0);
199
200        for _ in 0..3 {
201            allocator.allocate_id().unwrap();
202        }
203        assert_eq!(allocator.next_id.unwrap(), 8);
204        assert_eq!(allocator.free_id(8).unwrap_err(), Error::NeverAllocated(8));
205        assert_eq!(allocator.freed_ids.len(), 0);
206    }
207
208    #[test]
209    fn test_is_allocated() {
210        let mut allocator = IdAllocator::new(5, 23).unwrap();
211
212        assert!(!allocator.is_allocated(4));
213        assert!(!allocator.is_allocated(24));
214
215        assert!(!allocator.is_allocated(10));
216
217        let id = allocator.allocate_id().unwrap();
218        assert_eq!(id, 5);
219        assert!(allocator.is_allocated(5));
220        assert!(!allocator.is_allocated(6));
221
222        allocator.free_id(5).unwrap();
223        assert!(!allocator.is_allocated(5));
224
225        let id = allocator.allocate_id().unwrap();
226        assert_eq!(id, 5);
227        assert!(allocator.is_allocated(5));
228    }
229
230    #[test]
231    fn test_is_allocated_full_range() {
232        let mut allocator = IdAllocator::new(u32::MAX - 1, u32::MAX).unwrap();
233
234        assert!(!allocator.is_allocated(u32::MAX - 1));
235        assert!(!allocator.is_allocated(u32::MAX));
236
237        assert_eq!(allocator.allocate_id().unwrap(), u32::MAX - 1);
238        assert!(allocator.is_allocated(u32::MAX - 1));
239        assert!(!allocator.is_allocated(u32::MAX));
240
241        assert_eq!(allocator.allocate_id().unwrap(), u32::MAX);
242        assert!(allocator.next_id.is_none());
243        assert!(allocator.is_allocated(u32::MAX));
244
245        allocator.free_id(u32::MAX).unwrap();
246        assert!(!allocator.is_allocated(u32::MAX));
247        assert!(allocator.is_allocated(u32::MAX - 1));
248    }
249
250    #[test]
251    fn test_is_allocated_with_freed() {
252        let mut allocator = IdAllocator::new(5, 23).unwrap();
253        allocator.allocate_id().unwrap();
254        allocator.allocate_id().unwrap();
255
256        assert_eq!(allocator.next_id, Some(7));
257
258        assert!(allocator.is_allocated(6));
259        assert!(!allocator.is_allocated(7));
260
261        allocator.free_id(5).unwrap();
262        assert!(!allocator.is_allocated(5));
263        assert!(allocator.is_allocated(6));
264        assert!(!allocator.is_allocated(7));
265    }
266
267    #[test]
268    fn test_id_sanity_checks() {
269        let legacy_irq_allocator = IdAllocator::new(5, 23).unwrap();
270
271        assert!(!legacy_irq_allocator.id_in_range(4));
272        assert!(legacy_irq_allocator.id_in_range(6));
273        assert!(!legacy_irq_allocator.id_in_range(25));
274    }
275}
vm_allocator/id_allocator.rs

vm_allocator/
id_allocator.rs
