rust/src/libcore/task/local_data.rs

/*!

Task local data management

Allows storing boxes with arbitrary types inside, to be accessed
anywhere within a task, keyed by a pointer to a global finaliser
function. Useful for dynamic variables, singletons, and interfacing
with foreign code with bad callback interfaces.

To use, declare a monomorphic global function at the type to store,
and use it as the 'key' when accessing. See the 'tls' tests below for
examples.

Casting 'Arcane Sight' reveals an overwhelming aura of Transmutation
magic.

*/

use local_data_priv::{
    local_pop,
    local_get,
    local_set,
    local_modify
};

/**
 * Indexes a task-local data slot. The function's code pointer is used for
 * comparison. Recommended use is to write an empty function for each desired
 * task-local data slot (and use class destructors, not code inside the
 * function, if specific teardown is needed). DO NOT use multiple
 * instantiations of a single polymorphic function to index data of different
 * types; arbitrary type coercion is possible this way.
 *
 * One other exception is that this global state can be used in a destructor
 * context to create a circular @-box reference, which will crash during task
 * failure (see issue #3039).
 *
 * These two cases aside, the interface is safe.
 */
pub type LocalDataKey<T: Owned> = &fn(+v: @T);

/**
 * Remove a task-local data value from the table, returning the
 * reference that was originally created to insert it.
 */
pub unsafe fn local_data_pop<T: Owned>(
    key: LocalDataKey<T>) -> Option<@T> {

    local_pop(rt::rust_get_task(), key)
}
/**
 * Retrieve a task-local data value. It will also be kept alive in the
 * table until explicitly removed.
 */
pub unsafe fn local_data_get<T: Owned>(
    key: LocalDataKey<T>) -> Option<@T> {

    local_get(rt::rust_get_task(), key)
}
/**
 * Store a value in task-local data. If this key already has a value,
 * that value is overwritten (and its destructor is run).
 */
pub unsafe fn local_data_set<T: Owned>(
    key: LocalDataKey<T>, +data: @T) {

    local_set(rt::rust_get_task(), key, data)
}
/**
 * Modify a task-local data value. If the function returns 'None', the
 * data is removed (and its reference dropped).
 */
pub unsafe fn local_data_modify<T: Owned>(
    key: LocalDataKey<T>,
    modify_fn: fn(Option<@T>) -> Option<@T>) {

    local_modify(rt::rust_get_task(), key, modify_fn)
}

#[test]
pub fn test_tls_multitask() unsafe {
    fn my_key(+_x: @~str) { }
    local_data_set(my_key, @~"parent data");
    do task::spawn unsafe {
        assert local_data_get(my_key).is_none(); // TLS shouldn't carry over.
        local_data_set(my_key, @~"child data");
        assert *(local_data_get(my_key).get()) == ~"child data";
        // should be cleaned up for us
    }
    // Must work multiple times
    assert *(local_data_get(my_key).get()) == ~"parent data";
    assert *(local_data_get(my_key).get()) == ~"parent data";
    assert *(local_data_get(my_key).get()) == ~"parent data";
}

#[test]
pub fn test_tls_overwrite() unsafe {
    fn my_key(+_x: @~str) { }
    local_data_set(my_key, @~"first data");
    local_data_set(my_key, @~"next data"); // Shouldn't leak.
    assert *(local_data_get(my_key).get()) == ~"next data";
}

#[test]
pub fn test_tls_pop() unsafe {
    fn my_key(+_x: @~str) { }
    local_data_set(my_key, @~"weasel");
    assert *(local_data_pop(my_key).get()) == ~"weasel";
    // Pop must remove the data from the map.
    assert local_data_pop(my_key).is_none();
}

#[test]
pub fn test_tls_modify() unsafe {
    fn my_key(+_x: @~str) { }
    local_data_modify(my_key, |data| {
        match data {
            Some(@ref val) => fail ~"unwelcome value: " + *val,
            None       => Some(@~"first data")
        }
    });
    local_data_modify(my_key, |data| {
        match data {
            Some(@~"first data") => Some(@~"next data"),
            Some(@ref val)           => fail ~"wrong value: " + *val,
            None                 => fail ~"missing value"
        }
    });
    assert *(local_data_pop(my_key).get()) == ~"next data";
}

#[test]
pub fn test_tls_crust_automorestack_memorial_bug() unsafe {
    // This might result in a stack-canary clobber if the runtime fails to set
    // sp_limit to 0 when calling the cleanup extern - it might automatically
    // jump over to the rust stack, which causes next_c_sp to get recorded as
    // Something within a rust stack segment. Then a subsequent upcall (esp.
    // for logging, think vsnprintf) would run on a stack smaller than 1 MB.
    fn my_key(+_x: @~str) { }
    do task::spawn {
        unsafe { local_data_set(my_key, @~"hax"); }
    }
}

#[test]
pub fn test_tls_multiple_types() unsafe {
    fn str_key(+_x: @~str) { }
    fn box_key(+_x: @@()) { }
    fn int_key(+_x: @int) { }
    do task::spawn unsafe {
        local_data_set(str_key, @~"string data");
        local_data_set(box_key, @@());
        local_data_set(int_key, @42);
    }
}

#[test]
pub fn test_tls_overwrite_multiple_types() {
    fn str_key(+_x: @~str) { }
    fn box_key(+_x: @@()) { }
    fn int_key(+_x: @int) { }
    do task::spawn unsafe {
        local_data_set(str_key, @~"string data");
        local_data_set(int_key, @42);
        // This could cause a segfault if overwriting-destruction is done with
        // the crazy polymorphic transmute rather than the provided finaliser.
        local_data_set(int_key, @31337);
    }
}

#[test]
#[should_fail]
#[ignore(cfg(windows))]
pub fn test_tls_cleanup_on_failure() unsafe {
    fn str_key(+_x: @~str) { }
    fn box_key(+_x: @@()) { }
    fn int_key(+_x: @int) { }
    local_data_set(str_key, @~"parent data");
    local_data_set(box_key, @@());
    do task::spawn unsafe { // spawn_linked
        local_data_set(str_key, @~"string data");
        local_data_set(box_key, @@());
        local_data_set(int_key, @42);
        fail;
    }
    // Not quite nondeterministic.
    local_data_set(int_key, @31337);
    fail;
}