Exception safety in Rust: using transient droppers to prevent memory leaks

I have recently put in some documentation and code reviewing work to help with the 1.0.0 release of array-init. This crate allows user to provide an initializer function that will be used to fill an array. The twist compared to Rust's std is that you don't have to first initialize your array with some valid value before running your function on every entry.

If that's what you need, using this crate saves you from having to figure out two things:

1. How to properly use MaybeUninit and transmutations to manipulate uninitialized memory.
2. What to do in case the initializer fails, and you are left with only part of your entries initialized.

Point 2 is called exception safety and it is what I'll talk about in this post.

Exception safety

The Rustonomicon has a whole chapter dedicated to exception safety, which defines it as "being ready for unwinding". To understand what exception safety concretely means in the array-init case, let's look at the following simplified code¹ which initializes an array:

let mut uninit_array: MaybeUninit<[usize;5]> = MaybeUninit::uninit();
// pointer to array <-> pointer to first element
let mut ptr_i = uninit_array.as_mut_ptr() as *mut usize;
let array : [usize;5] = unsafe {
    for i in 0..5 {
        let value_i = i;
        ptr_i.write(value_i);
        ptr_i = ptr_i.add(1);
    }
    uninit_array.assume_init()
};
assert_eq!(array, [0,1,2,3,4]);

This code is ok as is, but what if we where initializing an array of T from a function f(i: usize) -> T? And what if this function f could fail and was instead of type f(i: usize) -> Option<T>? Then we could start with something like this

let mut uninit_array: MaybeUninit<[T;5]> = MaybeUninit::uninit();
// pointer to array <-> pointer to first element
let mut ptr_i = uninit_array.as_mut_ptr() as *mut T;
let array : Option<[T;5]> = unsafe {
    for i in 0..5 {
        let value_i = f(i).unwrap();
        ptr_i.write(value_i);
        ptr_i = ptr_i.add(1);
    }
    Some(uninit_array.assume_init())
};

But then, what happen if we panic because f returned a None? (Note that even if f could only return a T, the problem still exists because f may panic.)

First, no safety is violated: the array memory is in a inconsistent state but because None is returned, it is not observable and hence we are still sound: minimal exception safety is preserved.

If T does not need drop-glue, this code is also maximally exception safe: I'm not exactly sure when but uninit_array will be dropped at some point, though none of its element will be (a property of MaybeUninit). On the other hand, if T needs to be dropped to release some resources, then it will never be, and the memory that f allocated and the resources it took hold of will be leaked.

It is interesting here to pause and think about whether leaking objects is safe or not. Nowadays, it is pretty well indicated in the docs that leaking memory is not considered unsafe, but it wasn't always so clear. In any case, leaking memory when unwinding (and hence not being maximally exception-safe) is often considered normal.

Transient droppers

Now that the introduction is done, let's dive into the code itself. The code we are interested in, and that I detail after, is the following:

let mut uninit_array: MaybeUninit<[T;5]> = MaybeUninit::uninit();

struct TransientDropper {
    base_ptr: *mut T,
    initialized_count: usize,
}

impl Drop for TransientDropper {
    fn drop(self: &'_ mut Self) {
        unsafe {
            ptr::drop_in_place(slice::from_raw_parts_mut(
                self.base_ptr,
                self.initialized_count,
            ));
        }
    }
}

let mut ptr_i = uninit_array.as_mut_ptr() as *mut T;

let mut transient_dropper = TransientDropper {
    base_ptr: ptr_i,
    initialized_count: 0,
};
let array : Option<[T;5]> = unsafe {
    for i in 0..5 {
        let value_i = f(i).unwrap();
        ptr_i.write(value_i);
        ptr_i = ptr_i.add(1);
        transient_dropper.initialized_count += 1;
    }
    mem::forget(transient_dropper);
    Some(uninit_array.assume_init())
};

The key novelty is TransientDropper. The transient dropper will be used to transiently own the values we have already initialized: during initialization of the array, if a panic occurs, the transient dropper will be dropped, and with it the already initialized elements.

The TransientDropper is implemented as follow.

// TransientDropper has information akin to a slice:
// a base pointer, and the number of initialized
// elements from this pointer
struct TransientDropper {
    base_ptr: *mut T,
    initialized_count: usize,
}

// Here is the whole logic: when TransientDropper
// is dropped, we need to drop the initialized_count
// elements that have already be initialized.
// This is done by dropping the slice containing them all.
impl Drop for TransientDropper {
    fn drop(self: &'_ mut Self) {
        unsafe {
            ptr::drop_in_place(slice::from_raw_parts_mut(
                self.base_ptr,
                self.initialized_count,
            ));
        }
    }
}

and is used here

let array : Option<[T;5]> = unsafe {
    for i in 0..5 {
        let value_i = f(i).unwrap();
        ptr_i.write(value_i);
        ptr_i = ptr_i.add(1);
        // We have initialized one more
        // more element and hence increase
        // the initialized_count
        transient_dropper.initialized_count += 1;
    }
    // We need to forget the transient dropper,
    // see below.
    mem::forget(transient_dropper);
    Some(uninit_array.assume_init())
};

It is capital to forget the transient dropper before returning the array, otherwise the array elements have two different owners, that are both going to try to drop them when they go out of scope, resulting in a use after free.

If you want to test what happens if you remove parts of this code, you can find it on the playground.

Conclusion

Overall I find this solution much more elegant than something based on catch_unwind, and also more generic. If you don't want to implement it yourself, you may find the scopeguard crate interesting (note that I haven't tested it).

Also remarkable are tests for memory leak, which can be found on array-init's github. A crate with similar testing functionality is dropcheck (which I haven't tested either).

Finally, I'd like to warmly thank Daniel Henry-Mantilla, who wrote the array-init code for leak-free exception safety and was kind enough to discuss with me by mail and indicate me many links used in this article.