Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow reclaiming the current allocation #686

Merged
merged 3 commits into from
Jun 28, 2024
Merged

Conversation

shahn
Copy link
Contributor

@shahn shahn commented Mar 31, 2024

This is based on #680, where it was noted that it is hard to use BytesMut without additional allocations in some circumstances.

@shahn
Copy link
Contributor Author

shahn commented Mar 31, 2024

This is a draft to get some comments on whether such an API addition would be acceptable

src/bytes_mut.rs Outdated Show resolved Hide resolved
@shahn shahn marked this pull request as ready for review April 7, 2024 12:22
src/bytes_mut.rs Outdated Show resolved Hide resolved
src/bytes_mut.rs Outdated Show resolved Hide resolved
src/bytes_mut.rs Outdated Show resolved Hide resolved
@braddunbar
Copy link
Contributor

braddunbar commented Apr 8, 2024

This looks a lot like BytesMut::reserve with the allocating parts cut out, which sounds a lot like a special case of BytesMut::try_reserve(ORIGINAL_CAPACITY). What would y'all think of a more general implementation in that vein?

@shahn
Copy link
Contributor Author

shahn commented Apr 8, 2024

I think it is hard to know what the original capacity is, generally, as it could've been influenced by lots of hidden operations before.
In addition, try_reclaim also doesn't move any bytes internally, with the goal of being very cheap to call. To me, the allocation behaviour of the entire crate is sometimes a bit hard to predict, so a separate function for it seems sensible

@braddunbar
Copy link
Contributor

I think it is hard to know what the original capacity is, generally, as it could've been influenced by lots of hidden operations before.

That's true, but presumably you know how much data you'll need to reserve regardless (otherwise dynamic allocation seems like the best strategy). So it doesn't particularly matter if you can reclaim the entire buffer, only if it can fit the data you're receiving. Or am I missing the point here?

In addition, try_reclaim also doesn't move any bytes internally, with the goal of being very cheap to call.

Very cheap relative to what though? Based on your description in #680, your only other alternative would be to allocate anyway, right? And that's definitely expensive relative to moving bytes. If that's not the alternative, what is?

@shahn
Copy link
Contributor Author

shahn commented Apr 8, 2024

My usecase is that I have a couple of BytesMuts and I would only allocate more memory for the current BytesMut if I couldn't reclaim the other. Basically, switch between the buffers, filling them up, while the data is processed on another thread. Once all the handles have been dropped I could reclaim the allocation and start filling the other buffer, etc.

@shahn
Copy link
Contributor Author

shahn commented May 13, 2024

rebased this on top of current master, to make it mergable again. But I am also not sure if it is deemed merge-worthy or not :)

There is currently a test case failure because BytesMut::advance() changed from its documented behaviour, but if this is generally acceptable I would try to reconcile that.

@braddunbar
Copy link
Contributor

In my mind, reclaiming memory is always in service of then writing a specific amount of data so I would prefer a try_reserve approach here that allows both at once.

@shahn
Copy link
Contributor Author

shahn commented May 15, 2024

Wouldn't that just be reserve? Or maybe more specifically, where would the try_ part come in?

@braddunbar
Copy link
Contributor

reserve allocates if it can't reclaim enough space for the requested number of bytes. try_reserve would return false rather than allocate, giving you an opportunity to swap buffers instead.

@shahn
Copy link
Contributor Author

shahn commented May 16, 2024

Ah, I misunderstood. Yeah, that'd work as well for me! I will update the branch when I get a chance

@Darksonn
Copy link
Contributor

Naming nit: The name try_reserve is currently used on Vec for an operation that reallocates but returns an error on allocation failure instead of calling abort. So we shouldn't add a method of that name that does something else.

This fixes tokio-rs#680, where it was noted that it is hard to use
BytesMut without additional allocations in some circumstances.
@shahn
Copy link
Contributor Author

shahn commented May 17, 2024

Implemented the suggestion, I kept the name try_reclaim based on Darksonn's suggestion to avoid try_reserve. Thanks a lot for the review and comments

src/bytes_mut.rs Outdated Show resolved Hide resolved
This has the nice side effect of making the implementation simpler with
regards to unsafe code, as it mostly follows what `reserve` does.
@shahn
Copy link
Contributor Author

shahn commented Jun 9, 2024

Thanks for the reviews :) Added a commit implementing copying bytes if that is cheap according to reserve. This also simplified the implementation because I could mostly re-use reserve_inner(). Also I added some more testing to test the behaviour when copying bytes.

@shahn shahn requested a review from Darksonn June 25, 2024 22:06
@conradludgate
Copy link

We should fix this expect when allocate is false:
https://github.com/shahn/bytes/blob/26c14f000d42f86b8642c75604a4a93dbe603e1b/src/bytes_mut.rs#L669
currently b.try_reclaim(usize::MAX) would panic. while it's maybe an extreme scenario, it's probably wrong to panic.

@conradludgate
Copy link

Under the current strategy of inlining, there's only one large instance of reclaim_inner in the resulting binary. Does it maybe make sense to add an intermediate inline target for try_reclaim_inner that can optimise to a much smaller function?

Eg:

#[inline]
pub fn try_reclaim(&mut self, additional: usize) -> bool {
    let len = self.len();
    let rem = self.capacity() - len;

    if additional <= rem {
        return true;
    }

    self.try_reclaim_inner(additional)
}

fn try_reclaim_inner(&mut self, additional: usize) -> bool {
    // inlines
    self.reserve_inner(additional, false)
}

@shahn
Copy link
Contributor Author

shahn commented Jun 27, 2024

Thank you! Implemented a fix to no longer panic when trying to reclaim large amounts of storage. I am not sure about the inlining, tbh - we could also inline the helper function at the expense of making try_reclaim a bit larger. Did you play with it already to see what would make most sense?

@conradludgate
Copy link

Did you play with it already to see what would make most sense?

I only briefly scanned some generated assembly. No benchmarks or profiling. I think keeping what is currently here is fine and won't make a considerable different to performance. If you only use try_reclaim, a smaller function might be better on the instruction cache, but if you use both try_reclaim and reserve, then that optimisation is eradicated. I don't see there being a considerable different in the runtime of either case.

Copy link
Contributor

@Darksonn Darksonn left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you. Good catch with the panic.

@Darksonn Darksonn merged commit 8cc9407 into tokio-rs:master Jun 28, 2024
15 checks passed
@shahn shahn deleted the try_reclaim branch June 28, 2024 19:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

5 participants