use crate::SignalSubscription;
use crate::change::MapChange;
use crate::path::{PathRef, StorePath, Under};
use crate::primitives::signal::{SubscriptionMeta, forget, held, label};
use arc_swap::ArcSwap;
use dashmap::DashMap;
use rpds::RedBlackTreeMapSync;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use smol_str::{SmolStr, SmolStrBuilder};
use std::any;
use std::borrow::Borrow;
use std::fmt::{self, Debug, Display, Write as _};
use std::hash::Hash;
use std::panic::Location;
use std::str::FromStr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
pub trait MapEntryPath {
fn entry(&self, name: &str) -> StorePath;
}
impl MapEntryPath for StorePath {
fn entry(&self, name: &str) -> StorePath {
self.push(name)
}
}
pub type SubscriberAny<K, V> = Arc<dyn Fn(&MapChange<K, V>) + Send + Sync + 'static>;
pub type SubscriberKey<K, V> = Arc<dyn Fn(&MapChange<K, V>) + Send + Sync + 'static>;
pub enum EntryOf<K> {
Took(K),
ThePathItself,
NameWillNotRead,
NotThisMaps,
}
pub fn entry_of<K>(prefix: &StorePath, stored: PathRef<'_>) -> EntryOf<K>
where
K: ReactiveMapKey,
{
match stored.level_under(prefix) {
Under::Prefix => EntryOf::ThePathItself,
Under::Deeper(_) | Under::Outside => EntryOf::NotThisMaps,
Under::Entry(name) => match K::read(name.as_str()) {
Some(key) => EntryOf::Took(key),
None => EntryOf::NameWillNotRead,
},
}
}
pub trait ReactiveMapKey: AsRef<str> + Clone + Hash + Eq + Send + Sync + 'static {
fn read(name: &str) -> Option<Self>;
}
impl ReactiveMapKey for String {
fn read(name: &str) -> Option<Self> {
Some(name.to_string())
}
}
impl ReactiveMapKey for SmolStr {
fn read(name: &str) -> Option<Self> {
Some(SmolStr::new(name))
}
}
impl<T> ReactiveMapKey for Id<T>
where
T: Display + FromStr + Clone + Send + Sync + 'static,
{
fn read(name: &str) -> Option<Self> {
T::from_str(name).ok().map(Id::new)
}
}
#[derive(Clone)]
pub struct Id<T> {
spelled: SmolStr,
value: T,
}
impl<T: Display> Id<T> {
pub fn new(value: T) -> Self {
let mut out = SmolStrBuilder::new();
let _ = write!(out, "{value}");
Self {
spelled: out.finish(),
value,
}
}
}
impl<T> Id<T> {
pub fn get(&self) -> &T {
&self.value
}
pub fn into_inner(self) -> T {
self.value
}
}
impl<T> AsRef<str> for Id<T> {
fn as_ref(&self) -> &str {
&self.spelled
}
}
impl<T> Borrow<str> for Id<T> {
fn borrow(&self) -> &str {
&self.spelled
}
}
impl<T> PartialEq for Id<T> {
fn eq(&self, other: &Self) -> bool {
self.spelled == other.spelled
}
}
impl<T> Eq for Id<T> {}
impl<T> Hash for Id<T> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.spelled.hash(state);
}
}
impl<T> Display for Id<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.spelled)
}
}
impl<T: Debug> Debug for Id<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("Id").field(&self.value).finish()
}
}
impl<T> Serialize for Id<T> {
fn serialize<S: serde::Serializer>(&self, out: S) -> Result<S::Ok, S::Error> {
out.serialize_str(&self.spelled)
}
}
impl<'de, T: Display + FromStr> Deserialize<'de> for Id<T> {
fn deserialize<D: serde::Deserializer<'de>>(from: D) -> Result<Self, D::Error> {
let spelled = <std::borrow::Cow<'de, str>>::deserialize(from)?;
T::from_str(&spelled).map(Id::new).map_err(|_| {
serde::de::Error::custom(format!("{spelled:?} is not a {}", any::type_name::<T>()))
})
}
}
pub trait ReactiveMapValue:
Serialize + DeserializeOwned + Clone + Send + Sync + 'static + Default
{
}
impl<T: Serialize + DeserializeOwned + Clone + Send + Sync + 'static + Default> ReactiveMapValue
for T
{
}
pub struct MapCache<K, V> {
entries: ArcSwap<Snapshot<K, V>>,
settled: Mutex<Settled>,
}
impl<K: AsRef<str>, V> Default for MapCache<K, V> {
fn default() -> Self {
Self {
entries: ArcSwap::from_pointee(RedBlackTreeMapSync::new_sync()),
settled: Mutex::new(Settled::default()),
}
}
}
#[derive(Default)]
struct Settled {
keys: std::collections::HashMap<String, u64>,
cleared: u64,
forgotten: u64,
latest: u64,
}
const FORGET_AFTER: u64 = 4096;
impl Settled {
fn admits(&self, name: &str, at: u64) -> bool {
at >= self.cleared
&& match self.keys.get(name) {
Some(held) => at >= *held,
None => at > self.forgotten || self.forgotten == 0,
}
}
fn took(&mut self, name: &str, at: u64) {
self.keys.insert(name.to_string(), at);
self.latest = self.latest.max(at);
}
fn let_go_of_the_long_removed(&mut self, live: impl Fn(&str) -> bool) {
if self.keys.len() as u64 <= 2 * FORGET_AFTER {
return;
}
let horizon = self.latest.saturating_sub(FORGET_AFTER);
let mut forgotten = self.forgotten;
self.keys.retain(|name, held| {
let kept = *held > horizon || live(name);
if !kept {
forgotten = forgotten.max(*held);
}
kept
});
self.forgotten = forgotten;
}
}
#[derive(Clone)]
struct ByName<K>(K);
impl<K: AsRef<str>> Borrow<str> for ByName<K> {
fn borrow(&self) -> &str {
self.0.as_ref()
}
}
impl<K: AsRef<str>> PartialEq for ByName<K> {
fn eq(&self, other: &Self) -> bool {
self.0.as_ref() == other.0.as_ref()
}
}
impl<K: AsRef<str>> Eq for ByName<K> {}
impl<K: AsRef<str>> Ord for ByName<K> {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.0.as_ref().cmp(other.0.as_ref())
}
}
impl<K: AsRef<str>> PartialOrd for ByName<K> {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl<K: AsRef<str> + Clone, V: Clone> MapCache<K, V> {
pub fn new() -> Self {
Self::default()
}
pub fn get(&self, name: &str) -> Option<V> {
self.entries.load().get(name).cloned()
}
pub fn contains_key(&self, name: &str) -> bool {
self.entries.load().contains_key(name)
}
pub fn owned_key(&self, name: &str) -> Option<K> {
self.entries
.load()
.get_key_value(name)
.map(|(key, _)| key.0.clone())
}
pub fn len(&self) -> usize {
self.entries.load().size()
}
pub fn is_empty(&self) -> bool {
self.entries.load().is_empty()
}
pub fn clear(&self) {
self.entries
.store(Arc::new(RedBlackTreeMapSync::new_sync()));
}
pub fn view(&self) -> Entries<K, V> {
Entries {
held: self.entries.load_full(),
}
}
pub fn keys(&self) -> Walk<K, V, K> {
Walk::new(self.entries.load_full(), |key, _| key.0.clone())
}
pub fn entries(&self) -> Walk<K, V, (K, V)> {
Walk::new(self.entries.load_full(), |key, value| {
(key.0.clone(), value.clone())
})
}
}
type Snapshot<K, V> = RedBlackTreeMapSync<ByName<K>, V>;
type Held<K, V> = Arc<Snapshot<K, V>>;
type Pairs<'a, K, V> = <&'a Snapshot<K, V> as IntoIterator>::IntoIter;
pub struct Values<'a, K: AsRef<str>, V> {
pairs: Pairs<'a, K, V>,
}
impl<'a, K: AsRef<str>, V> Iterator for Values<'a, K, V> {
type Item = (&'a K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
self.pairs.next().map(|(key, value)| (&key.0, value))
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.pairs.size_hint()
}
}
impl<K: AsRef<str>, V> DoubleEndedIterator for Values<'_, K, V> {
fn next_back(&mut self) -> Option<Self::Item> {
self.pairs.next_back().map(|(key, value)| (&key.0, value))
}
}
pub struct Walk<K: AsRef<str> + 'static, V: 'static, T> {
pairs: Pairs<'static, K, V>,
take: fn(&ByName<K>, &V) -> T,
_held: Held<K, V>,
}
impl<K: AsRef<str>, V, T> Walk<K, V, T> {
fn new(held: Held<K, V>, take: fn(&ByName<K>, &V) -> T) -> Self {
let pairs =
unsafe { std::mem::transmute::<Pairs<'_, K, V>, Pairs<'static, K, V>>(held.iter()) };
Self {
pairs,
take,
_held: held,
}
}
}
impl<K: AsRef<str>, V, T> Iterator for Walk<K, V, T> {
type Item = T;
fn next(&mut self) -> Option<T> {
self.pairs
.next()
.map(|(key, value)| (self.take)(key, value))
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.pairs.size_hint()
}
}
impl<K: AsRef<str>, V, T> DoubleEndedIterator for Walk<K, V, T> {
fn next_back(&mut self) -> Option<T> {
self.pairs
.next_back()
.map(|(key, value)| (self.take)(key, value))
}
}
impl<K: AsRef<str>, V, T> ExactSizeIterator for Walk<K, V, T> {}
pub struct Entries<K, V> {
held: Held<K, V>,
}
impl<'e, K: AsRef<str>, V> IntoIterator for &'e Entries<K, V> {
type Item = (&'e K, &'e V);
type IntoIter = Values<'e, K, V>;
fn into_iter(self) -> Self::IntoIter {
Values {
pairs: self.held.iter(),
}
}
}
impl<K: AsRef<str>, V> Entries<K, V> {
pub fn iter(&self) -> impl DoubleEndedIterator<Item = (&K, &V)> {
self.held.iter().map(|(key, value)| (&key.0, value))
}
pub fn len(&self) -> usize {
self.held.size()
}
pub fn is_empty(&self) -> bool {
self.held.is_empty()
}
}
impl<K: AsRef<str> + Clone, V: Clone> MapCache<K, V> {
pub fn insert(&self, key: K, value: V) -> Option<V> {
let replaced = self
.entries
.rcu(|current| current.insert(ByName(key.clone()), value.clone()));
replaced.get(key.as_ref()).cloned()
}
pub fn remove(&self, name: &str) -> Option<(K, V)> {
let previous = self.entries.rcu(|current| current.remove(name));
previous
.get_key_value(name)
.map(|(key, value)| (key.0.clone(), value.clone()))
}
pub fn insert_settled(&self, key: K, value: V, at: u64) -> bool {
let mut settled = held(&self.settled);
if !settled.admits(key.as_ref(), at) {
return false;
}
settled.took(key.as_ref(), at);
self.insert(key, value);
true
}
pub fn remove_settled(&self, name: &str, at: u64) -> bool {
let mut settled = held(&self.settled);
if !settled.admits(name, at) {
return false;
}
settled.took(name, at);
self.remove(name);
let held = self.entries.load();
settled.let_go_of_the_long_removed(|name| held.contains_key(name));
true
}
pub fn clear_settled(&self, at: u64) -> bool {
let mut settled = held(&self.settled);
if at < settled.cleared {
return false;
}
settled.cleared = at;
settled.latest = settled.latest.max(at);
settled.keys.retain(|_, held| *held > at);
let later = &settled.keys;
self.entries.rcu(|current| {
current
.iter()
.filter(|(key, _)| later.contains_key(key.0.as_ref()))
.fold(RedBlackTreeMapSync::new_sync(), |kept, (key, value)| {
kept.insert(key.clone(), value.clone())
})
});
true
}
}
impl<K: AsRef<str> + Debug, V: Debug> Debug for MapCache<K, V> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let held = self.entries.load();
f.debug_map()
.entries(held.iter().map(|(key, value)| (&key.0, value)))
.finish()
}
}
pub struct ReactiveMapCore<K, V> {
pub subscribers_any: Arc<Mutex<Vec<(u64, SubscriberAny<K, V>, SubscriptionMeta)>>>,
pub subscribers_key: Arc<DashMap<K, Vec<(u64, SubscriberKey<K, V>, SubscriptionMeta)>>>,
pub next_id: Arc<AtomicU64>,
pub cache: Arc<MapCache<K, V>>,
}
impl<K, V> Clone for ReactiveMapCore<K, V> {
fn clone(&self) -> Self {
Self {
subscribers_any: self.subscribers_any.clone(),
subscribers_key: self.subscribers_key.clone(),
next_id: self.next_id.clone(),
cache: self.cache.clone(),
}
}
}
struct Counted<'a, T>(&'a Mutex<Vec<T>>);
impl<T> Debug for Counted<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.0.try_lock() {
Ok(list) => write!(f, "{}", list.len()),
Err(_) => f.write_str("<locked>"),
}
}
}
impl<K: AsRef<str> + Debug + Hash + Eq, V: Debug> Debug for ReactiveMapCore<K, V> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ReactiveMapCore")
.field("cache", &self.cache)
.field("subscribers_any", &Counted(&self.subscribers_any))
.finish()
}
}
impl<K: ReactiveMapKey, V: ReactiveMapValue> Default for ReactiveMapCore<K, V> {
fn default() -> Self {
Self::new()
}
}
impl<K: ReactiveMapKey, V: ReactiveMapValue> ReactiveMapCore<K, V> {
pub fn new() -> Self {
Self::with_capacity(0)
}
pub fn with_capacity(_entries: usize) -> Self {
Self {
subscribers_any: Arc::new(Mutex::new(Vec::new())),
subscribers_key: Arc::new(DashMap::new()),
next_id: Arc::new(AtomicU64::new(0)),
cache: Arc::new(MapCache::new()),
}
}
#[track_caller]
pub fn subscribe_any<F>(&self, callback: F) -> SignalSubscription
where
F: Fn(&MapChange<K, V>) + Send + Sync + 'static,
{
let location = Location::caller();
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
let meta = SubscriptionMeta {
id,
location,
name: None,
};
held(&self.subscribers_any).push((id, Arc::new(callback), meta));
let subs_for_name = self.subscribers_any.clone();
let set_name = Arc::new(move |name: &'static str| {
label(&mut held(&subs_for_name), id, name);
});
let subs_for_cleanup = self.subscribers_any.clone();
SignalSubscription::new(
id,
location,
set_name,
Arc::new(move |id| {
forget(&mut held(&subs_for_cleanup), id);
}),
)
}
#[track_caller]
pub fn subscribe_key<F>(&self, key: K, callback: F) -> SignalSubscription
where
F: Fn(&MapChange<K, V>) + Send + Sync + 'static,
{
let location = Location::caller();
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
let meta = SubscriptionMeta {
id,
location,
name: None,
};
self.subscribers_key
.entry(key.clone())
.or_default()
.push((id, Arc::new(callback), meta));
let subs_for_name = self.subscribers_key.clone();
let key_for_name = key.clone();
let set_name = Arc::new(move |name: &'static str| {
if let Some(mut list) = subs_for_name.get_mut(&key_for_name) {
label(&mut list, id, name);
}
});
let subs_for_cleanup = self.subscribers_key.clone();
SignalSubscription::new(
id,
location,
set_name,
Arc::new(move |id| {
if let Some(mut list) = subs_for_cleanup.get_mut(&key) {
forget(&mut list, id);
}
subs_for_cleanup.remove_if(&key, |_, list| list.is_empty());
}),
)
}
pub fn notify(&self, change: &MapChange<K, V>) {
let keyed: Vec<_> = match change.key() {
Some(k) => self
.subscribers_key
.get(k)
.map(|entries| {
entries
.iter()
.map(|(_, cb, meta)| (cb.clone(), *meta))
.collect()
})
.unwrap_or_default(),
None => self
.subscribers_key
.iter()
.flat_map(|entries| {
entries
.value()
.iter()
.map(|(_, cb, meta)| (cb.clone(), *meta))
.collect::<Vec<_>>()
})
.collect(),
};
let any: Vec<_> = held(&self.subscribers_any)
.iter()
.map(|(_, cb, meta)| (cb.clone(), *meta))
.collect();
for (cb, meta) in keyed {
tracing::trace!(
target: "amethystate",
subscription_id = meta.id,
name = meta.name,
location = format!("{}:{}", meta.location.file(), meta.location.line()),
"map signal emit → key subscription fire",
);
cb(change);
}
for (cb, meta) in any {
tracing::trace!(
target: "amethystate",
subscription_id = meta.id,
name = meta.name,
location = format!("{}:{}", meta.location.file(), meta.location.line()),
"map signal emit → any subscription fire",
);
cb(change);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_cache_lists_entries_in_the_order_the_store_keys_them() {
let cache = MapCache::<String, u8>::new();
let names = ["a1b", "a.b", "10", "9", "100", "a\\b", "cpu", "Cpu"];
for (n, name) in names.iter().enumerate() {
cache.insert((*name).to_string(), n as u8);
}
let under = StorePath::segment("m");
let mut by_store = names.to_vec();
by_store.sort_by_key(|name| under.push(name).key().as_bytes().to_vec());
assert_eq!(cache.keys().collect::<Vec<_>>(), by_store);
}
#[test]
fn keys_that_come_and_go_leave_a_bounded_record_behind() {
let cache = MapCache::<String, u8>::new();
for n in 0..10_000u64 {
cache.insert_settled(format!("pid{n}"), 1, 2 * n);
cache.remove_settled(&format!("pid{n}"), 2 * n + 1);
}
assert!(held(&cache.settled).keys.len() <= 2 * FORGET_AFTER as usize);
}
#[test]
fn a_removal_long_settled_still_keeps_a_late_insert_out() {
let cache = MapCache::<String, u8>::new();
cache.insert_settled("gone".to_string(), 1, 1);
cache.remove_settled("gone", 2);
for n in 0..3 * FORGET_AFTER {
cache.insert_settled(format!("other{n}"), 1, 3 + 2 * n);
cache.remove_settled(&format!("other{n}"), 4 + 2 * n);
}
assert_eq!(
(
cache.insert_settled("gone".to_string(), 7, 1),
cache.get("gone")
),
(false, None)
);
}
#[test]
fn a_name_holding_the_separator_is_one_entry() {
let cache = MapCache::<String, u8>::new();
cache.insert("dark.mode".to_string(), 1);
assert_eq!(cache.len(), 1);
assert_eq!(cache.get("dark.mode"), Some(1));
assert_eq!(cache.get("dark"), None);
assert_eq!(cache.get("dark\\.mode"), None);
assert_eq!(cache.keys().collect::<Vec<_>>(), ["dark.mode"]);
}
#[test]
fn a_dropped_key_subscription_leaves_no_entry_behind() {
let core = ReactiveMapCore::<String, u64>::new();
for i in 0..64u64 {
let key = format!("col{i}");
let sub = core.subscribe_key(key, |_| {});
drop(sub);
}
assert_eq!(core.subscribers_key.len(), 0);
}
#[test]
fn a_key_still_watched_by_someone_keeps_its_entry() {
let core = ReactiveMapCore::<String, u64>::new();
let first = core.subscribe_key("cpu".to_string(), |_| {});
let second = core.subscribe_key("cpu".to_string(), |_| {});
drop(first);
assert_eq!(core.subscribers_key.len(), 1);
drop(second);
assert_eq!(core.subscribers_key.len(), 0);
}
}