How to design a good generic container and zero-copy? Getting stuck because of MaybeUninit

help
1

Hi! I started implementing a fast single-threaded ringbuffer, similar to ringbuf. I would have used it directly, but it's missing one thing: zero-copy support for writing.

Here's my use case (the top library):

// Handler is part of a library, and the user supplies the App
struct<App> Handler<App>
where
    App: ApplicationHandler,
{
    app: App,
    ringbuf: MyRingbuf,
}

impl<App> Handler<App>
where ...
{
    // Some other library code
    func handle_data(&mut self, data: &[u8]) {
        let len = self.app.needed_length();

        // SAFETY: the app will initialize the slices before they are dropped
        let slices = unsafe { self.ringbuf.push_without_init(len); }
        // After this function, any write/read to self.ringbuf would be safe and correct
        self.app.handle(data, slices);
    }
}

trait ApplicationHandler {
    /// # Safety
    ///
    /// The implementer must ensure that the slices are initialized before returning
    unsafe fn handle(&mut self, data: &[u8], slices: (&mut [MaybeUninit<T>], &mut [MaybeUninit<T>]));
}

Here's my ringbuf, ideally as a different library (a dependency of the top library):

impl MyRingbuf {
    /// Advances the write_index as if `count` elements were pushed, without initializing them
    ///
    /// # Safety
    ///
    /// The caller must ensure that:
    /// - The returned slices are initialized before they are dropped
    unsafe push_without_init(&mut self, count: usize) -> (&mut [MaybeUninit<T>], &mut [MaybeUninit<T>]) { ... }
}

The app is doing some network processing, so it is desired that the libraries behind Handler and MyRingbuf have a small memory and processing overhead.

The app might copy the data into the slices, or do some modifications to it, then write it. I could've allocated a Vec in handle_data instead, give a mut ref to the ApplicationHandler, then push that to the ringbuf, but this would've meant making two copies, which I wanted to avoid.

However, now, the user of the library has a very big responsibility: to successfully initialize the slices. This might be fine, until you realize that a panic anywhere in the handle function would trigger a drop in Handler::ringbuf, which calls drop for all the elements between the read index and the write index, which includes the uninitialized slices, causing UB because of MaybeUninit::assume_init_drop. I think this shifts a lot of weight to the implementer of ApplicationHandler.

There are multiple ways of shifting the weight back to one of the libraries:

  • In MyRingBuf replace with two functions: one that gives the slices without increasing the index, another that commits the change. Now it's the top library's responsibility to commit after the data has been initialized by the app. Now, a panic in the application won't cause a UB, because the uncommited slices won't be dropped. There's still UB if the application
    doesn't initialize the slices, but at least, the "UB surface" is reduced now.
  • Instead of passing slice references to ApplicationHandler::handle, pass MyRingbuf or a view of it. This preserves the zero-copy, and removes the necessity of unsafe in the application. However it might lead to an API that's not uniform. In my top library, I actually have handlers at multiple layers, and all of them give access to some &mut [T], which makes a very easy to use API. That, and, there might be a higher chance that the ringbuf implementation is going to change.
  • Give up MaybeUninit and use a ringbuf that restricts the data to be Default or have a valid decoding of zero bytes.
  • Never drop the elements in the ringbuf. This is not UB, it just leaks memory.

My question to this forum is less about how to solve this particular problem. This is not specific to ringbuf, but to any generic container that uses MaybeUninit and wants to support zero-copy with zero overhead. Also, depending on the design of MyRingbuf some of its needs and particularities might leak through the top library to the ApplicationHandler, which is not ideal. My question is: How do I design a good library around these constraints? What is the first thing I should give up? Do you know a successful example?

  • Would you use a ringbuf that doesn't support MaybeUninit?
  • Would you use any kind of unsafe API? What about one that fails at panic?
  • Would you like to receive &mut [T], or specialized containers?

This is my first time writing complicated unsafe code, so any suggestion is welcomed.

2
  1. Catch the panic in Handler::handle_data.
  2. Create unsafe fn MyRingbuf::push_with(&mut self, n: usize, fill_callback: impl FnOnce(&mut [MaybeUninit<T>])).
3
  1. Yeah, probably the least-effort solution. I wonder if there are some gotchas about panic catching.
  2. From a state-machine pov, I think this one is close to having two functions in MyRingbuf with the "2-phase" commit way. I'm still not sure who would implement fill_callback - probably the Handler.
4

Hello, we are creating app that uses the same component. In my case, I am creating ring buffer based lock free MPMC channel with batch support. In my app, I do this

Because MaybeUninit does not implement Drop, the drop becomes the user's responsibility if do not want to leak memory on panic. By using written index, it can be obtained by subtracting the reader index with writer index, then in impl Drop I call assume_init_drop only between those index. This ensures I never touch the uninitialized indices. MaybeUninit allows us to drop specific indices, it does not have to be all indices

For reading, I do the same, by calling assume_init and creating a slice only between the written index

For writing, I return MaybeUninit mutable reference, then the user have to use .write() method. If the user use assignment with pointer dereference, because it tries to drop the old value which can be uninitialized it is UB (I think this is unavoidable if want to zero copy), whereas .write() overrides the old value (which makes it the user's responsibility to drop the memory if the value is a memory address metadata pointing to heap). What I do to handle this is, I call assume_init_drop on all values within the indices requested by the writer before writing (only the one that is within the written index), then I can write to the slice without memory leak. I do the drop earlier, not during writing to allow the writing to auto SIMD. But in my current design, the user must pass the total of how much they wrote to the MaybeUninit, passing bigger total than how many write is actually happened becomes UB

Edit : I am also looking for design that still allows zero copy but without risk in user space code :<. Right now I am thinking about providing macro that will expand to .write() and counter increment by 1 then I will use this counter value as the total writing that the user does, to remove the need of passing total write manually, but I am still making sure it does not prevent auto SIMD

The more I think more about my implementation, the more question pop up :>. I realised std method is still method, which should be the same like user space method in capability of handling copying or not. The .write() method in MaybeUninit is

pub const fn write(&mut self, val: T) -> &mut T {

}

Which means T can be copied (stack or stack metadata of heap, heap is moved) because it is cross method. But I do not yet for sure, will it be copy if entire call is inlined :thinking:

5

From what I see you can do another thing: leak written T’s by returning not a pair of chunks, but something which updates length only when dropped unless dropped when panicking. Simplified code (this is probably not how you actually increase rigbuffer length):

struct Slices<'a, T> {
    ringbuf_len: &'a mut usize,
    first_chunk: &'a mut [MaybeUninit<T>],
    second_chunk: &'a mut [MaybeUninit<T>],
}
impl<'a, T> Drop for Slices<'a, T> {
    fn drop(&mut self) {
        if std::thread::panicking() {
            // Leak T’s already written to avoid assuming T’s not yet written are initialized.
        } else {
            *self.ringbuf_len += self.first_chunk.len() + self.second_chunk.len();
        }
    }
}
impl<'a, T> Slices<'a, T> {
    fn first_chunk<'b>(&'b mut self) -> &'b mut [MaybeUninit<T>] {
        self.first_chunk
    }
    …
}

I would still not think this API is reasonable because unsafety of push_without_init function is questionable: your variant essentially requires for returned slices to be initialized before the call which is impossible (which you clearly understand based on the question), my variant makes push_without_init entirely safe and what is unsafe to call is drop which cannot be marked as such.

Maybe the following makes more sense: make push_without_init safe (and return my Slices struct), but in place of updating length in drop (in fact, remove impl Drop) make length update happen in impl<…> Slices<…> { unsafe fn finish(self). This has a bonus of not requiring std.

Note that this is actually your first suggestion of splitting functions (I personally would actually go with it), but with an additional convenience of not having caller store length.

In all cases, make not leaking memory headache of whoever wrote panicking function, but do not make panicking at the wrong moment UB.

As for the three last questions:

Would you use a ringbuf that doesn't support MaybeUninit?

If ringbuf does not use MaybeUninit in some form internally then it will probably not have good performance. And will probably leak this fact somewhere in their API making me avoid this.

Ringbuf providing MaybeUninit buffers is yet to become important in my projects.

Would you use any kind of unsafe API? What about one that fails at panic?

Unsafe APIs are fine. Putting extra non-trivial requirements on callers is a major inconvenience: I would equate such API to the cancellation safety requirement on selects: that is I will avoid using select just for that unless I really cannot get away without it.

Would you like to receive &mut [T], or specialized containers?

Depends on what “specialized” means.

6

This sounds like a reasonable trade off. But it means you have to update the written index after you have written to the buffer.

Oh, I see. That might be a reasonable trade off. If the user uses primitive data types (like u8, u16, f64 etc.), they won't even need to drop the memory. But it becomes a responsibility once you need something like Vec.

I guess this is done by the user, not by the ring buffer library.

I feel like this is a trade off I wouldn't do. Now, the user also has to return how much they have written, making the API a little more complicated. But this seems to be an invariant of MaybeUninint - someone has to say at some point "there are new elements here, so in the future you can assume that they're initialized".

Have you thought of using a proxy object? Instead of a MaybeUninint, have something called RingBufferSliceWriter, with a custom write method, and pass a mutable reference of that object to the user.

I'm not sure if inlining gives you less buffer copies. Maybe in this case it does. But I would anyway provide a function that does batch writing, like:

fn write_slice(&mut self, vals: &[T]) -> &mut [T] {}
7

Yeah. I updated that I do not need to save written index independently because I already do pure substraction in other code to get the written value

Tail - head == total written

The value passed by user is still used in different part for updating the producer's tail atomic index. Because the writing is done via mutable reference. It does not update the producer's index, so it has to be updated separately :<. I do not create writer abstraction that update the index implisitely because the writer will be called multiple times, that means the atomic operation is done multiple times which is overhead. So I have this design

let mut counter = 0
loop {
write()
counter += 1
}
update the atomic index once here using the counted value. This is so important, because without that 10.000 writes will translate to additional 10.000x atomic operation. With this approach, it is reduced to -> total atomic operation = total_data / batch_size. More bigger batch size, more fewer atomic, more higher performance

No, in my understanding it needs to be done in the library because user does not has access to the MaybeUninit that you use

Irealized I don't need to call assume init drop before writing. Currsntly I'm using assume init ref in consumer, I realized it does not re activate drop. So I change it to .assume_init_read(). I read it re activate drop, so I can remove the drop before write now because it is now unnecerery. With this my new design, I remove yet another unnececery operation :>. But I still place assume_init_drop in the impl drop of the buffer, to prevent memory leak when panic happen

I updated it that it is not to update written index now. The index is changed to all tail - head. But it is needed to update the producer's atomic usize, to reduce the total of atomic operation significantly because it can be scalled easily by increasing the batch size

What will the struct writer does in that design? The counter passed by user is also needed to optimize the amount of atomic operation. While only user know how many times they write in each round trip, so it is needed anyway :(. Previously I automate this by using the total write == batch_size. But it turns out to be not suitable in practice. Because in practice the data can be not divisable by batch_size, this makes the atomic index is updated with incorrect total write in the last writing, that made bug earlier :<. So now I make it flexible, user decide themself how many times they write

If I create writer abstraction that takes clossure and contains update the index +1. If the closure if FnOnce, then it comes back to unoptimized atomic update, because the index update is called at the same time with the writing, if the writing happens 10.000x, the atomic update also happens 10.000x. If the clossure is Fn, and expect user to pass loop, then only user knows how many write they do inside that loop. This makes me come to this design to user pass the total write they do to do atomic operation optimization and prevent the incorrect atomic index update bug when total_data % batch_size != 0

Thank you! You give me idea. I did not remember I can use slice since slice has .len() method. So, I can use the slice len to automatically update the atomic index, while the total writing happened also match it since I just need to iterate the slice until the end. It will also be compiled to SIMD just like the previous one

Though I have question. What is the use of returnjng &mut slice there if the T is the value that will be written to the buffer?

Another question is. It is cross function that receive reference, in order to write the value to the buffer I will need to dereference the reference like ptr.write(*val). Will the dereference will copy it to the writer function's stack first? I will put inline_always to the writer function anyway because it is hot function that will be called thousand of times because the writing can happen thousand or million time if it is IOT sensor data. I will see the assembly to make sure it :>

8

You are looking for std::ptr::copy_nonoverlapping.

9

Yeahhh I forgot that method :<

While the assembly will look same, both will be SIMD loop if len > SIMD lane, and common scalar if not. Memcpy will pick best SIMD at runtime when len is not known. In target-native, both is compiled to inlined SIMD, no function call anymore (but this has downside, can only operate at maximum performance in the CPU where it is compiled. Suitable for self hosted app, not suitable for generic app run on multi machine like game. To make target-native achieve max performance on multi target, it will need to compile separate build for each CPU target)

10

Hmm, now that I see that, I'm thinking of using that logic in the drop function of the ring buffer. And, if the thread is panicking, don't drop anything.

I think this idea is an instance of the "proxy objects" subclass. Meaning that you use an intermediary object between the user and the data structure. However, I find it special that it uses &mut ringbuf_len. At this point, my only concern with this kind of object is that it's "viral". If my code, that uses the ring buffer, is a library as well and I need to delegate the write, I need to expose that proxy object (aka struct Slice) to the user.

I agree, this sounds easier to use. It still has the advantage that the ring buffer itself doesn't hold any extra data. And, only Slices has the necessary data to make the safe updates, once they're done.

Yes, I agree. My particular use case is "create once, write a lot", but since my question was more generic, I wanted to see people's opinion on this.

I meant "proxy objects" (didn't have this term when I asked the question), which is related to the second solution in my question (Instead of passing slice ...). I would say that your Slice solution is a specialized container. The writer needs to know how about that container, even if the writer doesn't know that it writes to MyRingbuf. In contrast, &mut [T] is very common.

Thanks for your answer!