| use std::prelude::v1::*; |
| |
| use std::any::TypeId; |
| use std::cell::RefCell; |
| use std::hash::{BuildHasherDefault, Hasher}; |
| use std::collections::HashMap; |
| |
| use task_impl::with; |
| |
| /// A macro to create a `static` of type `LocalKey` |
| /// |
| /// This macro is intentionally similar to the `thread_local!`, and creates a |
| /// `static` which has a `with` method to access the data on a task. |
| /// |
| /// The data associated with each task local is per-task, so different tasks |
| /// will contain different values. |
| #[macro_export] |
| macro_rules! task_local { |
| (static $NAME:ident: $t:ty = $e:expr) => ( |
| static $NAME: $crate::task::LocalKey<$t> = { |
| fn __init() -> $t { $e } |
| fn __key() -> ::std::any::TypeId { |
| struct __A; |
| ::std::any::TypeId::of::<__A>() |
| } |
| $crate::task::LocalKey { |
| __init: __init, |
| __key: __key, |
| } |
| }; |
| ) |
| } |
| |
| pub type LocalMap = RefCell<HashMap<TypeId, |
| Box<Opaque>, |
| BuildHasherDefault<IdHasher>>>; |
| |
| pub fn local_map() -> LocalMap { |
| RefCell::new(HashMap::default()) |
| } |
| |
| pub trait Opaque: Send {} |
| impl<T: Send> Opaque for T {} |
| |
| /// A key for task-local data stored in a future's task. |
| /// |
| /// This type is generated by the `task_local!` macro and performs very |
| /// similarly to the `thread_local!` macro and `std::thread::LocalKey` types. |
| /// Data associated with a `LocalKey<T>` is stored inside of a future's task, |
| /// and the data is destroyed when the future is completed and the task is |
| /// destroyed. |
| /// |
| /// Task-local data can migrate between threads and hence requires a `Send` |
| /// bound. Additionally, task-local data also requires the `'static` bound to |
| /// ensure it lives long enough. When a key is accessed for the first time the |
| /// task's data is initialized with the provided initialization expression to |
| /// the macro. |
| #[derive(Debug)] |
| pub struct LocalKey<T> { |
| // "private" fields which have to be public to get around macro hygiene, not |
| // included in the stability story for this type. Can change at any time. |
| #[doc(hidden)] |
| pub __key: fn() -> TypeId, |
| #[doc(hidden)] |
| pub __init: fn() -> T, |
| } |
| |
| pub struct IdHasher { |
| id: u64, |
| } |
| |
| impl Default for IdHasher { |
| fn default() -> IdHasher { |
| IdHasher { id: 0 } |
| } |
| } |
| |
| impl Hasher for IdHasher { |
| fn write(&mut self, _bytes: &[u8]) { |
| // TODO: need to do something sensible |
| panic!("can only hash u64"); |
| } |
| |
| fn write_u64(&mut self, u: u64) { |
| self.id = u; |
| } |
| |
| fn finish(&self) -> u64 { |
| self.id |
| } |
| } |
| |
| impl<T: Send + 'static> LocalKey<T> { |
| /// Access this task-local key, running the provided closure with a |
| /// reference to the value. |
| /// |
| /// This function will access this task-local key to retrieve the data |
| /// associated with the current task and this key. If this is the first time |
| /// this key has been accessed on this task, then the key will be |
| /// initialized with the initialization expression provided at the time the |
| /// `task_local!` macro was called. |
| /// |
| /// The provided closure will be provided a shared reference to the |
| /// underlying data associated with this task-local-key. The data itself is |
| /// stored inside of the current task. |
| /// |
| /// # Panics |
| /// |
| /// This function can possibly panic for a number of reasons: |
| /// |
| /// * If there is not a current task. |
| /// * If the initialization expression is run and it panics |
| /// * If the closure provided panics |
| pub fn with<F, R>(&'static self, f: F) -> R |
| where F: FnOnce(&T) -> R |
| { |
| let key = (self.__key)(); |
| with(|task| { |
| let raw_pointer = { |
| let mut data = task.map.borrow_mut(); |
| let entry = data.entry(key).or_insert_with(|| { |
| Box::new((self.__init)()) |
| }); |
| &**entry as *const Opaque as *const T |
| }; |
| unsafe { |
| f(&*raw_pointer) |
| } |
| }) |
| } |
| } |