Skip to content

queue: add drain for SegQueue#1243

Open
jafar75 wants to merge 2 commits intocrossbeam-rs:masterfrom
jafar75:drain_seg_queue
Open

queue: add drain for SegQueue#1243
jafar75 wants to merge 2 commits intocrossbeam-rs:masterfrom
jafar75:drain_seg_queue

Conversation

@jafar75
Copy link
Copy Markdown
Contributor

@jafar75 jafar75 commented Mar 13, 2026

Summary

Closes #1228.

This PR adds a drain method to SegQueue, as requested in issue #1228. The implementation uses &mut self as suggested by @taiki-e, allowing all internal operations to use push_mut() and pop_mut(), bypassing atomic operations entirely, consistent with the approach introduced in #1191.


API

pub fn drain<R: RangeBounds<usize>>(&mut self, range: R) -> Drain<'_, T>

Accepts any RangeBounds<usize>, supporting all standard range variants:

Range Behavior
drain(..) Drain all elements
drain(..n) Drain first n elements, keep the rest
drain(a..b) Skip first a, drain next b-a, keep suffix
drain(a..) Skip first a, drain everything after

Implementation Notes

&mut self and push_mut() / pop_mut()

Since drain requires exclusive access, all internal operations use exclusive versions of push and pop which avoids CAS loops, atomic loads/stores, and backoff — the same reasoning as in #1191.

Performance note

The current implementation removes elements one by one via pop_mut. This is equivalent in performance to calling pop_mut N times manually, so the primary benefit of drain here is convenience and drop safety (elements outside the range are always preserved, and unconsumed elements in the range are dropped correctly when Drain is dropped). A potential optimization would be block-level bulk removal, grabbing entire blocks at once when the drain range covers them fully. It's possible to be done inside this PR, or left for a follow-up.

Prefix handling (drain(a..b))

At Drain creation, the first a elements are popped into a temporary SegQueue<T> called prefix. The drain range is then yielded via pop_mut() on the original queue. On Drain::drop, the suffix (whatever remains in the original queue after the range) is moved into prefix, and then mem::swap is called to restore the original queue to prefix + suffix in O(1).

Drop safety

If Drain is dropped before being fully consumed, the remaining elements within the range are dropped and blocks are freed via destroy_mut. Elements outside the range are always preserved.

size_hint and ExactSizeIterator

Credit to @laycookie (the original issue author) for suggesting size_hint / ExactSizeIterator in their draft implementation shared in #1228.


Testing

New tests cover:

  • Full drain (drain(..))
  • Prefix drain (drain(..n))
  • Range drain (drain(a..b))
  • Start-only drain (drain(a..))
  • Drop mid-iteration for all variants
  • Drop counter verification (no leaks, no double-drops)
  • Block boundary cases (BLOCK_CAP = 31)
  • Empty queue
  • Range exceeding queue length
  • Empty range (drain(n..n))
  • Queue reuse after draining

@laycookie
Copy link
Copy Markdown

I feel like tests should also define what should happen in case someone runs mem::forget on the drain struct, as not doing this now might bite back later. For reference what the regular Vec does according to The Rustonomicon

So what can we do? Well, we can pick a trivially consistent state: set the Vec’s len to be 0 when we start the iteration, and fix it up if necessary in the destructor. That way, if everything executes like normal we get the desired behavior with minimal overhead. But if someone has the audacity to mem::forget us in the middle of the iteration, all that does is leak even more (and possibly leave the Vec in an unexpected but otherwise consistent state). Since we’ve accepted that mem::forget is safe, this is definitely safe. We call leaks causing more leaks a leak amplification.

https://doc.rust-lang.org/nomicon/leaking.html

@jafar75
Copy link
Copy Markdown
Contributor Author

jafar75 commented Apr 16, 2026

Thanks @laycookie for raising the mem::forget concern. I've updated the implementation so that the queue is reset to empty at Drain creation. Now if mem::forget is called, the queue is left in a consistent empty state (leaked but not corrupted). A new drain_mem_forget test covers this, skipped under Miri via #[cfg_attr(miri, ignore)] since the leaks are intentional by design.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Development

Successfully merging this pull request may close these issues.

Would it make sense to impliment Drain on SegQueue

2 participants