Allocator as a student exercise

Continuing the discussion from Unsafe Rust - how do I inform the compiler of side effects?:

@godmar : I'd consider providing a raw-memory abstraction that has runtime safeguards against UB, but still requires the students to get their logic correct. Here's what I came up with as a first attempt; it compiles, but I haven't really checked to see if it's correct. It allows reading and writing at arbitrary locations, but will panic if the user tries to:

  • Store an improperly-aligned value
  • Read an uninitialized value
  • Read or write an allocation with a type other than was stored there
use std::mem;
use std::ops::Range;
use std::alloc;
use std::cell::RefCell;

use std::collections::HashMap;
use std::any::TypeId;

pub struct AllocationArena<const LEN:usize> {
    mem: *mut u8,
    // everything below here is for runtime verification only,
    // not necessary for production allocators
    allocations: RefCell<HashMap<Range<usize>, TypeId>>,    
}

impl<const LEN:usize> AllocationArena<LEN> {
    pub fn base(&self)->usize { self.mem as usize }
    pub fn new() -> Self {
        AllocationArena {
            mem: unsafe {
                alloc::alloc_zeroed(
                    alloc::Layout::from_size_align(LEN, 4096).unwrap()
                ) },
            allocations: RefCell::new(HashMap::new())
        }    
    }
    
    pub fn store<T:Sized + 'static + Copy>(&self, pos: Range<usize>, data:T) {
        let lo = pos.start;
        let hi = pos.end;
        assert_eq!(0, lo % mem::align_of::<T>());
        assert_eq!(lo - hi, mem::size_of::<T>());
        assert!(lo >= self.base());
        assert!(hi < self.base() + LEN);
        for alloc in self.allocations.borrow().keys() {
            assert!(alloc.start >= hi || alloc.end < lo);
        }
        self.allocations.borrow_mut().insert(pos, TypeId::of::<T>());
        unsafe {
            let dst = self.mem.offset(
                (lo - self.base()) as isize
            ) as *mut T;
            std::ptr::write(dst, data);
        }
    }
    
    pub fn read<T:Sized + 'static + Copy>(&self, pos: Range<usize>)->T {
        assert_eq!(self.allocations.borrow()[&pos], TypeId::of::<T>());
        unsafe {
            std::ptr::read(
                self.mem.offset((pos.start - self.base()) as isize)
                as *mut T
            )
        }
    }
    
    pub fn release(&self, pos: Range<usize>) {
        let _ = self.allocations.borrow_mut().remove(&pos).unwrap();
    }
}

Playground

2 Likes

Thanks for the idea.

Generally, an memory allocator would be given a growable region of memory (like a standard process heap) and would provide a malloc/free style interface that provides a block-based abstraction of memory.

Within the range of memory that the allocator manages at a given point in time, there will be allocated blocks (given to a client), and there will be free blocks (available for allocation). Free blocks in turn contain link elements for the intrusive data structures (linked list, rb tree) used to quickly find a block of appropriate size. The allocator updates the link elements contained in those blocks.

With this background, I'm trying to understand whether the "store/read" abstraction is useful or not. I assume their intended use is to help students remember where in the managed heap boundary tag headers and link elements are located. However, this may be too difficult to manage to be useful (?). Consider a common scenario: coalescing of a free block with its right neighbor. In this case, the right neighbor has a link element that's in a free list. Normally, the right neighbor would be removed from the free list and the information that it was in a free list (and where the link element was located) is simply "forgotten." Similarly, the old boundary tag footer of the block and the boundary tag header of the right block would be simply left as is. So store/read would require additional management that would not be present in a regular allocator. (But, perhaps this effort is worth it for debugging? In the C version of the assignment, we use an LLVM-based instrumentation tool to dynamically try to check student code for bugs in which the boundary tag/free list structure is likely violated. Essentially, the instrumenter keeps track internally of allocations and tries to deduce when the C code does store/read. For instance, we know that a boundary tag header is likely at payload - 4 or -8 and so can infer this from what's returned to the client on a malloc call.)

A second question is whether the store/read machinery can be compiled out (the student implementations are benchmarked and their performance is compared to that of a libc production allocator). When being benchmarked there should be no penalty. A potential concern here is the use of str::ptr::read str::ptr::write which makes a copy of T. (Of course, if T is just a word/pointer this would likely be zero overhead though I'm not sure yet about the ergonomics).

The goal here is to promote the most common sources of UB into runtime panics during development: reading invalid bit patterns and making unaligned / out-of-bound writes. I realized after writing this that release is redundant: Instead, you remove any old overlapping allocations during store, as they've been overwritten and are now invalid. With this change, you can "forget" a link header in exactly the same way as you otherwise would: "store" a [MaybeUninit<u8>;N] over the top of it, and it'll be automatically removed from the runtime tracking. (1)

The biggest downside of this system is that you have to load/store an entire type at a time, instead of just working with a single field. If you want your link headers to have 2-3 fields, you'd need to decide between reading and writing all the fields together (as a struct) every time or working with them individually every time (as usizes, etc.) (modulo whatever the optimizer decides to do).

What the runtime tracking is essentially doing is preventing accidental bit reinterpretation, as it will panic if you attempt to read a different type from a location than was originally written. This requires significant overhead compared to a real allocator, but that overhead is entirely diagnostic and contained within the "store/read" abstraction. It should be easy enough to define a #[cfg(...)] flag that removes the entire mechanism, leaving UB wherever the checked program would have panic!ed (2).

It may take some tweaking to come up with an appropriate interface, depending on how you expect the allocator to be written and what system facilities you want to emulate. I had assumed an allocator that manages a fixed-size memory space, but it shouldn't be too hard to emulate a variable-sized one similarly.


(1) It's probably a good idea to have a specialty method for this that doesn't require overwriting the entire block the way store does

(2) If you do this, these functions should properly be marked unsafe, as they won't necessarily have guardrails anymore.

Ideally I don't want to emulate anything, but rather implement an actual allocator that the compiler could use (but perhaps that's too ambitious).

Regarding your other comments, I think I get the idea. It would be, perhaps, nicer if this could be done with a macro as some kind of meta facility (similar to the C version that as I mentioned uses dynamic instrumentation.) Like so:

#[checked_store]
p.f = v

but I don't know enough Rust to see if that's possible (it would turn this into something like arena.store(address of p + offset of f, v) etc.

Regarding what's being written at a time. It's typically accessing bitfields (which I still need to research how to do in Rust), or list insertions that set both a next and a prev field, which the optimizer ought to optimize even when done as std::ptr::write - I'm making an optimistic guess here.

I'll call the allocations field something else to avoid confusion with the "allocations" made by the allocator's client.

Also, thinking out loud here, in C, we'd look at memory typically like so:

struct block {
     struct boundary_tag header;
     char payload[0];  // if block is allocated
     struct listelem elem; // link element if block is free
}

and then we operate on struct block *b pointers but read/write to &b->header or &b->elem. I think though that this should be compatible with your runtime-checked store/read idea. I would keep the release btw even if it's equivalent to "storing" a MaybeUninit.

1 Like

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.