Line data    Source code

       1              : pub(crate) type ThreadLocalStats<T> = Arc<Mutex<T>>;
       2              : pub(crate) type AllThreadLocalStats<T> = Arc<Mutex<Vec<ThreadLocalStats<T>>>>;
       3              : 
       4              : macro_rules! declare {
       5              :     ($THREAD_LOCAL_NAME:ident: $T:ty) => {
       6              :         thread_local! {
       7              :             pub static $THREAD_LOCAL_NAME: std::cell::RefCell<crate::util::tokio_thread_local_stats::ThreadLocalStats<$T>> = std::cell::RefCell::new(
       8              :                 std::sync::Arc::new(std::sync::Mutex::new(Default::default()))
       9              :             );
      10              :         }
      11              :     };
      12              : }
      13              : 
      14              : use std::sync::{Arc, Mutex};
      15              : 
      16              : pub(crate) use declare;
      17              : 
      18              : macro_rules! main {
      19              :     ($THREAD_LOCAL_NAME:ident, $main_impl:expr) => {{
      20              :         let main_impl = $main_impl;
      21              :         let all = Arc::new(Mutex::new(Vec::new()));
      22              : 
      23              :         let rt = tokio::runtime::Builder::new_multi_thread()
      24              :             .on_thread_start({
      25              :                 let all = Arc::clone(&all);
      26            0 :                 move || {
      27            0 :                     // pre-initialize the thread local stats by accessesing them
      28            0 :                     // (some stats like requests_stats::Stats are quite costly to initialize,
      29            0 :                     //  we don't want to pay that cost during the measurement period)
      30            0 :                     $THREAD_LOCAL_NAME.with(|stats| {
      31            0 :                         let stats: Arc<_> = Arc::clone(&*stats.borrow());
      32            0 :                         all.lock().unwrap().push(stats);
      33            0 :                     });
      34            0 :                 }
      35              :             })
      36              :             .enable_all()
      37              :             .build()
      38              :             .unwrap();
      39              : 
      40              :         let main_task = rt.spawn(main_impl(all));
      41              :         rt.block_on(main_task).unwrap()
      42              :     }};
      43              : }
      44              : 
      45              : pub(crate) use main;