Unsound arena implementation

[andresovela]

Hello,

I am working on a small arena implementation with reusable memory. Miri is not happy with my implementation so I'm here asking for help

The arena behaves like a bump arena, but when an allocation is dropped, the write index for the backing buffer moves back so that the freed memory can be reused. In principle this works, but Miri complains.

Miri is happy with this test:

#[test]
fn can_allocate_deallocated_memory() {
    let mut backing_storage = AlignedStruct([0u8; 1024]);
    let arena = StackArena::from_slice(&mut backing_storage.0);

    let a = arena.allocate::<u8>(1);
    drop(a);

    let b = arena.allocate::<u8>(1);
    drop(b);
}

but not with this one:

#[test]
fn passing_allocation_into_fn() {
    fn takes_allocation<'a, 'b, 'arena>(arena: &'b StackArena<'arena>, t: Allocation<'a, u8>) -> Allocation<'b, u8> where 'arena: 'b {
        drop(t);
        arena.allocate::<u8>(1)
    }

    let mut backing_storage = AlignedStruct([0u8; 1024]);
    let arena = StackArena::from_slice(&mut backing_storage.0);
    
    let a = arena.allocate::<u8>(1);
    let b = takes_allocation(&arena, a);
    drop(b);
}

In both tests, the arena creates both allocations using the same memory.

I've tried all sorts of combinations of lifetimes in takes_allocation(). I haven't been able to make Miri happy.

The code itself compiles and the tests pass.

This is what Miri says:

running 1 test
test arena::stack_arena::tests::passing_allocation_into_fn ... error: Undefined Behavior: not granting access to tag <423014> because that would remove [SharedReadOnly for <423770>] which is strongly protected
   --> /home/andy/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/slice/raw.rs:192:9
    |
192 |         &mut *ptr::slice_from_raw_parts_mut(data, len)
    |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ not granting access to tag <423014> because that would remove [SharedReadOnly for <423770>] which is strongly protected
    |
    = help: this indicates a potential bug in the program: it performed an invalid operation, but the Stacked Borrows rules it violated are still experimental
    = help: see https://github.com/rust-lang/unsafe-code-guidelines/blob/master/wip/stacked-borrows.md for further information
help: <423014> was created by a SharedReadWrite retag at offsets [0x0..0x400]
   --> crates/airten/src/arena/stack_arena.rs:72:22
    |
72  |             storage: slice.as_mut_ptr(),
    |                      ^^^^^^^^^^^^^^^^^^
help: <423770> is this argument
   --> crates/airten/src/arena/stack_arena.rs:545:76
    |
545 |         fn takes_allocation<'a, 'b, 'arena>(arena: &'b StackArena<'arena>, t: Allocation<'a, u8>) -> Allocation<'b, u8> where 'arena: 'b {
    |                                                                            ^
    = note: BACKTRACE (of the first span) on thread `arena::stack_ar`:
    = note: inside `std::slice::from_raw_parts_mut::<'_, u8>` at /home/andy/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/slice/raw.rs:192:9: 192:55
note: inside `arena::stack_arena::tests::passing_allocation_into_fn::takes_allocation::<'_, '_>`
   --> crates/airten/src/arena/stack_arena.rs:547:13
    |
547 |             arena.allocate::<u8>(1)
    |             ^^^^^^^^^^^^^^^^^^^^^^^
note: inside `arena::stack_arena::tests::passing_allocation_into_fn`
   --> crates/airten/src/arena/stack_arena.rs:554:18
    |
554 |         let _b = takes_allocation(&arena, a);
    |                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
note: inside closure
   --> crates/airten/src/arena/stack_arena.rs:544:36
    |
543 |     #[test]
    |     ------- in this procedural macro expansion
544 |     fn passing_allocation_into_fn() {

The relevant implementation of the arena is as follows:

pub(crate) struct StackArena<'a> {
    storage: *mut u8,
    size: usize,
    context: RefCell<Context>,
    _phantom: PhantomData<&'a mut ()>,
}

impl<'a> StackArena<'a> {
    /// Creates a new stack arena from a mutable slice of bytes.
    ///
    /// The slice must be aligned to 16 bytes.
    pub fn from_slice(slice: &'a mut [u8]) -> Self {
        assert!(
            slice.as_mut_ptr().align_offset(16) == 0,
            "Backing storage must be aligned to 16 bytes"
        );

        StackArena {
            storage: slice.as_mut_ptr(),
            size: slice.len(),
            context: RefCell::new(Context {
                item_count: 0,
                write_index: 0,
                allocation_map: 0,
                allocation_size: [0; usize::BITS as usize],
                maximum_memory_usage: 0,
            }),
            _phantom: PhantomData,
        }
    }

    #[track_caller]
    fn create_allocation<T: Allocatable>(&self, count: usize) -> (&'a mut [T], u8) {
        let allocation_size;
        let mut context = self.context.borrow_mut();

        assert!(count > 0, "Cannot allocate zero items");
        assert!(
            context.item_count < usize::BITS as usize,
            "Cannot allocate item: stack arena can only hold up to {} items",
            usize::BITS,
        );

        // SAFETY:
        // - The pointer is guaranteed to be aligned to the alignment of T
        // - The slice is guaranteed to be valid for the entire lifetime of the arena
        // - The arena is guaranteed to have enough space for the allocation
        // - The slice is guaranteed to hold valid values of type T
        let slice = unsafe {
            let ptr = self.storage.add(context.write_index);
            let offset = ptr.align_offset(core::mem::align_of::<T>());
            allocation_size = (core::mem::size_of::<T>() * count) + offset;

            assert!(
                self.size - context.write_index >= allocation_size,
                "Allocation does not fit in stack arena: {} bytes requested, {} available, arena size must be {}",
                allocation_size,
                self.size - context.write_index,
                context.write_index + allocation_size
            );

            let ptr = ptr.add(offset);
            core::slice::from_raw_parts_mut(ptr as *mut T, count)
        };

        // Mark the allocation as used and save its size
        let allocation_id = context.item_count;
        context.allocation_map |= 1 << allocation_id;
        context.allocation_size[allocation_id] = allocation_size;

        // Prepare for the next allocation
        context.item_count += 1;
        context.write_index += allocation_size;

        // Keep track of the maximum memory usage
        if context.write_index > context.maximum_memory_usage {
            context.maximum_memory_usage = context.write_index;
        }

        (slice, allocation_id as u8)
    }

    #[track_caller]
    pub fn allocate<T: Allocatable>(&'a self, count: usize) -> Allocation<'a, T> {
        let (slice, allocation_id) = self.create_allocation::<T>(count);
        slice.fill(T::default());
        Allocation {
            memory: slice,
            context: Some(AllocationContext::Stack(StackAllocationContext {
                id: allocation_id,
                context: &self.context,
            })),
        }
    }
}

// This is called from the Drop impl of an `Allocation` object that wraps the allocation
pub(super) fn drop_impl(stack_context: &StackAllocationContext) {
    let mut context = stack_context.context.borrow_mut();
    debug_assert!(context.item_count > 0);

    // Mark the allocation as unused
    context.allocation_map &= !(1 << stack_context.id);

    let free_slots = context.allocation_map.leading_zeros() as usize;
    let last_occupied_slot = usize::BITS as usize - free_slots;

    // Calculate the size of the allocations to free
    let mut size_to_free = 0;
    for i in last_occupied_slot..context.item_count {
        size_to_free += context.allocation_size[i];
        context.allocation_size[i] = 0;
    }

    // Clear all the slots that were freed
    let freed_slots = context.item_count - last_occupied_slot;
    context.item_count -= freed_slots;
    context.write_index -= size_to_free;
}

I've read into using ManuallyDrop and I've seen several arena impls using this, but I'm not yet sure how this helps solving this particular problem. If anyone can point me in the right direction I'd be very grateful

[2e71828]

@andresovela

I've tried all sorts of combinations of lifetimes in takes_allocation(). I haven't been able to make Miri happy.

IIUC, Changing lifetimes won't ever affect Miri's analysis-- It's validating the actual reads/writes/accesses that your program makes.

@andresovela

Undefined Behavior: not granting access to tag <423014> because that would remove [SharedReadOnly for <423770>] which is strongly protected

This indicates that you have an &mut/& aliasing problem somewhere (or similar). You might need to add an UnsafeCell somewhere to indicate that you'll be mutating behind a shared reference.

[simonbuchan]

Not at all an expert, but from my reading of that MIRI output your bug is that the returned allocation is considered derived from the allocator &self despite being created from a contained *mut by implicitly upgrading the ref to a mutable ref (which would be surprising to me if true!).

If so, that means the returned allocation keeps the allocator's self ref alive, which isn't a problem for the first test as the first allocation drops before the second allocate call does the same and attempts to upgrade the self ref, but the second test keeps the allocator ref alive until you call the second allocate and attempt to upgrade to a second mut ref to the allocator.

Not really sure about any of that, but it's the only way I can read that error. Wild guess - perhaps first copying storage out of self into a local would help? If nothing else, that unsafe block is far too long for my comfort: it looks like it only needs to cover the slice::from_raw_parts_mut() call itself. Also give cast a try instead of as *mut T in case something funky managed to sneak in?

[andresovela]

Hi, thanks for looking into it!

Wild guess - perhaps first copying storage out of self into a local would help?

Can you elaborate? Not sure what that would look like.

If nothing else, that unsafe block is far too long for my comfort: it looks like it only needs to cover the slice::from_raw_parts_mut() call itself.

ptr.add is also unsafe. I could wrap only the unsafe calls in unsafe blocks, but it feels like a lot of boilerplate for little benefit IMO, but I appreciate the feedback!

Also give cast a try instead of as *mut T in case something funky managed to sneak in?

Will do

[andresovela]

@2e71828

IIUC, Changing lifetimes won't ever affect Miri's analysis-- It's validating the actual reads/writes/accesses that your program makes.

That's good to know if true

[simonbuchan]

@andresovela

Can you elaborate? Not sure what that would look like

Simply let storage = self.storage; - I don't have any confidence with this though as add takes self by copy so it should look the same to MIRI. Only reason I'm thinking about it is lifetime extension occasionally kicks in when I don't expect it.

@andresovela

ptr.add is also unsafe. I could wrap only the unsafe calls in unsafe blocks, but it feels like a lot of boilerplate for little benefit IMO, but I appreciate the feedback!

Huh, my understanding was pointer operations were safe until you actually access them! Fair enough then, though generally it's still a good practice to use separate unsafe blocks each with their own SAFETY comments.

[khimru]

@simonbuchan

Huh, my understanding was pointer operations were safe until you actually access them!

Only if you use safe operations like wrapping_add.

[andresovela]

@simonbuchan

Simply let storage = self.storage;

I'm still not sure what you mean, at least I don't see how that would work.

The arena can create new allocations from a shared reference (maybe that's the problem? - context is mutably borrowed via a RefCell so I don't see how that'd be a problem), so moving storage wouldn't work I think.

[simonbuchan]

It's a pointer so it's Copy, unless I'm missing something? The (probably dumb) idea is to ensure you're working on a value that isn't contained in that shared reference, yeah.

[simonbuchan]

For overflow I guess that makes sense (though regular additions aren't unsafe!) but it also has an alignment requirement even though dereferencing has it too? Not sure what's the idea there, maybe if doing the offset could take advantage of the alignment?

[andresovela]

Ah I see, right. I'm still not sure how that would help, but I'll see what I can do with that idea. I'm currently reading the 'nomicon to see if there's something obvious I'm missing.

[simonbuchan]

Yeah, it's a bit squirrelly!

[jdahlstrom]

I don't see an alignment requirement, or rather, it's impossible to create an unaligned pointer by calling add, assuming self itself is aligned. The only safety requirements add has is that the addition cannot overflow and the result must be within the same allocation as self (where the latter essentially implies the former). And as the docs say, this UB can be avoided by using wrapping_add and the only reason to use add is that it can allow more optimizations (though I'm not sure how and what).

[andresovela]

@simonbuchan

Not at all an expert, but from my reading of that MIRI output your bug is that the returned allocation is considered derived from the allocator &self despite being created from a contained *mut by implicitly upgrading the ref to a mutable ref (which would be surprising to me if true!).

I took some time to try to really understand what the error message is saying, and my reading is this:

Miri is saying that calling arena.allocate in takes_allocation is UB because it's getting a &mut to the same memory that t is using. What I don't understand is why is miri complaining about t if t is explicitly dropped before calling arena.allocate:

fn takes_allocation<'a, 'b, 'arena>(arena: &'b StackArena<'arena>, t: Allocation<'a, u8>) -> Allocation<'b, u8> where 'arena: 'b {
    drop(t);
    arena.allocate::<u8>(1)
}

If t wasn't dropped yet, the arena would use a different slice of the backing storage to create the new allocation.

Is it really the case that Miri doesn't care about lifetimes? I have the feeling that Miri is considering t to be alive past the drop.

On another note, I read the 'nomicon now and I couldn't think of anything that needs handling that I'm not already handling, so I'm a bit stuck here now

[khimru]

@andresovela

Is it really the case that Miri doesn't care about lifetimes?

Obviously: lifetimes don't exist after compilation.

@andresovela

I have the feeling that Miri is considering t to be alive past the drop.

If t is Copy type then it would be alive after that, of course. Is it Copy?

[andresovela]

Allocation isn't Copy

[kpreid]

Could you post complete, compilable code (e.g. with the definitions of Context and Allocation, and all uses)? It will be easier to try to understand the problem here if we have the code to experiment with ourselves (on the Rust Playground or otherwise).

[andresovela]

I'll try to set up a minimal repro in the playground

[andresovela]

Here it is: Rust Playground

[kpreid]

I think, maybe, I see the general shape of the problem. I could be wrong. Get a second opinion.

You've passed a &'a [u8] into a function (wrapped in Allocation<'a, u8>, but that doesn't change things). When you call allocate(1) in that function, it tries to make an &mut [u8] to the same memory, which is UB. It doesn’t matter that drop(t) was called, because that runs some logic that doesn't make any difference to the state of the program’s borrows — the original borrow is always going to be at least until the function ends, given the presence of an &'a [u8] (or &'a mut [u8]).

The solution is to store memory: *mut [u8] in Allocation instead. That way, you’re not imposing the constraints on use that come with a reference.

(This wouldn't be a problem if allocations weren’t Droppable and they just lasted until some end-of-scope, because then you’re writing code that fits the shape the the lifetime and noalias rules are expecting.)