1 //! Index/slot allocator policies for the pooling allocator. 2 3 use crate::hash_map::{Entry, HashMap}; 4 use crate::prelude::*; 5 use crate::runtime::vm::CompiledModuleId; 6 use std::mem; 7 use std::sync::Mutex; 8 use wasmtime_environ::DefinedMemoryIndex; 9 10 /// A slot index. 11 #[derive(Hash, Clone, Copy, Debug, PartialEq, Eq)] 12 pub struct SlotId(pub u32); 13 14 impl SlotId { 15 /// The index of this slot. index(self) -> usize16 pub fn index(self) -> usize { 17 self.0 as usize 18 } 19 } 20 21 /// A simple index allocator. 22 /// 23 /// This index allocator doesn't do any module affinity or anything like that, 24 /// however it is built on top of the `ModuleAffinityIndexAllocator` to save 25 /// code (and code size). 26 #[derive(Debug)] 27 pub struct SimpleIndexAllocator(ModuleAffinityIndexAllocator); 28 29 impl SimpleIndexAllocator { new(capacity: u32) -> Self30 pub fn new(capacity: u32) -> Self { 31 SimpleIndexAllocator(ModuleAffinityIndexAllocator::new(capacity, 0)) 32 } 33 is_empty(&self) -> bool34 pub fn is_empty(&self) -> bool { 35 self.0.is_empty() 36 } 37 alloc(&self) -> Option<SlotId>38 pub fn alloc(&self) -> Option<SlotId> { 39 self.0.alloc(None) 40 } 41 42 /// Frees the `index` slot to be available for allocation elsewhere. 43 /// 44 /// The `bytes_resident` argument is a counter to keep track of how many 45 /// bytse are still resident in this slot, if any, for reporting later via 46 /// the [`Self::unused_bytes_resident`] method. free(&self, index: SlotId, bytes_resident: usize)47 pub(crate) fn free(&self, index: SlotId, bytes_resident: usize) { 48 self.0.free(index, bytes_resident); 49 } 50 51 /// Returns the number of previously-used slots in this allocator which are 52 /// not currently in use. 53 /// 54 /// Note that this acquires a `Mutex` for synchronization at this time to 55 /// read the internal counter information. unused_warm_slots(&self) -> u3256 pub fn unused_warm_slots(&self) -> u32 { 57 self.0.unused_warm_slots() 58 } 59 60 /// Returns the number of bytes that are resident in previously-used slots 61 /// in this allocator which are not currently in use. 62 /// 63 /// Note that this acquires a `Mutex` for synchronization at this time to 64 /// read the internal counter information. unused_bytes_resident(&self) -> usize65 pub fn unused_bytes_resident(&self) -> usize { 66 self.0.unused_bytes_resident() 67 } 68 69 #[cfg(test)] testing_freelist(&self) -> Vec<SlotId>70 pub(crate) fn testing_freelist(&self) -> Vec<SlotId> { 71 self.0.testing_freelist() 72 } 73 } 74 75 /// A particular defined memory within a particular module. 76 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] 77 pub struct MemoryInModule(pub CompiledModuleId, pub DefinedMemoryIndex); 78 79 /// An index allocator that has configurable affinity between slots and modules 80 /// so that slots are often reused for the same module again. 81 #[derive(Debug)] 82 pub struct ModuleAffinityIndexAllocator(Mutex<Inner>); 83 84 #[derive(Debug)] 85 struct Inner { 86 /// Maximum number of "unused warm slots" which will be allowed during 87 /// allocation. 88 /// 89 /// This is a user-configurable knob which can be used to influence the 90 /// maximum number of unused slots at any one point in time. A "warm slot" 91 /// is one that's considered having been previously allocated. 92 max_unused_warm_slots: u32, 93 94 /// Current count of "warm slots", or those that were previously allocated 95 /// which are now no longer in use. 96 /// 97 /// This is the size of the `warm` list. 98 unused_warm_slots: u32, 99 100 /// A linked list (via indices) which enumerates all "warm and unused" 101 /// slots, or those which have previously been allocated and then free'd. 102 warm: List, 103 104 /// Last slot that was allocated for the first time ever. 105 /// 106 /// This is initially 0 and is incremented during `pick_cold`. If this 107 /// matches `max_cold`, there are no more cold slots left. 108 last_cold: u32, 109 110 /// The state of any given slot. 111 /// 112 /// Records indices in the above list (empty) or two lists (with affinity), 113 /// and these indices are kept up-to-date to allow fast removal. 114 slot_state: Vec<SlotState>, 115 116 /// Affine slot management which tracks which slots are free and were last 117 /// used with the specified `CompiledModuleId`. 118 /// 119 /// The `List` here is appended to during deallocation and removal happens 120 /// from the tail during allocation. 121 module_affine: HashMap<MemoryInModule, List>, 122 123 /// Cache for the sum of the `bytes_resident` of all `UnusedWarm` slots. 124 unused_bytes_resident: usize, 125 } 126 127 /// A helper "linked list" data structure which is based on indices. 128 #[derive(Default, Debug)] 129 struct List { 130 head: Option<SlotId>, 131 tail: Option<SlotId>, 132 } 133 134 /// A helper data structure for an intrusive linked list, coupled with the 135 /// `List` type. 136 #[derive(Default, Debug, Copy, Clone)] 137 struct Link { 138 prev: Option<SlotId>, 139 next: Option<SlotId>, 140 } 141 142 #[derive(Clone, Debug)] 143 enum SlotState { 144 /// This slot is currently in use and is affine to the specified module's memory. 145 Used(Option<MemoryInModule>), 146 147 /// This slot is not currently used, and has never been used. 148 UnusedCold, 149 150 /// This slot is not currently used, but was previously allocated. 151 /// 152 /// The payload here is metadata about the lists that this slot is contained 153 /// within. 154 UnusedWarm(Unused), 155 } 156 157 impl SlotState { unwrap_unused(&mut self) -> &mut Unused158 fn unwrap_unused(&mut self) -> &mut Unused { 159 match self { 160 SlotState::UnusedWarm(u) => u, 161 _ => unreachable!(), 162 } 163 } 164 } 165 166 #[derive(Default, Copy, Clone, Debug)] 167 struct Unused { 168 /// Which module this slot was historically affine to, if any. 169 affinity: Option<MemoryInModule>, 170 171 /// Number of bytes that are part of `UnusedWarm` slots and are currently 172 /// kept resident (vs paged out). 173 bytes_resident: usize, 174 175 /// Metadata about the linked list for all slots affine to `affinity`. 176 affine_list_link: Link, 177 178 /// Metadata within the `warm` list of the main allocator. 179 unused_list_link: Link, 180 } 181 182 enum AllocMode { 183 ForceAffineAndClear, 184 AnySlot, 185 } 186 187 impl ModuleAffinityIndexAllocator { 188 /// Create the default state for this strategy. new(capacity: u32, max_unused_warm_slots: u32) -> Self189 pub fn new(capacity: u32, max_unused_warm_slots: u32) -> Self { 190 ModuleAffinityIndexAllocator(Mutex::new(Inner { 191 last_cold: 0, 192 max_unused_warm_slots, 193 unused_warm_slots: 0, 194 module_affine: HashMap::new(), 195 slot_state: (0..capacity).map(|_| SlotState::UnusedCold).collect(), 196 warm: List::default(), 197 unused_bytes_resident: 0, 198 })) 199 } 200 201 /// How many slots can this allocator allocate? len(&self) -> usize202 pub fn len(&self) -> usize { 203 let inner = self.0.lock().unwrap(); 204 inner.slot_state.len() 205 } 206 207 /// Are zero slots in use right now? is_empty(&self) -> bool208 pub fn is_empty(&self) -> bool { 209 let inner = self.0.lock().unwrap(); 210 !inner 211 .slot_state 212 .iter() 213 .any(|s| matches!(s, SlotState::Used(_))) 214 } 215 216 /// Allocate a new index from this allocator optionally using `id` as an 217 /// affinity request if the allocation strategy supports it. 218 /// 219 /// Returns `None` if no more slots are available. alloc(&self, for_memory: Option<MemoryInModule>) -> Option<SlotId>220 pub fn alloc(&self, for_memory: Option<MemoryInModule>) -> Option<SlotId> { 221 self._alloc(for_memory, AllocMode::AnySlot) 222 } 223 224 /// Attempts to allocate a guaranteed-affine slot to the module `id` 225 /// specified. 226 /// 227 /// Returns `None` if there are no slots affine to `id`. The allocation of 228 /// this slot will not record the affinity to `id`, instead simply listing 229 /// it as taken. This is intended to be used for clearing out all affine 230 /// slots to a module. alloc_affine_and_clear_affinity( &self, module_id: CompiledModuleId, memory_index: DefinedMemoryIndex, ) -> Option<SlotId>231 pub fn alloc_affine_and_clear_affinity( 232 &self, 233 module_id: CompiledModuleId, 234 memory_index: DefinedMemoryIndex, 235 ) -> Option<SlotId> { 236 self._alloc( 237 Some(MemoryInModule(module_id, memory_index)), 238 AllocMode::ForceAffineAndClear, 239 ) 240 } 241 _alloc(&self, for_memory: Option<MemoryInModule>, mode: AllocMode) -> Option<SlotId>242 fn _alloc(&self, for_memory: Option<MemoryInModule>, mode: AllocMode) -> Option<SlotId> { 243 let mut inner = self.0.lock().unwrap(); 244 let inner = &mut *inner; 245 246 // As a first-pass always attempt an affine allocation. This will 247 // succeed if any slots are considered affine to `module_id` (if it's 248 // specified). Failing that something else is attempted to be chosen. 249 let slot_id = inner.pick_affine(for_memory).or_else(|| { 250 match mode { 251 // If any slot is requested then this is a normal instantiation 252 // looking for an index. Without any affine candidates there are 253 // two options here: 254 // 255 // 1. Pick a slot amongst previously allocated slots 256 // 2. Pick a slot that's never been used before 257 // 258 // The choice here is guided by the initial configuration of 259 // `max_unused_warm_slots`. If our unused warm slots, which are 260 // likely all affine, is below this threshold then the affinity 261 // of the warm slots isn't tampered with and first a cold slot 262 // is chosen. If the cold slot allocation fails, however, a warm 263 // slot is evicted. 264 // 265 // The opposite happens when we're above our threshold for the 266 // maximum number of warm slots, meaning that a warm slot is 267 // attempted to be picked from first with a cold slot following 268 // that. Note that the warm slot allocation in this case should 269 // only fail of `max_unused_warm_slots` is 0, otherwise 270 // `pick_warm` will always succeed. 271 AllocMode::AnySlot => { 272 if inner.unused_warm_slots < inner.max_unused_warm_slots { 273 inner.pick_cold().or_else(|| inner.pick_warm()) 274 } else { 275 inner.pick_warm().or_else(|| { 276 debug_assert!(inner.max_unused_warm_slots == 0); 277 inner.pick_cold() 278 }) 279 } 280 } 281 282 // In this mode an affinity-based allocation is always performed 283 // as the purpose here is to clear out slots relevant to 284 // `module_id` during module teardown. This means that there's 285 // no consulting non-affine slots in this path. 286 AllocMode::ForceAffineAndClear => None, 287 } 288 })?; 289 290 let slot = &mut inner.slot_state[slot_id.index()]; 291 if let SlotState::UnusedWarm(Unused { bytes_resident, .. }) = slot { 292 inner.unused_bytes_resident -= *bytes_resident; 293 } 294 *slot = SlotState::Used(match mode { 295 AllocMode::ForceAffineAndClear => None, 296 AllocMode::AnySlot => for_memory, 297 }); 298 299 Some(slot_id) 300 } 301 free(&self, index: SlotId, bytes_resident: usize)302 pub(crate) fn free(&self, index: SlotId, bytes_resident: usize) { 303 let mut inner = self.0.lock().unwrap(); 304 let inner = &mut *inner; 305 let module_memory = match inner.slot_state[index.index()] { 306 SlotState::Used(module_memory) => module_memory, 307 _ => unreachable!(), 308 }; 309 310 // Bump the number of warm slots since this slot is now considered 311 // previously used. Afterwards append it to the linked list of all 312 // unused and warm slots. 313 inner.unused_warm_slots += 1; 314 let unused_list_link = inner 315 .warm 316 .append(index, &mut inner.slot_state, |s| &mut s.unused_list_link); 317 318 let affine_list_link = match module_memory { 319 // If this slot is affine to a particular module then append this 320 // index to the linked list for the affine module. Otherwise insert 321 // a new one-element linked list. 322 Some(module) => match inner.module_affine.entry(module) { 323 Entry::Occupied(mut e) => e 324 .get_mut() 325 .append(index, &mut inner.slot_state, |s| &mut s.affine_list_link), 326 Entry::Vacant(v) => { 327 v.insert(List::new(index)); 328 Link::default() 329 } 330 }, 331 332 // If this slot has no affinity then the affine link is empty. 333 None => Link::default(), 334 }; 335 336 inner.unused_bytes_resident += bytes_resident; 337 inner.slot_state[index.index()] = SlotState::UnusedWarm(Unused { 338 affinity: module_memory, 339 bytes_resident, 340 affine_list_link, 341 unused_list_link, 342 }); 343 } 344 345 /// Return the number of empty slots available in this allocator. 346 #[cfg(test)] num_empty_slots(&self) -> usize347 pub fn num_empty_slots(&self) -> usize { 348 let inner = self.0.lock().unwrap(); 349 let total_slots = inner.slot_state.len(); 350 (total_slots - inner.last_cold as usize) + inner.unused_warm_slots as usize 351 } 352 353 /// For testing only, we want to be able to assert what is on the single 354 /// freelist, for the policies that keep just one. 355 #[cfg(test)] testing_freelist(&self) -> Vec<SlotId>356 pub(crate) fn testing_freelist(&self) -> Vec<SlotId> { 357 let inner = self.0.lock().unwrap(); 358 inner 359 .warm 360 .iter(&inner.slot_state, |s| &s.unused_list_link) 361 .collect() 362 } 363 364 /// For testing only, get the list of all modules with at least one slot 365 /// with affinity for that module. 366 #[cfg(test)] testing_module_affinity_list(&self) -> Vec<MemoryInModule>367 pub(crate) fn testing_module_affinity_list(&self) -> Vec<MemoryInModule> { 368 let inner = self.0.lock().unwrap(); 369 inner.module_affine.keys().copied().collect() 370 } 371 372 /// Returns the number of previously-used slots in this allocator which are 373 /// not currently in use. 374 /// 375 /// Note that this acquires a `Mutex` for synchronization at this time to 376 /// read the internal counter information. unused_warm_slots(&self) -> u32377 pub fn unused_warm_slots(&self) -> u32 { 378 self.0.lock().unwrap().unused_warm_slots 379 } 380 381 /// Returns the number of bytes that are resident in previously-used slots 382 /// in this allocator which are not currently in use. 383 /// 384 /// Note that this acquires a `Mutex` for synchronization at this time to 385 /// read the internal counter information. unused_bytes_resident(&self) -> usize386 pub fn unused_bytes_resident(&self) -> usize { 387 self.0.lock().unwrap().unused_bytes_resident 388 } 389 } 390 391 impl Inner { 392 /// Attempts to allocate a slot already affine to `id`, returning `None` if 393 /// `id` is `None` or if there are no affine slots. pick_affine(&mut self, for_memory: Option<MemoryInModule>) -> Option<SlotId>394 fn pick_affine(&mut self, for_memory: Option<MemoryInModule>) -> Option<SlotId> { 395 // Note that the `tail` is chosen here of the affine list as it's the 396 // most recently used, which for affine allocations is what we want -- 397 // maximizing temporal reuse. 398 let ret = self.module_affine.get(&for_memory?)?.tail?; 399 self.remove(ret); 400 Some(ret) 401 } 402 pick_warm(&mut self) -> Option<SlotId>403 fn pick_warm(&mut self) -> Option<SlotId> { 404 // Insertions into the `unused` list happen at the `tail`, so the 405 // least-recently-used item will be at the head. That's our goal here, 406 // pick the least-recently-used slot since something "warm" is being 407 // evicted anyway. 408 let head = self.warm.head?; 409 self.remove(head); 410 Some(head) 411 } 412 remove(&mut self, slot: SlotId)413 fn remove(&mut self, slot: SlotId) { 414 // Decrement the size of the warm list, and additionally remove it from 415 // the `warm` linked list. 416 self.unused_warm_slots -= 1; 417 self.warm 418 .remove(slot, &mut self.slot_state, |u| &mut u.unused_list_link); 419 420 // If this slot is affine to a module then additionally remove it from 421 // that module's affinity linked list. Note that if the module's affine 422 // list is empty then the module's entry in the map is completely 423 // removed as well. 424 let module = self.slot_state[slot.index()].unwrap_unused().affinity; 425 if let Some(module) = module { 426 let mut list = match self.module_affine.entry(module) { 427 Entry::Occupied(e) => e, 428 Entry::Vacant(_) => unreachable!(), 429 }; 430 list.get_mut() 431 .remove(slot, &mut self.slot_state, |u| &mut u.affine_list_link); 432 433 if list.get_mut().head.is_none() { 434 list.remove(); 435 } 436 } 437 } 438 pick_cold(&mut self) -> Option<SlotId>439 fn pick_cold(&mut self) -> Option<SlotId> { 440 if (self.last_cold as usize) == self.slot_state.len() { 441 None 442 } else { 443 let ret = Some(SlotId(self.last_cold)); 444 self.last_cold += 1; 445 ret 446 } 447 } 448 } 449 450 impl List { 451 /// Creates a new one-element list pointing at `id`. new(id: SlotId) -> List452 fn new(id: SlotId) -> List { 453 List { 454 head: Some(id), 455 tail: Some(id), 456 } 457 } 458 459 /// Appends the `id` to this list whose links are determined by `link`. append( &mut self, id: SlotId, states: &mut [SlotState], link: fn(&mut Unused) -> &mut Link, ) -> Link460 fn append( 461 &mut self, 462 id: SlotId, 463 states: &mut [SlotState], 464 link: fn(&mut Unused) -> &mut Link, 465 ) -> Link { 466 // This `id` is the new tail... 467 let tail = mem::replace(&mut self.tail, Some(id)); 468 469 // If the tail was present, then update its `next` field to ourselves as 470 // we've been appended, otherwise update the `head` since the list was 471 // previously empty. 472 match tail { 473 Some(tail) => link(states[tail.index()].unwrap_unused()).next = Some(id), 474 None => self.head = Some(id), 475 } 476 Link { 477 prev: tail, 478 next: None, 479 } 480 } 481 482 /// Removes `id` from this list whose links are determined by `link`. remove( &mut self, id: SlotId, slot_state: &mut [SlotState], link: fn(&mut Unused) -> &mut Link, ) -> Unused483 fn remove( 484 &mut self, 485 id: SlotId, 486 slot_state: &mut [SlotState], 487 link: fn(&mut Unused) -> &mut Link, 488 ) -> Unused { 489 let mut state = *slot_state[id.index()].unwrap_unused(); 490 let next = link(&mut state).next; 491 let prev = link(&mut state).prev; 492 493 // If a `next` node is present for this link, then its previous was our 494 // own previous now. Otherwise we are the tail so the new tail is our 495 // previous. 496 match next { 497 Some(next) => link(slot_state[next.index()].unwrap_unused()).prev = prev, 498 None => self.tail = prev, 499 } 500 501 // Same as the `next` node, except everything is in reverse. 502 match prev { 503 Some(prev) => link(slot_state[prev.index()].unwrap_unused()).next = next, 504 None => self.head = next, 505 } 506 state 507 } 508 509 #[cfg(test)] iter<'a>( &'a self, states: &'a [SlotState], link: fn(&Unused) -> &Link, ) -> impl Iterator<Item = SlotId> + 'a510 fn iter<'a>( 511 &'a self, 512 states: &'a [SlotState], 513 link: fn(&Unused) -> &Link, 514 ) -> impl Iterator<Item = SlotId> + 'a { 515 let mut cur = self.head; 516 let mut prev = None; 517 std::iter::from_fn(move || { 518 if cur.is_none() { 519 assert_eq!(prev, self.tail); 520 } 521 let ret = cur?; 522 match &states[ret.index()] { 523 SlotState::UnusedWarm(u) => { 524 assert_eq!(link(u).prev, prev); 525 prev = Some(ret); 526 cur = link(u).next 527 } 528 _ => unreachable!(), 529 } 530 Some(ret) 531 }) 532 } 533 } 534 535 #[cfg(test)] 536 mod test { 537 use super::*; 538 use wasmtime_environ::EntityRef; 539 540 #[test] test_next_available_allocation_strategy()541 fn test_next_available_allocation_strategy() { 542 for size in 0..20 { 543 let state = ModuleAffinityIndexAllocator::new(size, 0); 544 assert_eq!(state.num_empty_slots(), usize::try_from(size).unwrap()); 545 for i in 0..size { 546 assert_eq!(state.num_empty_slots(), usize::try_from(size - i).unwrap()); 547 assert_eq!(state.alloc(None).unwrap().index(), i as usize); 548 } 549 assert!(state.alloc(None).is_none()); 550 } 551 } 552 553 #[test] test_affinity_allocation_strategy()554 fn test_affinity_allocation_strategy() { 555 let id1 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); 556 let id2 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); 557 let state = ModuleAffinityIndexAllocator::new(100, 100); 558 559 let index1 = state.alloc(Some(id1)).unwrap(); 560 assert_eq!(index1.index(), 0); 561 let index2 = state.alloc(Some(id2)).unwrap(); 562 assert_eq!(index2.index(), 1); 563 assert_ne!(index1, index2); 564 565 state.free(index1, 0); 566 assert_eq!(state.num_empty_slots(), 99); 567 568 // Allocate to the same `index1` slot again. 569 let index3 = state.alloc(Some(id1)).unwrap(); 570 assert_eq!(index3, index1); 571 state.free(index3, 0); 572 573 state.free(index2, 0); 574 575 // Both id1 and id2 should have some slots with affinity. 576 let affinity_modules = state.testing_module_affinity_list(); 577 assert_eq!(2, affinity_modules.len()); 578 assert!(affinity_modules.contains(&id1)); 579 assert!(affinity_modules.contains(&id2)); 580 581 // Now there is 1 free instance for id2 and 1 free instance 582 // for id1, and 98 empty. Allocate 100 for id2. The first 583 // should be equal to the one we know was previously used for 584 // id2. The next 99 are arbitrary. 585 assert_eq!(state.num_empty_slots(), 100); 586 let mut indices = vec![]; 587 for _ in 0..100 { 588 indices.push(state.alloc(Some(id2)).unwrap()); 589 } 590 assert!(state.alloc(None).is_none()); 591 assert_eq!(indices[0], index2); 592 assert_eq!(state.num_empty_slots(), 0); 593 594 for i in indices { 595 state.free(i, 0); 596 } 597 598 // Now there should be no slots left with affinity for id1. 599 let affinity_modules = state.testing_module_affinity_list(); 600 assert_eq!(1, affinity_modules.len()); 601 assert!(affinity_modules.contains(&id2)); 602 603 // Allocate an index we know previously had an instance but 604 // now does not (list ran empty). 605 let index = state.alloc(Some(id1)).unwrap(); 606 state.free(index, 0); 607 } 608 609 #[test] clear_affine()610 fn clear_affine() { 611 let id = CompiledModuleId::new(); 612 let memory_index = DefinedMemoryIndex::new(0); 613 614 for max_unused_warm_slots in [0, 1, 2] { 615 let state = ModuleAffinityIndexAllocator::new(100, max_unused_warm_slots); 616 617 let index1 = state.alloc(Some(MemoryInModule(id, memory_index))).unwrap(); 618 let index2 = state.alloc(Some(MemoryInModule(id, memory_index))).unwrap(); 619 state.free(index2, 0); 620 state.free(index1, 0); 621 assert!( 622 state 623 .alloc_affine_and_clear_affinity(id, memory_index) 624 .is_some() 625 ); 626 assert!( 627 state 628 .alloc_affine_and_clear_affinity(id, memory_index) 629 .is_some() 630 ); 631 assert_eq!( 632 state.alloc_affine_and_clear_affinity(id, memory_index), 633 None 634 ); 635 } 636 } 637 638 #[test] test_affinity_allocation_strategy_random()639 fn test_affinity_allocation_strategy_random() { 640 use rand::Rng; 641 let mut rng = rand::rng(); 642 643 let ids = std::iter::repeat_with(|| { 644 MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)) 645 }) 646 .take(10) 647 .collect::<Vec<_>>(); 648 let state = ModuleAffinityIndexAllocator::new(1000, 1000); 649 let mut allocated: Vec<SlotId> = vec![]; 650 let mut last_id = vec![None; 1000]; 651 652 let mut hits = 0; 653 let amt = if cfg!(miri) { 100 } else { 100_000 }; 654 for _ in 0..amt { 655 loop { 656 if !allocated.is_empty() && rng.random_bool(0.5) { 657 let i = rng.random_range(0..allocated.len()); 658 let to_free_idx = allocated.swap_remove(i); 659 state.free(to_free_idx, 0); 660 } else { 661 let id = ids[rng.random_range(0..ids.len())]; 662 let index = match state.alloc(Some(id)) { 663 Some(id) => id, 664 None => continue, 665 }; 666 if last_id[index.index()] == Some(id) { 667 hits += 1; 668 } 669 last_id[index.index()] = Some(id); 670 allocated.push(index); 671 } 672 break; 673 } 674 } 675 676 // 10% reuse would be random chance (because we have 10 module 677 // IDs). Check for at least double that to ensure some sort of 678 // affinity is occurring. 679 assert!( 680 hits > (amt / 5), 681 "expected at least 20000 (20%) ID-reuses but got {hits}" 682 ); 683 } 684 685 #[test] test_affinity_threshold()686 fn test_affinity_threshold() { 687 let id1 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); 688 let id2 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); 689 let id3 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); 690 let state = ModuleAffinityIndexAllocator::new(10, 2); 691 692 // Set some slot affinities 693 assert_eq!(state.alloc(Some(id1)), Some(SlotId(0))); 694 state.free(SlotId(0), 0); 695 assert_eq!(state.alloc(Some(id2)), Some(SlotId(1))); 696 state.free(SlotId(1), 0); 697 698 // Only 2 slots are allowed to be unused and warm, so we're at our 699 // threshold, meaning one must now be evicted. 700 assert_eq!(state.alloc(Some(id3)), Some(SlotId(0))); 701 state.free(SlotId(0), 0); 702 703 // pickup `id2` again, it should be affine. 704 assert_eq!(state.alloc(Some(id2)), Some(SlotId(1))); 705 706 // with only one warm slot available allocation for `id1` should pick a 707 // fresh slot 708 assert_eq!(state.alloc(Some(id1)), Some(SlotId(2))); 709 710 state.free(SlotId(1), 0); 711 state.free(SlotId(2), 0); 712 713 // ensure everything stays affine 714 assert_eq!(state.alloc(Some(id1)), Some(SlotId(2))); 715 assert_eq!(state.alloc(Some(id2)), Some(SlotId(1))); 716 assert_eq!(state.alloc(Some(id3)), Some(SlotId(0))); 717 718 state.free(SlotId(1), 0); 719 state.free(SlotId(2), 0); 720 state.free(SlotId(0), 0); 721 722 // LRU is 1, so that should be picked 723 assert_eq!( 724 state.alloc(Some(MemoryInModule( 725 CompiledModuleId::new(), 726 DefinedMemoryIndex::new(0) 727 ))), 728 Some(SlotId(1)) 729 ); 730 731 // Pick another LRU entry, this time 2 732 assert_eq!( 733 state.alloc(Some(MemoryInModule( 734 CompiledModuleId::new(), 735 DefinedMemoryIndex::new(0) 736 ))), 737 Some(SlotId(2)) 738 ); 739 740 // This should preserve slot `0` and pick up something new 741 assert_eq!( 742 state.alloc(Some(MemoryInModule( 743 CompiledModuleId::new(), 744 DefinedMemoryIndex::new(0) 745 ))), 746 Some(SlotId(3)) 747 ); 748 749 state.free(SlotId(1), 0); 750 state.free(SlotId(2), 0); 751 state.free(SlotId(3), 0); 752 753 // for good measure make sure id3 is still affine 754 assert_eq!(state.alloc(Some(id3)), Some(SlotId(0))); 755 } 756 757 #[test] test_freelist()758 fn test_freelist() { 759 let allocator = SimpleIndexAllocator::new(10); 760 assert_eq!(allocator.testing_freelist(), []); 761 let a = allocator.alloc().unwrap(); 762 assert_eq!(allocator.testing_freelist(), []); 763 allocator.free(a, 0); 764 assert_eq!(allocator.testing_freelist(), [a]); 765 assert_eq!(allocator.alloc(), Some(a)); 766 assert_eq!(allocator.testing_freelist(), []); 767 let b = allocator.alloc().unwrap(); 768 assert_eq!(allocator.testing_freelist(), []); 769 allocator.free(b, 0); 770 assert_eq!(allocator.testing_freelist(), [b]); 771 allocator.free(a, 0); 772 assert_eq!(allocator.testing_freelist(), [b, a]); 773 } 774 } 775