1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
use std::{
any::{Any, TypeId},
cell::RefCell,
collections::HashMap,
sync::{LazyLock, Mutex, mpsc::Sender},
};
// https://users.rust-lang.org/t/how-to-design-a-generic-thread-local/80213/6
pub struct ThreadLocalSender<T: 'static> {
mutex: &'static LazyLock<Mutex<Option<Sender<T>>>>,
}
thread_local! {
static THREAD_LOCAL_SENDER_REGISTRY: RefCell<HashMap<TypeId, Box<dyn Any>>> = RefCell::new(HashMap::new());
}
impl<T: 'static> ThreadLocalSender<T> {
pub const fn new(mutex: &'static LazyLock<Mutex<Option<Sender<T>>>>) -> Self {
ThreadLocalSender { mutex }
}
fn _try_send(&self, packet: &impl Fn() -> T, try_again: bool) -> Result<(), String> {
THREAD_LOCAL_SENDER_REGISTRY
.try_with(|registry| {
let id = TypeId::of::<T>();
let mut registry = registry.borrow_mut();
if !registry.contains_key(&id)
|| registry[&id]
.downcast_ref::<Option<Sender<T>>>()
.expect("Guaranteed by the initializer")
.is_none()
{
let initial_value = self.mutex.lock().unwrap().clone();
registry.insert(id, Box::new(initial_value));
}
registry[&id]
.downcast_ref::<Option<Sender<T>>>()
.expect("Guaranteed by the initializer")
.as_ref()
.ok_or("Empty thread-local sender mutex".to_string())?
.send(packet())
.or_else(|_| {
/*
* We end up in this case when the module has been unloaded then re-loaded.
* We don't clean up thread-local variables when the module exits. Therefore
* registry[id] might point to a Sender which is no longer connected to a
* channel.
*
* In that case we can simply remove registry[id] and try again, re-allocating
* a fresh thread-local Sender.
*/
registry.remove(&id);
drop(registry);
// Make sure we don't try again twice in a row. That would mean something else
// is broken.
if try_again {
self._try_send(packet, false)
} else {
Err("Could not send packet after re-cloning sender".to_string())
}
})
})
.unwrap_or_else(|_| {
// AccessError -> we are likely in the TLS destructor for this pthread
// Let's just fall back to using the global mutex.
self.mutex
.lock()
.map_err(|_| "Could not acquire mutex".to_string())?
.as_ref()
.ok_or("Empty sender mutex".to_string())?
.send(packet())
.map_err(|_| "Could not send packet".to_string())
})
}
pub fn try_send(&self, packet: &impl Fn() -> T) -> Result<(), String> {
self._try_send(packet, true)
}
}