How safe is this TLS abstraction?

[WorldSEnder]

Thread local storage can only be accessed while the thread is still alive, whereas the associated LocalKey can have a 'static lifetime. At least that was given to me as the reason why the API only supports LocalKey::with and LocalKey::try_with.

But can we not tie the guarantee that "the thread is still alive" to another struct that we force to be not-Send and not-Sync (so that you can't store it in a static or send it away)? Behold, then,

/// Small helper to only change lifetimes. Caller must hold up invariants for the cast
unsafe fn promote_lifetime<'a, 'b, T: ?Sized>(r: &'a T) -> &'b T {
    unsafe { &*(r as *const T) }
}
struct NaughtyLocal<'key /* always 'static atm, but allow coercions */, T: ?Sized> {
    _not_send_sync: PhantomData<*mut u8>,
    // SAFETY: only use this as &'a T when there is some &'a Self around
    the_t: &'key T,
}
impl<T: 'static> NaughtyLocal<'static, T> {
    pub fn new(key: &'static LocalKey<T>) -> Self
    {
        key.with(|value| NaughtyLocal {
            _not_send_sync: PhantomData,
            // SAFETY: The data is alive for the lifetime of the surrounding thread.
            // But, since this datatype is not send nor sync, it itself proves
            // that the thread is still alive!
            the_t: unsafe { promote_lifetime(value) },
        })
    }
}
impl<'key, T: ?Sized> Deref for NaughtyLocal<'key, T> {
    type Target = T;
    fn deref(&self) -> &T {
        self.the_t
    }
}

The reason to have this is two-fold: I always disliked how LocalKey::with forced you to nest a code-block. Secondly, that API doesn't play well with async code. Let me show:

thread_local! {
    static FRANKENSTEIN: String = String::from("It lives!");
}
fn foo() {
    FRANKENSTEIN.with(|monster| {
        println!("{monster}");
    });
    // vs (no closure, no extra scope)
    let monster = NaughtyLocal::new(&FRANKENSTEIN);
    print!("{}", *monster);
}

and, when using it with async

type SomeRefFuture = /* as an example, see async_once_cell::Lazy */;
impl<'a> IntoFuture for Pin<&'a SomeRefFuture> {
    type Output = &'a i32;
    // ... 
}
thread_local! {
    static FUT: SomeRefFuture = SomeRefFuture::new();
}
async fn bar() -> i32 {
    let future = NaughtyLocal::new(&FUT);
    // SAFETY: we promise not to move the future
    let future = unsafe { Pin::new_unchecked(future) };
    future.as_ref().await
    // This future is not Send nor Sync
}

I'm not even sure how I would write async fn bar without the abstraction (a manual poll_fn perhaps? Still, I suppose one should make sure that the future can't move to other threads).

All of this however, is on a shaky footing and depends on the safety of NaughtyLocal. Pointing out flaws, criticism of the idea and are welcome

EDIT: Playground link for completeness

[RustyYato]

It's unsound as written.

let monster = NaughtyLocal::new(&FRANKENSTEIN);
let m: &'static NaughtyLocal<String> = Box::leak(Box::new(monster));
let text: &'static str = m; // via deref coercion

Now text can be sent to another thread and outlive the current thread. Thus causing a use after free.

[WorldSEnder]

Well shux, but expected. Was too good to be true.

[WorldSEnder]

Actually, we should be able to use the old "stack marker" trick here to restrict the lifetime of the NaughtyLocal to something sensible:

pub struct StackMarkerImpl {
    _phantom: PhantomData<(*mut u8, ::std::marker::PhantomPinned)>,
}
impl StackMarkerImpl {
    /// # Safety: only okay if you use this on the stack.
    unsafe fn new() -> Self {
        Self { _phantom : PhantomData }
    }
}
type StackMarker<'s> = Pin<&'s mut StackMarkerImpl>;
macro_rules! stack_marker { 
    ($id: ident) => { 
        let mut $id = unsafe { $crate::StackMarkerImpl::new() };
        let $id: StackMarker<'_> = pin!($id);
    };
}

and make sure that such a marker is alive when you construct the value

impl<'key, T: 'static> NaughtyLocal<'key, T> {
    pub fn new(key: &'static LocalKey<T>, marker: &StackMarker<'key>) -> Self { ... }
}

Slightly less convinient to use, since you need to take care to have a stack marker around, but at least, the async example is actually implementable (I'm still not sure how I would do that).

async fn bar() -> i32 {
    stack_marker!(marker);
    let future = NaughtyLocal::new(&FUT, &marker);
    // SAFETY: we promise not to move the future
    let future = unsafe { Pin::new_unchecked(future) };
    *future.as_ref().await
}

For posteriority, again, a playground link

[RustyYato]

I thought about that but I wasn't sure how well that works with async functions/blocks. I.e. maybe there's some weird way to go from &Future -> &NaughtyLocal. And that would allow me to pull this trick one level up. But I think this isn't a real concern because I don't think you can access those locals.

So I think the stack marker approach is probably safe.

[RustyYato]

I think coroutines (#![feature(coroutines, yield_expr)]) may allow you to yield a reference to local variables in the future, so that's something to look out for. So while this is sound now, it may become unsound in the future if something like that happens. Because then we could leak the coroutine, get a reference to the NaughtyLocal<T> local (via coroutine yield), coerce it to a reference to a T: Sync type, and send it to another thread (as a &'static T).

The reason this would be allowed for coroutines is to support lending iterators, which allow yielding items that point to their local state. As a contrived example, something like this:

gen {
    let mut buffer = Vec::with_capacity(...);
    
    // We can yield references to local data because the coroutine
    // state keeps the data alive across yield points
    loop {
        buffer.clear();
        buffer.extend(some_channel.by_ref().take(...));
        // NOTE: this doesn't work right now, but it may in the future
        yield buffer.as_mut_slice();
    }
}

[WorldSEnder]

may allow you to yield a reference to local variables in the future

Definitely something to look out for, although since the StackMarker is infected with a local lifetime, I can't quite convince myself why it would be okay to yield it to the caller. You certainly can't name the 'key lifetime outside, and also can't just prolong it to the lifetime of the generator. So I don't think it's as trivial to go from a &'static mut {generator} to a NaughtyLocal<'???, T> via yield. '??? can not become 'static just because the value got yielded?

[RustyYato]

Yes, it's definitely not as trivial, and may not work out as I think it will (especially because it's not implemented yet, so I can't test it!).

I think being able to name the lifetime is a red herring, since the coroutine itself doesn't need to cross a function boundary. Also, coroutines could be adapted to use GATs to represent the Yield associated type, so we don't need to be able to name the lifetime (and this would be needed to support lending iterators). If Yield was changed to a GAT, then we could return the coroutine (impl Coroutine<Yield<'a> = &'a NaughtyLocal<'a>>, maybe something like this?).