1use crate::hash_map::{Entry, HashMap};
4use crate::prelude::*;
5use crate::runtime::vm::CompiledModuleId;
6use std::mem;
7use std::sync::Mutex;
8use wasmtime_environ::DefinedMemoryIndex;
9
10#[derive(Hash, Clone, Copy, Debug, PartialEq, Eq)]
12pub struct SlotId(pub u32);
13
14impl SlotId {
15 pub fn index(self) -> usize {
17 self.0 as usize
18 }
19}
20
21#[derive(Debug)]
27pub struct SimpleIndexAllocator(ModuleAffinityIndexAllocator);
28
29impl SimpleIndexAllocator {
30 pub fn new(capacity: u32) -> Self {
31 SimpleIndexAllocator(ModuleAffinityIndexAllocator::new(capacity, 0))
32 }
33
34 pub fn is_empty(&self) -> bool {
35 self.0.is_empty()
36 }
37
38 pub fn alloc(&self) -> Option<SlotId> {
39 self.0.alloc(None)
40 }
41
42 pub(crate) fn free(&self, index: SlotId, bytes_resident: usize) {
48 self.0.free(index, bytes_resident);
49 }
50
51 pub fn unused_warm_slots(&self) -> u32 {
57 self.0.unused_warm_slots()
58 }
59
60 pub fn unused_bytes_resident(&self) -> usize {
66 self.0.unused_bytes_resident()
67 }
68
69 #[cfg(test)]
70 pub(crate) fn testing_freelist(&self) -> Vec<SlotId> {
71 self.0.testing_freelist()
72 }
73}
74
75#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
77pub struct MemoryInModule(pub CompiledModuleId, pub DefinedMemoryIndex);
78
79#[derive(Debug)]
82pub struct ModuleAffinityIndexAllocator(Mutex<Inner>);
83
84#[derive(Debug)]
85struct Inner {
86 max_unused_warm_slots: u32,
93
94 unused_warm_slots: u32,
99
100 warm: List,
103
104 last_cold: u32,
109
110 slot_state: Vec<SlotState>,
115
116 module_affine: HashMap<MemoryInModule, List>,
122
123 unused_bytes_resident: usize,
125}
126
127#[derive(Default, Debug)]
129struct List {
130 head: Option<SlotId>,
131 tail: Option<SlotId>,
132}
133
134#[derive(Default, Debug, Copy, Clone)]
137struct Link {
138 prev: Option<SlotId>,
139 next: Option<SlotId>,
140}
141
142#[derive(Clone, Debug)]
143enum SlotState {
144 Used(Option<MemoryInModule>),
146
147 UnusedCold,
149
150 UnusedWarm(Unused),
155}
156
157impl SlotState {
158 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)]
167struct Unused {
168 affinity: Option<MemoryInModule>,
170
171 bytes_resident: usize,
174
175 affine_list_link: Link,
177
178 unused_list_link: Link,
180}
181
182enum AllocMode {
183 ForceAffineAndClear,
184 AnySlot,
185}
186
187impl ModuleAffinityIndexAllocator {
188 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 pub fn len(&self) -> usize {
203 let inner = self.0.lock().unwrap();
204 inner.slot_state.len()
205 }
206
207 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 pub fn alloc(&self, for_memory: Option<MemoryInModule>) -> Option<SlotId> {
221 self._alloc(for_memory, AllocMode::AnySlot)
222 }
223
224 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
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 let slot_id = inner.pick_affine(for_memory).or_else(|| {
250 match mode {
251 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 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
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 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 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 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 #[cfg(test)]
347 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 #[cfg(test)]
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 #[cfg(test)]
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 pub fn unused_warm_slots(&self) -> u32 {
378 self.0.lock().unwrap().unused_warm_slots
379 }
380
381 pub fn unused_bytes_resident(&self) -> usize {
387 self.0.lock().unwrap().unused_bytes_resident
388 }
389}
390
391impl Inner {
392 fn pick_affine(&mut self, for_memory: Option<MemoryInModule>) -> Option<SlotId> {
395 let ret = self.module_affine.get(&for_memory?)?.tail?;
399 self.remove(ret);
400 Some(ret)
401 }
402
403 fn pick_warm(&mut self) -> Option<SlotId> {
404 let head = self.warm.head?;
409 self.remove(head);
410 Some(head)
411 }
412
413 fn remove(&mut self, slot: SlotId) {
414 self.unused_warm_slots -= 1;
417 self.warm
418 .remove(slot, &mut self.slot_state, |u| &mut u.unused_list_link);
419
420 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
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
450impl List {
451 fn new(id: SlotId) -> List {
453 List {
454 head: Some(id),
455 tail: Some(id),
456 }
457 }
458
459 fn append(
461 &mut self,
462 id: SlotId,
463 states: &mut [SlotState],
464 link: fn(&mut Unused) -> &mut Link,
465 ) -> Link {
466 let tail = mem::replace(&mut self.tail, Some(id));
468
469 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 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 match next {
497 Some(next) => link(slot_state[next.index()].unwrap_unused()).prev = prev,
498 None => self.tail = prev,
499 }
500
501 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)]
510 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)]
536mod test {
537 use super::*;
538 use wasmtime_environ::EntityRef;
539
540 #[test]
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]
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 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 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 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 let affinity_modules = state.testing_module_affinity_list();
600 assert_eq!(1, affinity_modules.len());
601 assert!(affinity_modules.contains(&id2));
602
603 let index = state.alloc(Some(id1)).unwrap();
606 state.free(index, 0);
607 }
608
609 #[test]
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]
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 assert!(
680 hits > (amt / 5),
681 "expected at least 20000 (20%) ID-reuses but got {hits}"
682 );
683 }
684
685 #[test]
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 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 assert_eq!(state.alloc(Some(id3)), Some(SlotId(0)));
701 state.free(SlotId(0), 0);
702
703 assert_eq!(state.alloc(Some(id2)), Some(SlotId(1)));
705
706 assert_eq!(state.alloc(Some(id1)), Some(SlotId(2)));
709
710 state.free(SlotId(1), 0);
711 state.free(SlotId(2), 0);
712
713 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 assert_eq!(
724 state.alloc(Some(MemoryInModule(
725 CompiledModuleId::new(),
726 DefinedMemoryIndex::new(0)
727 ))),
728 Some(SlotId(1))
729 );
730
731 assert_eq!(
733 state.alloc(Some(MemoryInModule(
734 CompiledModuleId::new(),
735 DefinedMemoryIndex::new(0)
736 ))),
737 Some(SlotId(2))
738 );
739
740 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 assert_eq!(state.alloc(Some(id3)), Some(SlotId(0)));
755 }
756
757 #[test]
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}