use std::cell::RefCell;
use std::collections::HashMap;
use std::sync::Arc;
#[cfg(not(feature = "parking_lot"))]
use std::sync::RwLock;
#[cfg(feature = "parking_lot")]
use parking_lot::RwLock;
type LabelValue = Arc<String>;
#[derive(Debug, Clone, Default)]
struct LabelScope {
pairs: Option<Arc<HashMap<String, LabelValue>>>,
}
impl LabelScope {
fn set(&self, key: String, value: LabelValue) -> Self {
let mut new_pairs = match self.pairs {
None => HashMap::new(),
Some(ref old_pairs) => old_pairs.as_ref().clone(),
};
new_pairs.insert(key, value);
LabelScope {
pairs: Some(Arc::new(new_pairs)),
}
}
fn unset(&self, key: &str) -> Self {
match self.pairs {
None => self.clone(),
Some(ref old_pairs) => {
let mut new_pairs = old_pairs.as_ref().clone();
if new_pairs.remove(key).is_some() {
if new_pairs.is_empty() {
LabelScope { pairs: None }
} else {
LabelScope {
pairs: Some(Arc::new(new_pairs)),
}
}
} else {
self.clone()
}
}
}
}
fn get(&self, key: &str) -> Option<LabelValue> {
match &self.pairs {
None => None,
Some(pairs) => pairs.get(key).cloned(),
}
}
fn collect(&self, map: &mut HashMap<String, LabelValue>) {
if let Some(pairs) = &self.pairs {
map.extend(pairs.as_ref().clone().into_iter())
}
}
}
lazy_static! {
static ref APP_LABELS: RwLock<LabelScope> = RwLock::new(LabelScope::default());
}
thread_local! {
static THREAD_LABELS: RefCell<LabelScope> = RefCell::new(LabelScope::default());
}
pub struct ThreadLabel;
impl ThreadLabel {
pub fn get(key: &str) -> Option<Arc<String>> {
THREAD_LABELS.with(|map| map.borrow().get(key))
}
pub fn set<S: Into<String>>(key: S, value: S) {
THREAD_LABELS.with(|map| {
let new = { map.borrow().set(key.into(), Arc::new(value.into())) };
*map.borrow_mut() = new;
});
}
pub fn unset(key: &str) {
THREAD_LABELS.with(|map| {
let new = { map.borrow().unset(key) };
*map.borrow_mut() = new;
});
}
fn collect(map: &mut HashMap<String, LabelValue>) {
THREAD_LABELS.with(|mop| mop.borrow().collect(map));
}
}
pub struct AppLabel;
impl AppLabel {
pub fn get(key: &str) -> Option<Arc<String>> {
read_lock!(APP_LABELS).get(key)
}
pub fn set<S: Into<String>>(key: S, value: S) {
let b = { read_lock!(APP_LABELS).set(key.into(), Arc::new(value.into())) };
*write_lock!(APP_LABELS) = b;
}
pub fn unset(key: &str) {
let b = { read_lock!(APP_LABELS).unset(key) };
*write_lock!(APP_LABELS) = b;
}
fn collect(map: &mut HashMap<String, LabelValue>) {
read_lock!(APP_LABELS).collect(map)
}
}
#[derive(Debug, Clone)]
pub struct Labels {
scopes: Vec<LabelScope>,
}
impl From<HashMap<String, LabelValue>> for Labels {
fn from(map: HashMap<String, LabelValue>) -> Self {
Labels {
scopes: vec![LabelScope {
pairs: Some(Arc::new(map)),
}],
}
}
}
impl Default for Labels {
#[inline]
fn default() -> Self {
Labels { scopes: vec![] }
}
}
impl Labels {
pub fn save_context(&mut self) {
self.scopes
.push(THREAD_LABELS.with(|map| map.borrow().clone()));
self.scopes.push(read_lock!(APP_LABELS).clone());
}
pub fn lookup(&self, key: &str) -> Option<LabelValue> {
fn lookup_current_context(key: &str) -> Option<LabelValue> {
ThreadLabel::get(key).or_else(|| AppLabel::get(key))
}
match self.scopes.len() {
0 => lookup_current_context(key),
1 => self.scopes[0]
.get(key)
.or_else(|| lookup_current_context(key)),
_ => {
for src in &self.scopes {
if let Some(label_value) = src.get(key) {
return Some(label_value);
}
}
None
}
}
}
pub fn into_map(mut self) -> HashMap<String, LabelValue> {
let mut map = HashMap::new();
match self.scopes.len() {
0 => {
AppLabel::collect(&mut map);
ThreadLabel::collect(&mut map);
}
1 => {
AppLabel::collect(&mut map);
ThreadLabel::collect(&mut map);
self.scopes[0].collect(&mut map);
}
_ => {
self.scopes.reverse();
for src in self.scopes {
src.collect(&mut map)
}
}
}
map
}
}
#[cfg(test)]
pub mod test {
use super::*;
use std::sync::Mutex;
lazy_static! {
static ref TEST_SEQUENCE: Mutex<()> = Mutex::new(());
}
#[test]
fn context_labels() {
let _lock = TEST_SEQUENCE.lock().expect("Test Sequence");
AppLabel::set("abc", "456");
ThreadLabel::set("abc", "123");
assert_eq!(
Arc::new("123".into()),
labels!().lookup("abc").expect("ThreadLabel Value")
);
ThreadLabel::unset("abc");
assert_eq!(
Arc::new("456".into()),
labels!().lookup("abc").expect("AppLabel Value")
);
AppLabel::unset("abc");
assert_eq!(true, labels!().lookup("abc").is_none());
}
#[test]
fn labels_macro() {
let _lock = TEST_SEQUENCE.lock().expect("Test Sequence");
let labels = labels! {
"abc" => "789",
"xyz" => "123"
};
assert_eq!(
Arc::new("789".into()),
labels.lookup("abc").expect("Label Value")
);
assert_eq!(
Arc::new("123".into()),
labels.lookup("xyz").expect("Label Value")
);
}
#[test]
fn value_labels() {
let _lock = TEST_SEQUENCE.lock().expect("Test Sequence");
let labels = labels! { "abc" => "789" };
assert_eq!(
Arc::new("789".into()),
labels.lookup("abc").expect("Label Value")
);
AppLabel::set("abc", "456");
assert_eq!(
Arc::new("789".into()),
labels.lookup("abc").expect("Label Value")
);
ThreadLabel::set("abc", "123");
assert_eq!(
Arc::new("789".into()),
labels.lookup("abc").expect("Label Value")
);
}
}