Handling External Ownership of Buffer and Safe std::mem::forget


#1

I’m looking for general design guidance around patterns for handling mem::forget in situations where a buffer is owned by some asynchronous library.

I’m contributing to the Rust MPI library rsmpi. MPI is a standard for scalable programs in scientific computing.

One of the features of MPI is its asynchronous communication routines. E.g.:

int MPI_Isend(const void *buf, int count, MPI_Datatype datatype, int dest, int tag,
              MPI_Comm comm, MPI_Request *request);

int MPI_Irecv(void *buf, int count, MPI_Datatype datatype, int source,
              int tag, MPI_Comm comm, MPI_Request *request);

MPI_Isend (the I stands for “immediate”), takes a buffer buf of size count and type datatype, sends it to dest in the communicator comm and returns a request. Conceptually, request owns a borrow of buf until the user completes the request by calling MPI_Wait(request), at which point MPI has release the buffer back to the user. MPI_Irecv does the same, but with a mutable borrow.

A naive Rust-managed version of MPI_Request would look like this:

struct Request<'a> {
    request: MPI_Request,
    phantom: PhantomData<Cell<&'a ()>>,
}

impl<'a> Drop for Request<'a> {
    fn drop(&mut self) {
        assert!(request == MPI_REQUEST_NULL,
        "The user did not complete the request before it went out of scope! \
         This is unsafe because an uncompleted Request may still use its attached buffers.");
    }
}

Then ideally, something like this would be safe:

{
    let mut my_recv_buffer = [0i32; 8];
    let request = comm.process_at_rank(src).immediate_recv(my_recv_buffer);

    // If code does not take this branch, for whatever reason, the code should still
    // be safe because the program will crash rather than continue with a potential
    // use after free.
    if ... {
        request.wait();
    }
}

Unfortunately this hypothetical API would not be safe - in light of std::mem::forget, the following code could be written completely in safe code.

{
    let mut my_recv_buffer = [0i32; 8];
    let request = comm.process_at_rank(src).immediate_recv(&mut my_recv_buffer);

    std::mem::forget(request);
}
// Uh-oh! The request persists, but the buffer is no longer borrowed!
// These APIs are not actually safe!

Because forget causes the Request destructor to not run, it can defeat our guarantee that the request does not yield ownership of the buffer before it is completed.

rsmpi solves the problem in the following way - all “immediate” functions take a scope parameter, where the scope is an un-forgettable type. E.g.

let mut recv_buffer = [0i32; 8];
mpi::request::scope(|scope: &LocalScope| {
    let request = world.process_at_rank(src)
        .immediate_recv(scope, &mut recv_buffer);
});

LocalScope is defined such that each Request registers with scope. At the end of the lambda, which is a FnOnce, the code panics if there are any registered Requests that haven’t been completed. It goes without saying that the lifetime of any buffers must be greater than the lifetime of scope.

One of the changes I’m making is to remove the scope field, and instead allow the Request to take ownership of the buffers that the request owns. This makes the API a little less arduous if you just want to pass Vec<T> as the send or receive buffer. Unfortunately, you still need the scope concept if you want to use a buffer borrow. e.g. there would be some API like let scoped_buffer = scope.attach(send_buffer) with the same semantics we currently have. This change would mean for a substantial amount of code, it’s possible to avoid using scope.

Now my question - is there a better code pattern for this besides scope? It’s kind of frustrating that the ownership semantics of Rust can perfectly guarantee that the request’s buffers outlive the request, but due to mem::forget, cannot guarantee that the Request is safely “cleaned up” before yielding ownership of the buffers.


#2

On the surface this sounds similar to the old situation with scoped thread guards where drop blocked the current thread until the background thread finished. That was deemed unsound for the same reason: drop() may not be called. Similar functionality was created in crossbeam but it uses a scope concept not unlike yours (in theory).

However, I don’t quite understand something about your case. If the request takes ownership of the buffer, then how can the buffer leak back out without someone calling something on request to get it back?


#3

I just remembered reading this blog post by @japaric a while back - it talks about a similar situation in the context of DMA transfers. If you want to support reading into borrowed buffers then the callback approach in the blog is a good alternative. But perhaps you can use owned buffers and move them in and out of the request operations (rather than borrowing).


#4

Sorry if it wasn’t clear - the scope is intended to guard borrowed buffers, but I recently made a change to remove the need for scopes if you allow the request to own the buffer (Vec<T>, Box<[T]>, etc.).

Yeah, looks like @japaric’s solution to the problem is substantially similar.

This particular corner of Rust is confusing to me because it feels like the borrow checker isn’t doing its job, but I suppose as far as the borrow checker is concerned, it can’t really see the opaque borrow of the buffer that is hidden inside of MPI_Request (and that will mutated on what is most likely another thread), so I suppose it makes sense that it can’t really enforce this particular borrowing rule.


#5

Yeah, if it works for your cases then taking owned buffers seems more straightforward and bulletproof. This is the approach that futures/tokio take - lots of methods consume values, and you get them back only when the operation completes.

Your Request<'a> causes the borrowed buffer to move into Request so from that standpoint borrowck knows its owned by Request. So long as Request is alive, you cannot use the original buffer. But once the Request is gone, the mutable borrow ends and the original buffer is live again, ready to be re-used. If you move the buffer, however, then this problem goes away since you’d need to explicitly get the buffer back from the Request, and that would be the place where you can ensure the operation completes.


#6

Well, we still want to support using borrowed buffers for performance reasons, so we’ll need to keep the scope concept for those purposes. But I agree, allowing the Request to own the buffers in cases where having to do an allocation (or re-use an existing allocation) isn’t going to have a performance penalty solves the problem nicely. I suppose you could imagine an efficient domain-specific datatype that allows a user to explicitly check out overlapping sections of a buffer (kind of a more advanced Rc), and those sections could only be checked out as mutable if all other immutable checkouts have been explicitly returned. Such a datatype would allow you to safely forget the request (though, of course, this would cause the check out to leak).

Buuut using the scope type for now allows us to let the borrow checker do the hard work for us.


#7

Maybe it’s a terrible idea, but perhaps you can forge the buffer’s lifetime to be 'static and then use the same idea that’s in @japaric’s blog post: transfer ownership over a &'static mut [u8]. The only way to get the buffer back is to take it out of the Request - the Request going out of scope is not sufficient to make the original buffer reusable.

I think your scope idea is cleaner (and safer), but wanted to throw this option out on the table.