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::{PollStrategy, PollWith},
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, S: PollStrategy, E: ?Sized, EMut: ?Sized> {
23    pub(super) selector: &'s mut Selector<S>,
24    pub(super) ext: &'e E,
25    pub(super) ext_mut: &'emut mut EMut,
26}
27
28impl<'e, S, E, EMut> Stream for WithExt<'_, 'e, '_, S, E, EMut>
29where
30    S: PollStrategy,
31    S: PollWith<'e, E, EMut>,
32    E: ?Sized,
33    EMut: ?Sized,
34{
35    type Item = S::Progress;
36
37    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
38        let this = self.get_mut();
39        this.selector
40            .poll_next_inner(this.ext, this.ext_mut, |_| (), cx)
41            .map(|opt| opt.map(|(result, ())| result))
42    }
43}
44
45/// Borrowed [`Stream`] that will poll the inner [`Selector`] and attach the origin task [`Id`] to every item.
46///
47/// Created with [`Selector::with_id`].
48pub struct WithId<'s, S: PollStrategy> {
49    pub(super) selector: &'s mut Selector<S>,
50}
51
52impl<S> Stream for WithId<'_, S>
53where
54    S: PollStrategy,
55    S: PollWith<'static, (), ()>,
56{
57    type Item = (S::Progress, Id<S::Pollable>);
58
59    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
60        self.get_mut()
61            .selector
62            .poll_next_inner(&(), &mut (), |task| Id(Arc::downgrade(task)), cx)
63    }
64}
65
66/// Borrowed [`Stream`] that will poll the inner [`Selector`], pass the extensions the inner tasks,
67/// and attach the origin task [`Id`] to every item.
68///
69/// Created with [`Selector::with_ext_and_id`].
70///
71/// **Important:** before polling the tasks with different extension types, see the wakeups [section](Selector#wakeups).
72pub struct WithExtAndId<'s, 'e, 'emut, S: PollStrategy, E: ?Sized, EMut: ?Sized> {
73    pub(super) selector: &'s mut Selector<S>,
74    pub(super) ext: &'e E,
75    pub(super) ext_mut: &'emut mut EMut,
76}
77
78impl<'e, S, E, EMut> Stream for WithExtAndId<'_, 'e, '_, S, E, EMut>
79where
80    S: PollStrategy,
81    S: PollWith<'e, E, EMut>,
82    E: ?Sized,
83    EMut: ?Sized,
84{
85    type Item = (S::Progress, Id<S::Pollable>);
86
87    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
88        let this = self.get_mut();
89        this.selector
90            .poll_next_inner(this.ext, this.ext_mut, |task| Id(Arc::downgrade(task)), cx)
91    }
92}
async_selector/selector/ext.rs

async_selector/selector/
ext.rs
