futures_util/stream/futures_unordered/
iter.rs

1use super::task::Task;
2use super::Arc;
3use super::FuturesUnordered;
4use core::marker::PhantomData;
5use core::pin::Pin;
6use core::ptr;
7use core::sync::atomic::Ordering::Relaxed;
8
9/// Mutable iterator over all futures in the unordered set.
10#[derive(Debug)]
11pub struct IterPinMut<'a, Fut> {
12    pub(super) task: *const Task<Fut>,
13    pub(super) len: usize,
14    pub(super) _marker: PhantomData<&'a mut FuturesUnordered<Fut>>,
15}
16
17/// Mutable iterator over all futures in the unordered set.
18#[derive(Debug)]
19pub struct IterMut<'a, Fut: Unpin>(pub(super) IterPinMut<'a, Fut>);
20
21/// Immutable iterator over all futures in the unordered set.
22#[derive(Debug)]
23pub struct IterPinRef<'a, Fut> {
24    pub(super) task: *const Task<Fut>,
25    pub(super) len: usize,
26    pub(super) pending_next_all: *mut Task<Fut>,
27    pub(super) _marker: PhantomData<&'a FuturesUnordered<Fut>>,
28}
29
30/// Immutable iterator over all the futures in the unordered set.
31#[derive(Debug)]
32pub struct Iter<'a, Fut: Unpin>(pub(super) IterPinRef<'a, Fut>);
33
34/// Owned iterator over all futures in the unordered set.
35#[derive(Debug)]
36pub struct IntoIter<Fut: Unpin> {
37    pub(super) len: usize,
38    pub(super) inner: FuturesUnordered<Fut>,
39}
40
41impl<Fut: Unpin> Iterator for IntoIter<Fut> {
42    type Item = Fut;
43
44    fn next(&mut self) -> Option<Self::Item> {
45        // `head_all` can be accessed directly and we don't need to spin on
46        // `Task::next_all` since we have exclusive access to the set.
47        let task = self.inner.head_all.get_mut();
48
49        if (*task).is_null() {
50            return None;
51        }
52
53        unsafe {
54            let head = *task;
55
56            // Moving out of the future is safe because it is `Unpin`
57            let future = (*(*head).future.get()).take().unwrap();
58
59            // Mutable access to a previously shared `FuturesUnordered` implies
60            // that the other threads already released the object before the
61            // current thread acquired it, so relaxed ordering can be used and
62            // valid `next_all` checks can be skipped.
63            let next = (*head).next_all.load(Relaxed);
64            *task = next;
65            if !next.is_null() {
66                *(*next).prev_all.get() = ptr::null_mut();
67            }
68            self.len -= 1;
69
70            let arc = Arc::from_raw(head);
71            arc.next_all.store(self.inner.pending_next_all(), Relaxed);
72            *arc.prev_all.get() = ptr::null_mut();
73            self.inner.release_task(arc);
74
75            Some(future)
76        }
77    }
78
79    fn size_hint(&self) -> (usize, Option<usize>) {
80        (self.len, Some(self.len))
81    }
82}
83
84impl<Fut: Unpin> ExactSizeIterator for IntoIter<Fut> {}
85
86impl<'a, Fut> Iterator for IterPinMut<'a, Fut> {
87    type Item = Pin<&'a mut Fut>;
88
89    fn next(&mut self) -> Option<Self::Item> {
90        if self.task.is_null() {
91            return None;
92        }
93
94        unsafe {
95            let future = (*(*self.task).future.get()).as_mut().unwrap();
96
97            // Mutable access to a previously shared `FuturesUnordered` implies
98            // that the other threads already released the object before the
99            // current thread acquired it, so relaxed ordering can be used and
100            // valid `next_all` checks can be skipped.
101            let next = (*self.task).next_all.load(Relaxed);
102            self.task = next;
103            self.len -= 1;
104            Some(Pin::new_unchecked(future))
105        }
106    }
107
108    fn size_hint(&self) -> (usize, Option<usize>) {
109        (self.len, Some(self.len))
110    }
111}
112
113impl<Fut> ExactSizeIterator for IterPinMut<'_, Fut> {}
114
115impl<'a, Fut: Unpin> Iterator for IterMut<'a, Fut> {
116    type Item = &'a mut Fut;
117
118    fn next(&mut self) -> Option<Self::Item> {
119        self.0.next().map(Pin::get_mut)
120    }
121
122    fn size_hint(&self) -> (usize, Option<usize>) {
123        self.0.size_hint()
124    }
125}
126
127impl<Fut: Unpin> ExactSizeIterator for IterMut<'_, Fut> {}
128
129impl<'a, Fut> Iterator for IterPinRef<'a, Fut> {
130    type Item = Pin<&'a Fut>;
131
132    fn next(&mut self) -> Option<Self::Item> {
133        if self.task.is_null() {
134            return None;
135        }
136
137        unsafe {
138            let future = (*(*self.task).future.get()).as_ref().unwrap();
139
140            // Relaxed ordering can be used since acquire ordering when
141            // `head_all` was initially read for this iterator implies acquire
142            // ordering for all previously inserted nodes (and we don't need to
143            // read `len_all` again for any other nodes).
144            let next = (*self.task).spin_next_all(self.pending_next_all, Relaxed);
145            self.task = next;
146            self.len -= 1;
147            Some(Pin::new_unchecked(future))
148        }
149    }
150
151    fn size_hint(&self) -> (usize, Option<usize>) {
152        (self.len, Some(self.len))
153    }
154}
155
156impl<Fut> ExactSizeIterator for IterPinRef<'_, Fut> {}
157
158impl<'a, Fut: Unpin> Iterator for Iter<'a, Fut> {
159    type Item = &'a Fut;
160
161    fn next(&mut self) -> Option<Self::Item> {
162        self.0.next().map(Pin::get_ref)
163    }
164
165    fn size_hint(&self) -> (usize, Option<usize>) {
166        self.0.size_hint()
167    }
168}
169
170impl<Fut: Unpin> ExactSizeIterator for Iter<'_, Fut> {}
171
172// SAFETY: `IterPinRef` yields `Pin<&Fut>` shared references. Sending these
173// to another thread requires `Fut: Sync` since both the sender and receiver
174// can concurrently access the same `Fut` through their shared references.
175unsafe impl<Fut: Sync> Send for IterPinRef<'_, Fut> {}
176unsafe impl<Fut: Sync> Sync for IterPinRef<'_, Fut> {}
177
178// SAFETY: we do nothing thread-local and there is no interior mutability,
179// so the usual structural `Send`/`Sync` apply.
180unsafe impl<Fut: Send> Send for IterPinMut<'_, Fut> {}
181unsafe impl<Fut: Sync> Sync for IterPinMut<'_, Fut> {}
182
183unsafe impl<Fut: Send + Unpin> Send for IntoIter<Fut> {}
184unsafe impl<Fut: Sync + Unpin> Sync for IntoIter<Fut> {}
futures_util/stream/futures_unordered/iter.rs

futures_util/stream/futures_unordered/
iter.rs
