1 // Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2 // SPDX-License-Identifier: Apache-2.0 OR BSD-3-Clause
3 
4 //! Bitmap backend implementation based on atomic integers.
5 
6 use std::sync::atomic::{AtomicU64, Ordering};
7 
8 use crate::bitmap::{Bitmap, RefSlice, WithBitmapSlice};
9 
10 #[cfg(feature = "backend-mmap")]
11 use crate::mmap::NewBitmap;
12 
13 /// `AtomicBitmap` implements a simple bit map on the page level with test and set operations.
14 /// It is page-size aware, so it converts addresses to page numbers before setting or clearing
15 /// the bits.
16 #[derive(Debug)]
17 pub struct AtomicBitmap {
18     map: Vec<AtomicU64>,
19     size: usize,
20     page_size: usize,
21 }
22 
23 #[allow(clippy::len_without_is_empty)]
24 impl AtomicBitmap {
25     /// Create a new bitmap of `byte_size`, with one bit per page. This is effectively
26     /// rounded up, and we get a new vector of the next multiple of 64 bigger than `bit_size`.
new(byte_size: usize, page_size: usize) -> Self27     pub fn new(byte_size: usize, page_size: usize) -> Self {
28         let mut num_pages = byte_size / page_size;
29         if byte_size % page_size > 0 {
30             num_pages += 1;
31         }
32 
33         // Adding one entry element more just in case `num_pages` is not a multiple of `64`.
34         let map_size = num_pages / 64 + 1;
35         let map: Vec<AtomicU64> = (0..map_size).map(|_| AtomicU64::new(0)).collect();
36 
37         AtomicBitmap {
38             map,
39             size: num_pages,
40             page_size,
41         }
42     }
43 
44     /// Is bit `n` set? Bits outside the range of the bitmap are always unset.
is_bit_set(&self, index: usize) -> bool45     pub fn is_bit_set(&self, index: usize) -> bool {
46         if index < self.size {
47             (self.map[index >> 6].load(Ordering::Acquire) & (1 << (index & 63))) != 0
48         } else {
49             // Out-of-range bits are always unset.
50             false
51         }
52     }
53 
54     /// Is the bit corresponding to address `addr` set?
is_addr_set(&self, addr: usize) -> bool55     pub fn is_addr_set(&self, addr: usize) -> bool {
56         self.is_bit_set(addr / self.page_size)
57     }
58 
59     /// Set a range of `len` bytes starting at `start_addr`. The first bit set in the bitmap
60     /// is for the page corresponding to `start_addr`, and the last bit that we set corresponds
61     /// to address `start_addr + len - 1`.
set_addr_range(&self, start_addr: usize, len: usize)62     pub fn set_addr_range(&self, start_addr: usize, len: usize) {
63         // Return early in the unlikely event that `len == 0` so the `len - 1` computation
64         // below does not underflow.
65         if len == 0 {
66             return;
67         }
68 
69         let first_bit = start_addr / self.page_size;
70         // Handle input ranges where `start_addr + len - 1` would otherwise overflow an `usize`
71         // by ignoring pages at invalid addresses.
72         let last_bit = start_addr.saturating_add(len - 1) / self.page_size;
73         for n in first_bit..=last_bit {
74             if n >= self.size {
75                 // Attempts to set bits beyond the end of the bitmap are simply ignored.
76                 break;
77             }
78             self.map[n >> 6].fetch_or(1 << (n & 63), Ordering::SeqCst);
79         }
80     }
81 
82     /// Get the length of the bitmap in bits (i.e. in how many pages it can represent).
len(&self) -> usize83     pub fn len(&self) -> usize {
84         self.size
85     }
86 
87     /// Atomically get and reset the dirty page bitmap.
get_and_reset(&self) -> Vec<u64>88     pub fn get_and_reset(&self) -> Vec<u64> {
89         self.map
90             .iter()
91             .map(|u| u.fetch_and(0, Ordering::SeqCst))
92             .collect()
93     }
94 
95     /// Reset all bitmap bits to 0.
reset(&self)96     pub fn reset(&self) {
97         for it in self.map.iter() {
98             it.store(0, Ordering::Release);
99         }
100     }
101 }
102 
103 impl Clone for AtomicBitmap {
clone(&self) -> Self104     fn clone(&self) -> Self {
105         let map = self
106             .map
107             .iter()
108             .map(|i| i.load(Ordering::Acquire))
109             .map(AtomicU64::new)
110             .collect();
111         AtomicBitmap {
112             map,
113             size: self.size,
114             page_size: self.page_size,
115         }
116     }
117 }
118 
119 impl<'a> WithBitmapSlice<'a> for AtomicBitmap {
120     type S = RefSlice<'a, Self>;
121 }
122 
123 impl Bitmap for AtomicBitmap {
mark_dirty(&self, offset: usize, len: usize)124     fn mark_dirty(&self, offset: usize, len: usize) {
125         self.set_addr_range(offset, len)
126     }
127 
dirty_at(&self, offset: usize) -> bool128     fn dirty_at(&self, offset: usize) -> bool {
129         self.is_addr_set(offset)
130     }
131 
slice_at(&self, offset: usize) -> <Self as WithBitmapSlice>::S132     fn slice_at(&self, offset: usize) -> <Self as WithBitmapSlice>::S {
133         RefSlice::new(self, offset)
134     }
135 }
136 
137 impl Default for AtomicBitmap {
default() -> Self138     fn default() -> Self {
139         AtomicBitmap::new(0, 0x1000)
140     }
141 }
142 
143 #[cfg(feature = "backend-mmap")]
144 impl NewBitmap for AtomicBitmap {
with_len(len: usize) -> Self145     fn with_len(len: usize) -> Self {
146         let page_size;
147 
148         #[cfg(unix)]
149         {
150             // SAFETY: There's no unsafe potential in calling this function.
151             page_size = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) };
152         }
153 
154         #[cfg(windows)]
155         {
156             use winapi::um::sysinfoapi::{GetSystemInfo, LPSYSTEM_INFO, SYSTEM_INFO};
157 
158             // It's safe to initialize this object from a zeroed memory region.
159             let mut sysinfo: SYSTEM_INFO = unsafe { std::mem::zeroed() };
160 
161             // It's safe to call this method as the pointer is based on the address
162             // of the previously initialized `sysinfo` object.
163             unsafe { GetSystemInfo(&mut sysinfo as LPSYSTEM_INFO) };
164 
165             page_size = sysinfo.dwPageSize;
166         }
167 
168         // The `unwrap` is safe to use because the above call should always succeed on the
169         // supported platforms, and the size of a page will always fit within a `usize`.
170         AtomicBitmap::new(len, usize::try_from(page_size).unwrap())
171     }
172 }
173 
174 #[cfg(test)]
175 mod tests {
176     use super::*;
177 
178     use crate::bitmap::tests::test_bitmap;
179 
180     #[test]
test_bitmap_basic()181     fn test_bitmap_basic() {
182         // Test that bitmap size is properly rounded up.
183         let a = AtomicBitmap::new(1025, 128);
184         assert_eq!(a.len(), 9);
185 
186         let b = AtomicBitmap::new(1024, 128);
187         assert_eq!(b.len(), 8);
188         b.set_addr_range(128, 129);
189         assert!(!b.is_addr_set(0));
190         assert!(b.is_addr_set(128));
191         assert!(b.is_addr_set(256));
192         assert!(!b.is_addr_set(384));
193 
194         #[allow(clippy::redundant_clone)]
195         let copy_b = b.clone();
196         assert!(copy_b.is_addr_set(256));
197         assert!(!copy_b.is_addr_set(384));
198 
199         b.reset();
200         assert!(!b.is_addr_set(128));
201         assert!(!b.is_addr_set(256));
202         assert!(!b.is_addr_set(384));
203 
204         b.set_addr_range(128, 129);
205         let v = b.get_and_reset();
206 
207         assert!(!b.is_addr_set(128));
208         assert!(!b.is_addr_set(256));
209         assert!(!b.is_addr_set(384));
210 
211         assert_eq!(v.len(), 1);
212         assert_eq!(v[0], 0b110);
213     }
214 
215     #[test]
test_bitmap_out_of_range()216     fn test_bitmap_out_of_range() {
217         let b = AtomicBitmap::new(1024, 1);
218         // Set a partial range that goes beyond the end of the bitmap
219         b.set_addr_range(768, 512);
220         assert!(b.is_addr_set(768));
221         // The bitmap is never set beyond its end.
222         assert!(!b.is_addr_set(1024));
223         assert!(!b.is_addr_set(1152));
224     }
225 
226     #[test]
test_bitmap_impl()227     fn test_bitmap_impl() {
228         let b = AtomicBitmap::new(0x2000, 128);
229         test_bitmap(&b);
230     }
231 }
232