tentacle 0.7.5

Minimal implementation for a multiplexed p2p network framework.
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

//! Since tokio does not make the coop module public, a similar collaborative yield
//! strategy had to be implemented manually

use std::{
    cell::RefCell,
    task::{Context, Poll},
};

thread_local! {
    static CURRENT: RefCell<u8> = const { RefCell::new(128) };
}

/// Returns `Poll::Pending` if the current task has exceeded its budget and should yield.
///
/// User can insert this logic into your own implementation of future to actively yield the execution state
#[inline]
pub fn poll_proceed(cx: &mut Context<'_>) -> Poll<()> {
    CURRENT.with(|cell| {
        let mut budget = cell.borrow_mut();

        *budget -= 1;
        if *budget != 0 {
            Poll::Ready(())
        } else {
            *budget = 128;
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    })
}

#[cfg(test)]
mod test {
    use super::*;
    use futures::{executor::block_on, future::poll_fn};
    use std::thread;

    fn get() -> u8 {
        CURRENT.with(|cell| *cell.borrow())
    }

    #[test]
    fn test_budget() {
        assert_eq!(get(), 128);
        block_on(poll_fn(poll_proceed));
        assert_eq!(get(), 127);

        thread::spawn(|| {
            assert_eq!(get(), 128);
            block_on(async {
                for _ in 0..2 {
                    poll_fn(poll_proceed).await
                }
            });
            assert_eq!(get(), 126);
        })
        .join()
        .unwrap();
        assert_eq!(get(), 127);
        block_on(async {
            for _ in 0..127 {
                poll_fn(poll_proceed).await
            }
        });
        assert_eq!(get(), 127);
    }
}