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}