Skip to main content

async_selector/selector/
ext.rs

1//! Wrapper types for [`Selector`] that allow for customizing polling behavior.
2
3use std::{
4    pin::Pin,
5    sync::Arc,
6    task::{Context, Poll},
7};
8
9use futures::Stream;
10
11use crate::{
12    pollable::PollProxy,
13    selector::{Id, Selector},
14};
15
16/// Borrowed [`Stream`] that will poll the inner [`Selector`]
17/// and pass the extensions to the inner tasks.
18///
19/// Created with [`Selector::with_ext`].
20///
21/// **Important:** before polling the tasks with different extension types, see the wakeups [section](Selector#wakeups).
22pub struct WithExt<'s, 'e, 'emut, T, P, E: ?Sized, EMut: ?Sized> {
23    pub(super) selector: &'s mut Selector<T, P>,
24    pub(super) ext: &'e E,
25    pub(super) ext_mut: &'emut mut EMut,
26}
27
28impl<'e, T, P, E, EMut> Stream for WithExt<'_, 'e, '_, T, P, E, EMut>
29where
30    P: PollProxy<'e, T, E, EMut>,
31    E: ?Sized,
32    EMut: ?Sized,
33{
34    type Item = P::Progress;
35
36    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
37        let this = self.get_mut();
38        this.selector
39            .poll_next_inner(this.ext, this.ext_mut, |_| (), cx)
40            .map(|opt| opt.map(|(result, ())| result))
41    }
42}
43
44/// Borrowed [`Stream`] that will poll the inner [`Selector`] and attach the origin task [`Id`] to every item.
45///
46/// Created with [`Selector::with_id`].
47pub struct WithId<'s, T, P> {
48    pub(super) selector: &'s mut Selector<T, P>,
49}
50
51impl<T, P> Stream for WithId<'_, T, P>
52where
53    P: PollProxy<'static, T, (), ()>,
54{
55    type Item = (P::Progress, Id<T>);
56
57    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
58        self.get_mut()
59            .selector
60            .poll_next_inner(&(), &mut (), |task| Id(Arc::downgrade(task)), cx)
61    }
62}
63
64/// Borrowed [`Stream`] that will poll the inner [`Selector`], pass the extensions the inner tasks,
65/// and attach the origin task [`Id`] to every item.
66///
67/// Created with [`Selector::with_ext_and_id`].
68///
69/// **Important:** before polling the tasks with different extension types, see the wakeups [section](Selector#wakeups).
70pub struct WithExtAndId<'s, 'e, 'emut, T, P, E: ?Sized, EMut: ?Sized> {
71    pub(super) selector: &'s mut Selector<T, P>,
72    pub(super) ext: &'e E,
73    pub(super) ext_mut: &'emut mut EMut,
74}
75
76impl<'e, T, P, E, EMut> Stream for WithExtAndId<'_, 'e, '_, T, P, E, EMut>
77where
78    P: PollProxy<'e, T, E, EMut>,
79    E: ?Sized,
80    EMut: ?Sized,
81{
82    type Item = (P::Progress, Id<T>);
83
84    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
85        let this = self.get_mut();
86        this.selector
87            .poll_next_inner(this.ext, this.ext_mut, |task| Id(Arc::downgrade(task)), cx)
88    }
89}