为了确保安全性,Rust在零抽象的基础上,限制了很多易导致潜在bug的操作,比如直接指针操作,以及长生命周期等等,但是在某些情况下,如果涉及到操作比较底层的数据,往往难免自行维护内存。
本文主要以layout库来进行手动内存维护演示。
假设要实现一个内存分配器,但是需要满足以下条件:
- 手动从堆中申请内存
- 手动执行内存回收
- 内存空间能够被复用
由于要手动回收内存,可行的方案是实现一个智能指针包裹裸指针,为了将裸指针指向的内存区视为待分配的数据结构,可以采用强制类型转换。
use std::alloc::{alloc, dealloc, Layout};
use std::ops::{Deref, DerefMut};
use std::fmt;
use fmt::Debug;
use std::collections::HashMap;
#[derive(Debug)]
struct Header {
a: i32,
}
struct Cap<T>(*mut u8, Layout, std::marker::PhantomData<T>);
impl<T> Cap<T> {
fn new() -> Self {
let layout = Layout::new::<T>();
unsafe {
let ptr = alloc(layout);
Cap(ptr, layout, std::marker::PhantomData)
}
}
}
impl<T: Debug> fmt::Display for Cap<T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
unsafe {
write!(f, "({:?})", & *(self.0 as *mut T))
}
}
}
impl<T> Drop for Cap<T> {
fn drop(&mut self) {}
}
impl<T> Deref for Cap<T> {
type Target = T;
fn deref(&self) -> &T {
unsafe {
& *(self.0 as *mut T)
}
}
}
impl<T> DerefMut for Cap<T> {
fn deref_mut(&mut self) -> &mut T {
unsafe {
&mut *(self.0 as *mut T)
}
}
}
struct MemCache<T> {
cache_map: HashMap<usize, Vec<Cap<T>>>,
}
impl<T> MemCache<T> {
fn new() ->Self {
MemCache {
cache_map: HashMap::new(),
}
}
fn alloc(&mut self) -> Cap<T> {
let obj_size = std::mem::size_of::<T>();
match self.cache_map.get_mut(&obj_size) {
Some(bucket) => {
if bucket.is_empty() {
return Cap::<T>::new()
} else {
let ele = bucket.remove(0);
return ele
}
},
_ => {
return Cap::<T>::new();
},
}
}
fn drop(&mut self, d: Cap<T>) {
let obj_size = std::mem::size_of::<T>();
match self.cache_map.get_mut(&obj_size) {
Some(bucket) => {
bucket.push(d);
}
_ => {
let mut b = Vec::new();
b.push(d);
self.cache_map.insert(obj_size, b);
}
};
}
}
fn main() {
let mut cache = MemCache::<Header>::new();
let mut t1 = cache.alloc();
t1.a = 38;
println!("Cap {:#}", t1);
cache.drop(t1);
let mut t2 = cache.alloc();
t2.a = 36;
println!("Cap {:#}", t2);
cache.drop(t2);
}
主要思路如下:
- 用智能指针封装裸指针,在deref时,将裸指针转换为具体的类型
- 释放智能指针时,用一个HashMap缓存住,在分配阶段如果有缓存则用缓存
上述的实现不能自动执行drop,需要手动执行,所以在Cap
的Drop
实现是一个空函数。
为了实现自动释放,需要在Cap
中记录下MemCache
,参考实现如下:
use std::fmt;
use std::mem;
use std::ops::{Deref, DerefMut};
use std::ptr;
use crossbeam_queue::SegQueue;
use stable_deref_trait::StableDeref;
pub trait Allocate {
fn alloc() -> Self;
}
#[derive(Debug)]
struct Header {
a: i32,
}
impl Allocate for Header {
fn alloc() -> Self {
Header {
a: 89,
}
}
}
pub struct BytePool<T>
where
T: Allocate,
{
small_bask: SegQueue<T>,
}
pub struct ByteBuffer<'a, T: Allocate> {
data: mem::ManuallyDrop<T>,
pool: &'a BytePool<T>,
}
impl<T: Allocate + fmt::Debug> fmt::Debug for ByteBuffer<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{:?}", self.data.deref())
}
}
impl<T: Allocate + fmt::Debug> fmt::Display for ByteBuffer<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{:?}", self.data.deref())
}
}
impl<T: Allocate> Default for BytePool<T> {
fn default() -> Self {
BytePool::<T> {
small_bask: SegQueue::new(),
}
}
}
impl<T: Allocate> BytePool<T> {
pub fn new() -> Self {
BytePool::default()
}
pub fn alloc(&self) -> ByteBuffer<'_, T> {
let list = &self.small_bask;
if let Some(el) = list.pop() {
return ByteBuffer::new(el, self);
}
let data = T::alloc();
ByteBuffer::new(data, self)
}
fn push_raw_block(&self, buffer: T) {
self.small_bask.push(buffer);
}
}
impl<'a, T: Allocate> Drop for ByteBuffer<'a, T> {
fn drop(&mut self) {
let data = mem::ManuallyDrop::into_inner(unsafe { ptr::read(&self.data) });
self.pool.push_raw_block(data);
}
}
impl<'a, T: Allocate> ByteBuffer<'a, T> {
fn new(data: T, pool: &'a BytePool<T>) -> Self {
ByteBuffer {
data: mem::ManuallyDrop::new(data),
pool,
}
}
pub fn size(&self) -> usize {
std::mem::size_of::<T>()
}
}
impl<'a, T: Allocate> Deref for ByteBuffer<'a, T> {
type Target = T;
#[inline]
fn deref(&self) -> &Self::Target {
self.data.deref()
}
}
impl<'a, T: Allocate> DerefMut for ByteBuffer<'a, T> {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
self.data.deref_mut()
}
}
unsafe impl<'a, T: StableDeref + Allocate> StableDeref for ByteBuffer<'a, T> {}
fn main() {
let p = BytePool::<Header>::new();
{
let mut t = p.alloc();
t.a = 1;
println!("{:#}", t);
}
let mut t = p.alloc();
t.a = 1;
println!("{:#}", t);
}