swage_thp/
lib.rs

1//! Transparent Huge Pages (THP) memory allocator.
2//!
3//! This crate provides a memory allocator that uses Linux Transparent Huge Pages
4//! to obtain 2MB physically contiguous memory blocks. THP must be enabled in the
5//! kernel configuration.
6//!
7//! Implements the [`swage_core::allocator::ConsecAllocator`] trait.
8//!
9//! # Platform Requirements
10//!
11//! - x86_64 Linux with THP support enabled
12//! - THP should be set to "always" or "madvise" mode
13
14#![warn(missing_docs)]
15
16use std::ptr::null_mut;
17
18use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
19use itertools::max;
20use log::{debug, log_enabled, warn};
21use swage_core::allocator::ConsecAllocator;
22use swage_core::memory::{
23    ConsecBlocks, GetConsecPfns, PfnResolver, TimerError, construct_memory_tuple_timer,
24};
25use swage_core::util::Size::MB;
26use swage_core::util::{NamedProgress, Size};
27use swage_core::{memory::Memory, util::PAGE_SIZE};
28use thiserror::Error;
29
30/// THP allocator. This allocator uses Linux Transparent Huge Pages to obtain 2MB physically contiguous memory blocks.
31pub struct THP {
32    conflict_threshold: u64,
33    progress: Option<MultiProgress>,
34}
35
36impl THP {
37    /// Constructor for THP allocator
38    pub fn new(conflict_threshold: u64, progress: Option<MultiProgress>) -> Self {
39        THP {
40            conflict_threshold,
41            progress,
42        }
43    }
44}
45
46const ALIGN_SIZE: Size = MB(2);
47
48impl THP {
49    /// allocate a 2 MB physically aligned memory block.
50    fn allocate_2m_aligned(size: Size) -> Result<Memory, std::io::Error> {
51        let aligned = unsafe {
52            libc::mmap(
53                null_mut(),
54                size.bytes(),
55                libc::PROT_READ | libc::PROT_WRITE,
56                libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
57                -1,
58                0,
59            )
60        };
61        if aligned == libc::MAP_FAILED {
62            return Err(std::io::Error::last_os_error());
63        }
64        unsafe { libc::memset(aligned, 0, size.bytes()) };
65        if unsafe { libc::madvise(aligned, size.bytes(), libc::MADV_COLLAPSE) } != 0 {
66            return Err(std::io::Error::last_os_error());
67        }
68        unsafe { libc::mlock(aligned, PAGE_SIZE) };
69        if log_enabled!(log::Level::Debug)
70            && let Ok(consecs) = (aligned, size.bytes()).consec_pfns()
71        {
72            debug!("Aligned PFNs: {:?}", consecs);
73        }
74        assert_eq!(aligned as usize & (ALIGN_SIZE.bytes() - 1), 0);
75        assert_eq!(
76            aligned.pfn().unwrap_or_default().as_usize() & (ALIGN_SIZE.bytes() - 1),
77            0
78        );
79        Ok(Memory::new(aligned as *mut u8, size.bytes()))
80    }
81}
82
83/// Errors that can happen during THP allocation
84#[derive(Debug, Error)]
85#[allow(missing_docs)]
86pub enum Error {
87    #[error(transparent)]
88    IoError(#[from] std::io::Error),
89    #[error(transparent)]
90    TimerError(#[from] TimerError),
91    #[error("Size must be a multiple of {0}")]
92    SizeError(Size),
93}
94
95impl ConsecAllocator for THP {
96    type Error = Error;
97
98    fn block_size(&self) -> swage_core::util::Size {
99        Size::GB(1)
100    }
101
102    fn alloc_consec_blocks(
103        &mut self,
104        size: swage_core::util::Size,
105    ) -> Result<swage_core::memory::ConsecBlocks, Self::Error> {
106        if size.bytes() == 0 || !size.bytes().is_multiple_of(ALIGN_SIZE.bytes()) {
107            return Err(Error::SizeError(ALIGN_SIZE));
108        }
109        let mut blocks: Vec<Memory> = vec![];
110        let required_blocks =
111            max([size.bytes() / self.block_size().bytes(), 1]).expect("empty iter");
112        let timer = construct_memory_tuple_timer()?;
113        let p = self.progress.as_ref().map(|p| {
114            p.add(
115                ProgressBar::new(required_blocks as u64)
116                    .with_style(ProgressStyle::named_bar("Allocating blocks")),
117            )
118        });
119        let mut garbage = vec![];
120        while blocks.len() < required_blocks {
121            let block = Self::allocate_2m_aligned(size)?;
122
123            // check for same bank
124            if let Some(last_block) = blocks.last() {
125                let timing = unsafe {
126                    timer.time_subsequent_access_from_ram(block.ptr, last_block.ptr, 10000)
127                };
128                let same_bank = timing >= self.conflict_threshold;
129                if !same_bank {
130                    warn!(
131                        "Bank check failed: {} < {} for blocks {:?} and {:?}",
132                        timing, self.conflict_threshold, block, last_block
133                    );
134                    block.log_pfns(log::Level::Warn);
135                    last_block.log_pfns(log::Level::Warn);
136                    garbage.push(block);
137                    continue;
138                }
139            }
140            if let Some(p) = &p {
141                p.inc(1);
142            }
143            blocks.push(block);
144        }
145        for block in garbage {
146            block.dealloc();
147        }
148        Ok(ConsecBlocks::new(blocks))
149    }
150}
swage_thp/lib.rs

swage_thp/
lib.rs
