Implement a task_local! macro

Intended to be very similar to the thread_local!() macro and is useful for
storing data for the lifetime of a future.

Author: alexcrichton

Cargo.toml

@@ -15,8 +15,7 @@ composability, and iterator-like interfaces.
 
 [dependencies]
 log = { version = "0.3", default-features = false }
-typemap = { version = "0.3.3", optional = true }
 
 [features]
-use_std = ["typemap"]
+use_std = []
 default = ["use_std"]

src/lib.rs

@@ -203,8 +203,6 @@ pub use select::{Select, SelectNext};
 pub use then::Then;
 
 if_std! {
-    extern crate typemap;
-
     mod lock;
     mod slot;
     pub mod task;

src/task/data.rs

@@ -0,0 +1,128 @@
+use std::prelude::v1::*;
+
+use std::any::TypeId;
+use std::cell::RefCell;
+use std::hash::{BuildHasherDefault, Hasher};
+use std::collections::HashMap;
+
+/// 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.
+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)();
+        super::with(|_, data| {
+            let raw_pointer = {
+                let mut data = data.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)
+            }
+        })
+    }
+}

src/task/mod.rs

@@ -26,24 +26,24 @@
 //! function is similar to the standard library's `thread::park` method where it
 //! returns a handle to wake up a task at a later date (via an `unpark` method).
 
-mod unpark_mutex;
-mod task_rc;
-pub use self::task_rc::TaskRc;
-
 use {BoxFuture, Poll};
 
 use std::prelude::v1::*;
 
-use std::thread;
-use std::cell::{Cell, RefCell};
+use std::cell::Cell;
 use std::sync::Arc;
 use std::sync::atomic::{Ordering, AtomicUsize, ATOMIC_USIZE_INIT};
+use std::thread;
 
 use self::unpark_mutex::UnparkMutex;
 
-use typemap::{SendMap, TypeMap};
+mod unpark_mutex;
+mod task_rc;
+mod data;
+pub use self::task_rc::TaskRc;
+pub use self::data::LocalKey;
 
-thread_local!(static CURRENT_TASK: Cell<(*const Task, *const TaskData)> = {
+thread_local!(static CURRENT_TASK: Cell<(*const Task, *const data::LocalMap)> = {
     Cell::new((0 as *const _, 0 as *const _))
 });
 
@@ -57,10 +57,10 @@ fn fresh_task_id() -> usize {
     return id
 }
 
-fn set<F, R>(task: &Task, data: &TaskData, f: F) -> R
+fn set<F, R>(task: &Task, data: &data::LocalMap, f: F) -> R
     where F: FnOnce() -> R
 {
-    struct Reset((*const Task, *const TaskData));
+    struct Reset((*const Task, *const data::LocalMap));
     impl Drop for Reset {
         fn drop(&mut self) {
             CURRENT_TASK.with(|c| c.set(self.0));
@@ -74,7 +74,7 @@ fn set<F, R>(task: &Task, data: &TaskData, f: F) -> R
     })
 }
 
-fn with<F: FnOnce(&Task, &TaskData) -> R, R>(f: F) -> R {
+fn with<F: FnOnce(&Task, &data::LocalMap) -> R, R>(f: F) -> R {
     let (task, data) = CURRENT_TASK.with(|c| c.get());
     assert!(!task.is_null(), "no Task is currently running");
     assert!(!data.is_null());
@@ -148,49 +148,6 @@ pub fn with_unpark_event<F, R>(event: UnparkEvent, f: F) -> R
     })
 }
 
-/// Access the current task's local data.
-///
-/// Each task has its own set of local data, which is required to be `Send` but
-/// not `Sync`. Futures within a task can access this data at any point.
-///
-/// # Panics
-///
-/// This function will panic if a task is not currently being executed. That
-/// is, this method can be dangerous to call outside of an implementation of
-/// `poll`.
-///
-/// It will also panic if called in the context of another `with_local_data` or
-/// `with_local_data_mut`.
-pub fn with_local_data<F, R>(f: F) -> R
-    where F: FnOnce(&SendMap) -> R
-{
-    with(|_, data| {
-        f(&*data.borrow())
-    })
-}
-
-/// Access the current task's local data, mutably.
-///
-/// Each task has its own set of local data, which is required to be `Send` but
-/// not `Sync`. Futures within a task can access this data at any point.
-///
-/// # Panics
-///
-/// This function will panic if a task is not currently being executed. That
-/// is, this method can be dangerous to call outside of an implementation of
-/// `poll`.
-///
-/// It will also panic if called in the context of another `with_local_data` or
-/// `with_local_data_mut`.
-pub fn with_local_data_mut<F, R>(f: F) -> R
-    where F: FnOnce(&mut SendMap) -> R
-{
-    with(|_, data| {
-        f(&mut *data.borrow_mut())
-    })
-}
-
-
 /// A handle to a "task", which represents a single lightweight "thread" of
 /// execution driving a future to completion.
 ///
@@ -222,11 +179,9 @@ enum TaskKind {
     Local(thread::Thread),
 }
 
-type TaskData = RefCell<SendMap>;
-
 struct MutexInner {
     future: BoxFuture<(), ()>,
-    task_data: TaskData,
+    task_data: data::LocalMap,
 }
 
 /// A task intended to be run directly on a local thread.
@@ -239,7 +194,7 @@ struct MutexInner {
 /// `std::thread::Thread::unpark` for the thread that entered the task.
 pub struct ThreadTask {
     id: usize,
-    task_data: TaskData,
+    task_data: data::LocalMap,
 }
 
 /// A handle for running an executor-bound task.
@@ -254,7 +209,7 @@ impl ThreadTask {
     pub fn new() -> ThreadTask {
         ThreadTask {
             id: fresh_task_id(),
-            task_data: RefCell::new(TypeMap::custom()),
+            task_data: data::local_map(),
         }
     }
 
@@ -364,13 +319,13 @@ impl Task {
     ///
     /// Does not actually begin task execution; use the `unpark` method to do
     /// so.
-    pub fn new() -> Task {
+    pub fn new(exec: Arc<Executor>, future: BoxFuture<(), ()>) -> Task {
         Task {
             id: fresh_task_id(),
             kind: TaskKind::Executor {
                 mutex: Arc::new(UnparkMutex::new(MutexInner {
                     future: future,
-                    task_data: RefCell::new(TypeMap::custom()),
+                    task_data: data::local_map(),
                 })),
                 exec: exec
             },
@@ -408,15 +363,6 @@ impl Task {
             TaskKind::Local(ref thread) => thread.unpark()
         }
     }
-
-    /// Determines whether the underlying future has ever polled with a final
-    /// result (or panicked), and thus terminated.
-    pub fn is_done(&self) -> bool {
-        match self.kind {
-            TaskKind::Executor { ref mutex, .. } => mutex.is_complete(),
-            _ => false,
-        }
-    }
 }
 
 impl Events {

src/task/unpark_mutex.rs

@@ -141,10 +141,4 @@ impl<D> UnparkMutex<D> {
     pub unsafe fn complete(&self) {
         self.status.store(COMPLETE, SeqCst);
     }
-
-
-    /// Determine whether task execution has completed.
-    pub fn is_complete(&self) -> bool {
-        self.status.load(SeqCst) == COMPLETE
-    }
 }