1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
use std::cell::Cell;
use std::future::Future;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::thread;
use std::time::{Duration, Instant};

use futures_lite::pin;
use once_cell::sync::Lazy;
use waker_fn::waker_fn;

use crate::reactor::Reactor;

/// Number of currently active `block_on()` invocations.
static BLOCK_ON_COUNT: AtomicUsize = AtomicUsize::new(0);

/// Unparker for the "async-io" thread.
static UNPARKER: Lazy<parking::Unparker> = Lazy::new(|| {
    let (parker, unparker) = parking::pair();

    // Spawn a helper thread driving the reactor.
    //
    // Note that this thread is not exactly necessary, it's only here to help push things
    // forward if there are no `Parker`s around or if `Parker`s are just idling and never
    // parking.
    thread::Builder::new()
        .name("async-io".to_string())
        .spawn(move || main_loop(parker))
        .expect("cannot spawn async-io thread");

    unparker
});

/// Initializes the "async-io" thread.
pub(crate) fn init() {
    Lazy::force(&UNPARKER);
}

/// The main loop for the "async-io" thread.
fn main_loop(parker: parking::Parker) {
    // The last observed reactor tick.
    let mut last_tick = 0;
    // Number of sleeps since this thread has called `react()`.
    let mut sleeps = 0u64;

    loop {
        let tick = Reactor::get().ticker();

        if last_tick == tick {
            let reactor_lock = if sleeps >= 10 {
                // If no new ticks have occurred for a while, stop sleeping and spinning in
                // this loop and just block on the reactor lock.
                Some(Reactor::get().lock())
            } else {
                Reactor::get().try_lock()
            };

            if let Some(mut reactor_lock) = reactor_lock {
                log::trace!("main_loop: waiting on I/O");
                reactor_lock.react(None).ok();
                last_tick = Reactor::get().ticker();
                sleeps = 0;
            }
        } else {
            last_tick = tick;
        }

        if BLOCK_ON_COUNT.load(Ordering::SeqCst) > 0 {
            // Exponential backoff from 50us to 10ms.
            let delay_us = [50, 75, 100, 250, 500, 750, 1000, 2500, 5000]
                .get(sleeps as usize)
                .unwrap_or(&10_000);

            log::trace!("main_loop: sleeping for {} us", delay_us);
            if parker.park_timeout(Duration::from_micros(*delay_us)) {
                log::trace!("main_loop: notified");

                // If notified before timeout, reset the last tick and the sleep counter.
                last_tick = Reactor::get().ticker();
                sleeps = 0;
            } else {
                sleeps += 1;
            }
        }
    }
}

/// Blocks the current thread on a future, processing I/O events when idle.
///
/// # Examples
///
/// ```
/// use async_io::Timer;
/// use std::time::Duration;
///
/// async_io::block_on(async {
///     // This timer will likely be processed by the current
///     // thread rather than the fallback "async-io" thread.
///     Timer::after(Duration::from_millis(1)).await;
/// });
/// ```
pub fn block_on<T>(future: impl Future<Output = T>) -> T {
    log::trace!("block_on()");

    // Increment `BLOCK_ON_COUNT` so that the "async-io" thread becomes less aggressive.
    BLOCK_ON_COUNT.fetch_add(1, Ordering::SeqCst);

    // Make sure to decrement `BLOCK_ON_COUNT` at the end and wake the "async-io" thread.
    let _guard = CallOnDrop(|| {
        BLOCK_ON_COUNT.fetch_sub(1, Ordering::SeqCst);
        UNPARKER.unpark();
    });

    // Parker and unparker for notifying the current thread.
    let (p, u) = parking::pair();
    // This boolean is set to `true` when the current thread is blocked on I/O.
    let io_blocked = Arc::new(AtomicBool::new(false));

    thread_local! {
        // Indicates that the current thread is polling I/O, but not necessarily blocked on it.
        static IO_POLLING: Cell<bool> = Cell::new(false);
    }

    // Prepare the waker.
    let waker = waker_fn({
        let io_blocked = io_blocked.clone();
        move || {
            if u.unpark() {
                // Check if waking from another thread and if currently blocked on I/O.
                if !IO_POLLING.with(Cell::get) && io_blocked.load(Ordering::SeqCst) {
                    Reactor::get().notify();
                }
            }
        }
    });
    let cx = &mut Context::from_waker(&waker);
    pin!(future);

    loop {
        // Poll the future.
        if let Poll::Ready(t) = future.as_mut().poll(cx) {
            log::trace!("block_on: completed");
            return t;
        }

        // Check if a notification was received.
        if p.park_timeout(Duration::from_secs(0)) {
            log::trace!("block_on: notified");

            // Try grabbing a lock on the reactor to process I/O events.
            if let Some(mut reactor_lock) = Reactor::get().try_lock() {
                // First let wakers know this parker is processing I/O events.
                IO_POLLING.with(|io| io.set(true));
                let _guard = CallOnDrop(|| {
                    IO_POLLING.with(|io| io.set(false));
                });

                // Process available I/O events.
                reactor_lock.react(Some(Duration::from_secs(0))).ok();
            }
            continue;
        }

        // Try grabbing a lock on the reactor to wait on I/O.
        if let Some(mut reactor_lock) = Reactor::get().try_lock() {
            // Record the instant at which the lock was grabbed.
            let start = Instant::now();

            loop {
                // First let wakers know this parker is blocked on I/O.
                IO_POLLING.with(|io| io.set(true));
                io_blocked.store(true, Ordering::SeqCst);
                let _guard = CallOnDrop(|| {
                    IO_POLLING.with(|io| io.set(false));
                    io_blocked.store(false, Ordering::SeqCst);
                });

                // Check if a notification has been received before `io_blocked` was updated
                // because in that case the reactor won't receive a wakeup.
                if p.park_timeout(Duration::from_secs(0)) {
                    log::trace!("block_on: notified");
                    break;
                }

                // Wait for I/O events.
                log::trace!("block_on: waiting on I/O");
                reactor_lock.react(None).ok();

                // Check if a notification has been received.
                if p.park_timeout(Duration::from_secs(0)) {
                    log::trace!("block_on: notified");
                    break;
                }

                // Check if this thread been handling I/O events for a long time.
                if start.elapsed() > Duration::from_micros(500) {
                    log::trace!("block_on: stops hogging the reactor");

                    // This thread is clearly processing I/O events for some other threads
                    // because it didn't get a notification yet. It's best to stop hogging the
                    // reactor and give other threads a chance to process I/O events for
                    // themselves.
                    drop(reactor_lock);

                    // Unpark the "async-io" thread in case no other thread is ready to start
                    // processing I/O events. This way we prevent a potential latency spike.
                    UNPARKER.unpark();

                    // Wait for a notification.
                    p.park();
                    break;
                }
            }
        } else {
            // Wait for an actual notification.
            log::trace!("block_on: sleep until notification");
            p.park();
        }
    }
}

/// Runs a closure when dropped.
struct CallOnDrop<F: Fn()>(F);

impl<F: Fn()> Drop for CallOnDrop<F> {
    fn drop(&mut self) {
        (self.0)();
    }
}