use std::cell::RefCell;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::thread::LocalKey;
use dashmap::mapref::entry::Entry as DashEntry;
use dashmap::DashMap;
use fbinit::FacebookInit;
use crate::stat_types::BoxHistogram;
use crate::stat_types::BoxLocalCounter;
use crate::stat_types::BoxLocalHistogram;
use crate::stat_types::BoxLocalTimeseries;
use crate::stat_types::BoxSingletonCounter;
use crate::stat_types::Counter;
use crate::stat_types::Histogram;
use crate::stat_types::SingletonCounter;
use crate::stat_types::Timeseries;
pub struct DynamicStat<T, TStatType> {
map: RefCell<HashMap<String, TStatType>>,
key_generator: fn(&T) -> String,
stat_generator: fn(&str) -> TStatType,
}
impl<T, TStatType> DynamicStat<T, TStatType> {
pub fn new(key_generator: fn(&T) -> String, stat_generator: fn(&str) -> TStatType) -> Self {
DynamicStat {
map: RefCell::new(HashMap::new()),
key_generator,
stat_generator,
}
}
fn get_or_default<F, V>(&self, args: T, cb: F) -> V
where
F: FnOnce(&TStatType) -> V,
{
let key = (self.key_generator)(&args);
let mut map = self.map.borrow_mut();
match map.entry(key) {
Entry::Occupied(occ) => cb(occ.get()),
Entry::Vacant(vac) => {
let stat = (self.stat_generator)(vac.key());
cb(vac.insert(stat))
}
}
}
}
pub struct DynamicStatSync<T, TStatType> {
map: DashMap<String, TStatType>,
key_generator: fn(&T) -> String,
stat_generator: fn(&str) -> TStatType,
}
impl<T, TStatType> DynamicStatSync<T, TStatType> {
pub fn new(key_generator: fn(&T) -> String, stat_generator: fn(&str) -> TStatType) -> Self {
Self {
map: DashMap::new(),
key_generator,
stat_generator,
}
}
fn get_or_default<F, V>(&self, args: T, cb: F) -> V
where
F: FnOnce(&TStatType) -> V,
{
let key = (self.key_generator)(&args);
match self.map.entry(key) {
DashEntry::Occupied(occ) => cb(occ.get()),
DashEntry::Vacant(vac) => {
let stat = (self.stat_generator)(vac.key());
cb(&vac.insert(stat))
}
}
}
}
pub trait DynamicCounter<'a, T> {
fn increment_value(&'a self, value: i64, args: T);
}
impl<'a, T> DynamicCounter<'a, T> for DynamicStat<T, BoxLocalCounter> {
fn increment_value(&'a self, value: i64, args: T) {
self.get_or_default(args, |s| s.increment_value(value));
}
}
impl<T> DynamicCounter<'static, T> for LocalKey<DynamicStat<T, BoxLocalCounter>> {
fn increment_value(&'static self, value: i64, args: T) {
self.with(|s| s.increment_value(value, args));
}
}
pub trait DynamicTimeseries<'a, T> {
fn add_value(&'a self, value: i64, args: T);
fn add_value_aggregated(&'a self, value: i64, nsamples: u32, args: T);
}
impl<'a, T> DynamicTimeseries<'a, T> for DynamicStat<T, BoxLocalTimeseries> {
fn add_value(&'a self, value: i64, args: T) {
self.get_or_default(args, |s| s.add_value(value));
}
fn add_value_aggregated(&'a self, value: i64, nsamples: u32, args: T) {
self.get_or_default(args, |s| s.add_value_aggregated(value, nsamples));
}
}
impl<T> DynamicTimeseries<'static, T> for LocalKey<DynamicStat<T, BoxLocalTimeseries>> {
fn add_value(&'static self, value: i64, args: T) {
self.with(|s| s.add_value(value, args));
}
fn add_value_aggregated(&'static self, value: i64, nsamples: u32, args: T) {
self.with(|s| s.add_value_aggregated(value, nsamples, args));
}
}
pub trait DynamicHistogram<'a, T> {
fn add_value(&'a self, value: i64, args: T);
fn add_repeated_value(&'a self, value: i64, nsamples: u32, args: T);
fn flush(&self) {}
}
impl<'a, T> DynamicHistogram<'a, T> for DynamicStat<T, BoxLocalHistogram> {
fn add_value(&'a self, value: i64, args: T) {
self.get_or_default(args, |s| s.add_value(value));
}
fn add_repeated_value(&'a self, value: i64, nsamples: u32, args: T) {
self.get_or_default(args, |s| s.add_repeated_value(value, nsamples));
}
}
impl<'a, T> DynamicHistogram<'a, T> for DynamicStatSync<T, BoxHistogram> {
fn add_value(&'a self, value: i64, args: T) {
self.get_or_default(args, |s| s.add_value(value));
}
fn add_repeated_value(&'a self, value: i64, nsamples: u32, args: T) {
self.get_or_default(args, |s| s.add_repeated_value(value, nsamples));
}
fn flush(&self) {
for item in self.map.iter() {
item.value().flush();
}
}
}
impl<T> DynamicHistogram<'static, T> for LocalKey<DynamicStat<T, BoxLocalHistogram>> {
fn add_value(&'static self, value: i64, args: T) {
self.with(|s| s.add_value(value, args));
}
fn add_repeated_value(&'static self, value: i64, nsamples: u32, args: T) {
self.with(|s| s.add_repeated_value(value, nsamples, args));
}
}
pub trait DynamicSingletonCounter<'a, T> {
fn set_value(&'a self, fb: FacebookInit, value: i64, args: T);
fn get_value(&'a self, fb: FacebookInit, args: T) -> Option<i64>;
fn increment_value(&'a self, fb: FacebookInit, value: i64, args: T);
}
impl<'a, T> DynamicSingletonCounter<'a, T> for DynamicStat<T, BoxSingletonCounter> {
fn set_value(&'a self, fb: FacebookInit, value: i64, args: T) {
self.get_or_default(args, |s| s.set_value(fb, value));
}
fn get_value(&'a self, fb: FacebookInit, args: T) -> Option<i64> {
self.get_or_default(args, |s| s.get_value(fb))
}
fn increment_value(&'a self, fb: FacebookInit, value: i64, args: T) {
self.get_or_default(args, |s| s.increment_value(fb, value))
}
}
impl<T> DynamicSingletonCounter<'static, T> for LocalKey<DynamicStat<T, BoxSingletonCounter>> {
fn set_value(&'static self, fb: FacebookInit, value: i64, args: T) {
self.with(|s| s.set_value(fb, value, args))
}
fn get_value(&'static self, fb: FacebookInit, args: T) -> Option<i64> {
self.with(|s| s.get_value(fb, args))
}
fn increment_value(&'static self, fb: FacebookInit, value: i64, args: T) {
self.with(|s| s.increment_value(fb, value, args))
}
}