task-local-with-timeouts.rs

use rand::Rng;
use std::default::Default;
use std::{cell::RefCell, collections::HashMap, time::Duration};
use tokio::{self, select, task, task_local, time::timeout};

fn rand_gen_millis() -> (u64, u64) {
    let mut rng = rand::thread_rng();
    let t1 = rng.gen_range(0..10);
    let t2 = rng.gen_range(0..10);
    (t1, t2)
}

fn main() {
    let (t1, t2) = rand_gen_millis();
    let rt = tokio::runtime::Runtime::new().expect("failed to start async runtime");
    let produce_ten_or_one = async move {
        let handler = task::spawn(async move {
            task_local! {
                pub static WRITER: RefCell<HashMap<String, i64>>;
            }
            let request_handler = || {
                WRITER.with(|writer| {
                    *writer.borrow_mut().entry("value".to_string()).or_default() += 3;
                    writer.borrow().get("value").cloned()
                })
            };
            WRITER
                .scope(Default::default(), async move {
                    let a1 = timeout(Duration::from_millis(t1), async {
                        WRITER.with(|writer| {
                            writer.borrow_mut().insert("value".to_string(), 1);
                            writer.borrow().get("value").cloned()
                        })
                    });
                    let a2 = timeout(Duration::from_millis(t2), async {
                        WRITER.with(|writer| {
                            writer.borrow_mut().insert("value".to_string(), 7);
                        });
                        WRITER.with(|writer| {
                            let v = writer.borrow().get("value").cloned();
                            assert_eq!(v, Some(7));
                        });
                        request_handler()
                    });
                    select! {
                        v1 = a1 => {
                            v1.ok().flatten()
                        }
                        v2 = a2 => {
                            v2.ok().flatten()
                        }
                    }
                })
                .await
        });
        let result = handler.await.expect("failed to execute handler");
        dbg!(result);
        assert!(result == Some(10) || result == Some(1));
    };
    rt.block_on(produce_ten_or_one);
}