@@ -116,8 +116,9 @@ use asyncio_impl::*;
116116#[ cfg( feature = "async-io" ) ]
117117mod asyncio_impl {
118118 use std:: {
119- io,
119+ future , io,
120120 os:: fd:: { BorrowedFd , RawFd } ,
121+ pin:: pin,
121122 task:: Poll ,
122123 } ;
123124
@@ -129,7 +130,7 @@ mod asyncio_impl {
129130 impl Impl {
130131 pub fn new ( fd : RawFd ) -> io:: Result < Self > {
131132 let fd = unsafe { BorrowedFd :: borrow_raw ( fd) } ;
132- Async :: new ( fd) . map ( Self )
133+ Async :: new_nonblocking ( fd) . map ( Self )
133134 }
134135
135136 pub async fn asyncify < T > (
@@ -138,13 +139,33 @@ mod asyncio_impl {
138139 ) -> io:: Result < T > {
139140 loop {
140141 match op ( ) {
141- Poll :: Pending => self . 0 . readable ( ) . await ?,
142+ Poll :: Pending => optimistic ( self . 0 . readable ( ) ) . await ?,
142143 Poll :: Ready ( res) => return res,
143144 }
144145 }
145146 }
146147 }
147148
149+ // This "optimization" is copied from async-io.
150+ // async-io is apparently very buggy (see smol-rs/async-io#78), so it ends up being required for
151+ // things to work right.
152+ // Specifically, the `.readable()` future is permanently `Pending`, even after the reactor
153+ // schedules the future again, so `asyncify` would just never complete.
154+ async fn optimistic ( fut : impl Future < Output = io:: Result < ( ) > > ) -> io:: Result < ( ) > {
155+ let mut polled = false ;
156+ let mut fut = pin ! ( fut) ;
157+
158+ future:: poll_fn ( |cx| {
159+ if !polled {
160+ polled = true ;
161+ fut. as_mut ( ) . poll ( cx)
162+ } else {
163+ Poll :: Ready ( Ok ( ( ) ) )
164+ }
165+ } )
166+ . await
167+ }
168+
148169 #[ cfg( test) ]
149170 pub struct Runtime ;
150171
@@ -169,14 +190,95 @@ pub struct Impl;
169190#[ cfg( doc) ]
170191pub struct Runtime ;
171192
172- /// Calls `f` with an instance of the selected async runtime.
173- ///
174- /// Allows writing async-runtime-agnostic tests.
175- ///
176- /// The only supported API is `runtime.block_on(future)`.
177193#[ cfg( test) ]
178- pub fn with_runtime < R > ( f : impl FnOnce ( & Runtime ) -> io:: Result < R > ) -> io:: Result < R > {
179- let rt = Runtime :: new ( ) ?;
180- let _guard = rt. enter ( ) ;
181- f ( & rt)
194+ pub mod test {
195+ use std:: { fmt, future, panic:: resume_unwind, pin:: pin, sync:: mpsc, thread} ;
196+
197+ use super :: * ;
198+
199+ pub struct AsyncTest < F , U > {
200+ future : F ,
201+ unblocker : U ,
202+ allowed_polls : usize ,
203+ }
204+
205+ impl < F , U > AsyncTest < F , U > {
206+ pub fn new ( future : F , unblocker : U ) -> Self {
207+ Self {
208+ future,
209+ unblocker,
210+ allowed_polls : 1 ,
211+ }
212+ }
213+
214+ /// Sets the number of allowed future polls after the `unblocker` has been run.
215+ ///
216+ /// By default, this is 1, expecting the future to complete immediately after the waker has
217+ /// been notified.
218+ /// Higher values may be needed if the API-under-test is system-global and may have to
219+ /// process some irrelevant events until it becomes `Ready`.
220+ pub fn allowed_polls ( mut self , allowed_polls : usize ) -> Self {
221+ self . allowed_polls = allowed_polls;
222+ self
223+ }
224+
225+ /// Polls `future`, expecting `Poll::Pending`. Then runs `unblocker`, and expects the waker to
226+ /// be invoked and the `future` to be `Poll::Ready`.
227+ pub fn run < T > ( self ) -> io:: Result < T >
228+ where
229+ F : Future < Output = io:: Result < T > > + Send ,
230+ F :: Output : Send ,
231+ U : FnOnce ( ) -> io:: Result < ( ) > ,
232+ T : fmt:: Debug ,
233+ {
234+ let ( sender, recv) = mpsc:: sync_channel ( 0 ) ;
235+ thread:: scope ( |s| {
236+ let h = s. spawn ( move || -> io:: Result < _ > {
237+ let rt = Runtime :: new ( ) ?;
238+ let _guard = rt. enter ( ) ;
239+ let mut fut = pin ! ( self . future) ;
240+ let mut poll_count = 0 ;
241+
242+ rt. block_on ( future:: poll_fn ( |cx| {
243+ if poll_count == 0 {
244+ match fut. as_mut ( ) . poll ( cx) {
245+ Poll :: Ready ( val) => {
246+ panic ! ( "expected future to be `Pending`, but it is `Ready({val:?})`" )
247+ }
248+ Poll :: Pending => {
249+ // Waker is now scheduled to be woken when the event of interest occurs.
250+ println ! ( "future is pending; scheduling wakeup" ) ;
251+ poll_count += 1 ;
252+ sender. send ( ( ) ) . unwrap ( ) ;
253+ return Poll :: Pending ;
254+ }
255+ }
256+ } else {
257+ // This is called when the `Waker` has been woken up.
258+ match fut. as_mut ( ) . poll ( cx) {
259+ Poll :: Ready ( out) => Poll :: Ready ( out) ,
260+ Poll :: Pending => {
261+ if poll_count >= self . allowed_polls {
262+ panic ! ( "future still `Pending` after {poll_count} polls" ) ;
263+ }
264+ poll_count += 1 ;
265+ Poll :: Pending
266+ }
267+ }
268+ }
269+ } ) )
270+ } ) ;
271+
272+ recv. recv ( ) . unwrap ( ) ;
273+
274+ // We've been signaled to invoke `unblocker`.
275+ ( self . unblocker ) ( ) ?;
276+
277+ match h. join ( ) {
278+ Ok ( res) => res,
279+ Err ( payload) => resume_unwind ( payload) ,
280+ }
281+ } )
282+ }
283+ }
182284}
0 commit comments