monoio/macros/
scoped_tls.rs

1// Copyright 2014-2015 The Rust Project Developers. See the COPYRIGHT
2// file at the top-level directory of this distribution and at
3// http://rust-lang.org/COPYRIGHT.
4//
5// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8// option. This file may not be copied, modified, or distributed
9// except according to those terms.
10
11//! Scoped thread-local storage
12//!
13//! This module provides the ability to generate *scoped* thread-local
14//! variables. In this sense, scoped indicates that thread local storage
15//! actually stores a reference to a value, and this reference is only placed
16//! in storage for a scoped amount of time.
17//!
18//! There are no restrictions on what types can be placed into a scoped
19//! variable, but all scoped variables are initialized to the equivalent of
20//! null. Scoped thread local storage is useful when a value is present for a known
21//! period of time and it is not required to relinquish ownership of the
22//! contents.
23
24#![deny(missing_docs, warnings)]
25
26use std::{cell::Cell, marker, thread::LocalKey};
27
28/// The macro. See the module level documentation for the description and examples.
29#[macro_export]
30macro_rules! scoped_thread_local {
31    ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty) => (
32        $(#[$attrs])*
33        $vis static $name: $crate::macros::scoped_tls::ScopedKey<$ty> = $crate::macros::scoped_tls::ScopedKey {
34            inner: {
35                thread_local!(static FOO: ::std::cell::Cell<*const ()> = {
36                    ::std::cell::Cell::new(::std::ptr::null())
37                });
38                &FOO
39            },
40            _marker: ::std::marker::PhantomData,
41        };
42    )
43}
44
45/// Type representing a thread local storage key corresponding to a reference
46/// to the type parameter `T`.
47///
48/// Keys are statically allocated and can contain a reference to an instance of
49/// type `T` scoped to a particular lifetime. Keys provides two methods, `set`
50/// and `with`, both of which currently use closures to control the scope of
51/// their contents.
52pub struct ScopedKey<T> {
53    #[doc(hidden)]
54    pub inner: &'static LocalKey<Cell<*const ()>>,
55    #[doc(hidden)]
56    pub _marker: marker::PhantomData<T>,
57}
58
59unsafe impl<T> Sync for ScopedKey<T> {}
60
61impl<T> ScopedKey<T> {
62    /// Inserts a value into this scoped thread local storage slot for a
63    /// duration of a closure.
64    ///
65    /// While `cb` is running, the value `t` will be returned by `get` unless
66    /// this function is called recursively inside of `cb`.
67    ///
68    /// Upon return, this function will restore the previous value, if any
69    /// was available.
70    pub fn set<F, R>(&'static self, t: &T, f: F) -> R
71    where
72        F: FnOnce() -> R,
73    {
74        struct Reset {
75            key: &'static LocalKey<Cell<*const ()>>,
76            val: *const (),
77        }
78        impl Drop for Reset {
79            fn drop(&mut self) {
80                self.key.with(|c| c.set(self.val));
81            }
82        }
83        let prev = self.inner.with(|c| {
84            let prev = c.get();
85            c.set(t as *const T as *const ());
86            prev
87        });
88        let _reset = Reset {
89            key: self.inner,
90            val: prev,
91        };
92        f()
93    }
94
95    /// Gets a value out of this scoped variable.
96    ///
97    /// This function takes a closure which receives the value of this
98    /// variable.
99    pub fn with<F, R>(&'static self, f: F) -> R
100    where
101        F: FnOnce(&T) -> R,
102    {
103        let val = self.inner.with(|c| c.get());
104        assert!(
105            !val.is_null(),
106            "cannot access a scoped thread local variable without calling `set` first"
107        );
108        unsafe { f(&*(val as *const T)) }
109    }
110
111    /// Gets a value out of this scoped variable.
112    pub fn try_with<F, R>(&'static self, f: F) -> R
113    where
114        F: FnOnce(Option<&T>) -> R,
115    {
116        let val = self.inner.with(|c| c.get());
117        if val.is_null() {
118            f(None)
119        } else {
120            unsafe { f(Some(&*(val as *const T))) }
121        }
122    }
123
124    /// Test whether this TLS key has been `set` for the current thread.
125    #[inline]
126    pub fn is_set(&'static self) -> bool {
127        self.inner.with(|c| !c.get().is_null())
128    }
129}