1use crate::SignalSubscription;
2use crate::change::MapChange;
3use crate::path::{PathRef, StorePath, Under};
4use crate::primitives::signal::{SubscriptionMeta, forget, held, label};
5use arc_swap::ArcSwap;
6use dashmap::DashMap;
7use rpds::RedBlackTreeMapSync;
8use serde::de::DeserializeOwned;
9use serde::{Deserialize, Serialize};
10use smol_str::{SmolStr, SmolStrBuilder};
11use std::any;
12use std::borrow::Borrow;
13use std::fmt::{self, Debug, Display, Write as _};
14use std::hash::Hash;
15use std::panic::Location;
16use std::str::FromStr;
17use std::sync::atomic::{AtomicU64, Ordering};
18use std::sync::{Arc, Mutex};
19
20pub trait MapEntryPath {
22 fn entry(&self, name: &str) -> StorePath;
24}
25
26impl MapEntryPath for StorePath {
27 fn entry(&self, name: &str) -> StorePath {
28 self.push(name)
29 }
30}
31
32pub type SubscriberAny<K, V> = Arc<dyn Fn(&MapChange<K, V>) + Send + Sync + 'static>;
33pub type SubscriberKey<K, V> = Arc<dyn Fn(&MapChange<K, V>) + Send + Sync + 'static>;
34
35pub enum EntryOf<K> {
53 Took(K),
55
56 ThePathItself,
59
60 NameWillNotRead,
62
63 NotThisMaps,
67}
68
69pub fn entry_of<K>(prefix: &StorePath, stored: PathRef<'_>) -> EntryOf<K>
71where
72 K: ReactiveMapKey,
73{
74 match stored.level_under(prefix) {
75 Under::Prefix => EntryOf::ThePathItself,
76 Under::Deeper(_) | Under::Outside => EntryOf::NotThisMaps,
77 Under::Entry(name) => match K::read(name.as_str()) {
78 Some(key) => EntryOf::Took(key),
79 None => EntryOf::NameWillNotRead,
80 },
81 }
82}
83
84pub trait ReactiveMapKey: AsRef<str> + Clone + Hash + Eq + Send + Sync + 'static {
85 fn read(name: &str) -> Option<Self>;
91}
92
93impl ReactiveMapKey for String {
94 fn read(name: &str) -> Option<Self> {
95 Some(name.to_string())
96 }
97}
98
99impl ReactiveMapKey for SmolStr {
100 fn read(name: &str) -> Option<Self> {
101 Some(SmolStr::new(name))
102 }
103}
104
105impl<T> ReactiveMapKey for Id<T>
106where
107 T: Display + FromStr + Clone + Send + Sync + 'static,
108{
109 fn read(name: &str) -> Option<Self> {
110 T::from_str(name).ok().map(Id::new)
111 }
112}
113
114#[derive(Clone)]
127pub struct Id<T> {
128 spelled: SmolStr,
129 value: T,
130}
131
132impl<T: Display> Id<T> {
133 pub fn new(value: T) -> Self {
134 let mut out = SmolStrBuilder::new();
135 let _ = write!(out, "{value}");
136
137 Self {
138 spelled: out.finish(),
139 value,
140 }
141 }
142}
143
144impl<T> Id<T> {
145 pub fn get(&self) -> &T {
147 &self.value
148 }
149
150 pub fn into_inner(self) -> T {
151 self.value
152 }
153}
154
155impl<T> AsRef<str> for Id<T> {
156 fn as_ref(&self) -> &str {
157 &self.spelled
158 }
159}
160
161impl<T> Borrow<str> for Id<T> {
165 fn borrow(&self) -> &str {
166 &self.spelled
167 }
168}
169
170impl<T> PartialEq for Id<T> {
173 fn eq(&self, other: &Self) -> bool {
174 self.spelled == other.spelled
175 }
176}
177
178impl<T> Eq for Id<T> {}
179
180impl<T> Hash for Id<T> {
181 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
182 self.spelled.hash(state);
183 }
184}
185
186impl<T> Display for Id<T> {
187 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188 f.write_str(&self.spelled)
189 }
190}
191
192impl<T: Debug> Debug for Id<T> {
193 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
194 f.debug_tuple("Id").field(&self.value).finish()
195 }
196}
197
198impl<T> Serialize for Id<T> {
202 fn serialize<S: serde::Serializer>(&self, out: S) -> Result<S::Ok, S::Error> {
203 out.serialize_str(&self.spelled)
204 }
205}
206
207impl<'de, T: Display + FromStr> Deserialize<'de> for Id<T> {
208 fn deserialize<D: serde::Deserializer<'de>>(from: D) -> Result<Self, D::Error> {
209 let spelled = <std::borrow::Cow<'de, str>>::deserialize(from)?;
210
211 T::from_str(&spelled).map(Id::new).map_err(|_| {
212 serde::de::Error::custom(format!("{spelled:?} is not a {}", any::type_name::<T>()))
213 })
214 }
215}
216
217pub trait ReactiveMapValue:
218 Serialize + DeserializeOwned + Clone + Send + Sync + 'static + Default
219{
220}
221impl<T: Serialize + DeserializeOwned + Clone + Send + Sync + 'static + Default> ReactiveMapValue
222 for T
223{
224}
225
226pub struct MapCache<K, V> {
255 entries: ArcSwap<Snapshot<K, V>>,
256 settled: Mutex<Settled>,
257}
258
259impl<K: AsRef<str>, V> Default for MapCache<K, V> {
260 fn default() -> Self {
261 Self {
262 entries: ArcSwap::from_pointee(RedBlackTreeMapSync::new_sync()),
263 settled: Mutex::new(Settled::default()),
264 }
265 }
266}
267
268#[derive(Default)]
282struct Settled {
283 keys: std::collections::HashMap<String, u64>,
284 cleared: u64,
285 forgotten: u64,
286 latest: u64,
287}
288
289const FORGET_AFTER: u64 = 4096;
290
291impl Settled {
292 fn admits(&self, name: &str, at: u64) -> bool {
293 at >= self.cleared
294 && match self.keys.get(name) {
295 Some(held) => at >= *held,
296 None => at > self.forgotten || self.forgotten == 0,
297 }
298 }
299
300 fn took(&mut self, name: &str, at: u64) {
301 self.keys.insert(name.to_string(), at);
302 self.latest = self.latest.max(at);
303 }
304
305 fn let_go_of_the_long_removed(&mut self, live: impl Fn(&str) -> bool) {
306 if self.keys.len() as u64 <= 2 * FORGET_AFTER {
307 return;
308 }
309
310 let horizon = self.latest.saturating_sub(FORGET_AFTER);
311 let mut forgotten = self.forgotten;
312 self.keys.retain(|name, held| {
313 let kept = *held > horizon || live(name);
314 if !kept {
315 forgotten = forgotten.max(*held);
316 }
317 kept
318 });
319 self.forgotten = forgotten;
320 }
321}
322
323#[derive(Clone)]
331struct ByName<K>(K);
332
333impl<K: AsRef<str>> Borrow<str> for ByName<K> {
334 fn borrow(&self) -> &str {
335 self.0.as_ref()
336 }
337}
338
339impl<K: AsRef<str>> PartialEq for ByName<K> {
340 fn eq(&self, other: &Self) -> bool {
341 self.0.as_ref() == other.0.as_ref()
342 }
343}
344
345impl<K: AsRef<str>> Eq for ByName<K> {}
346
347impl<K: AsRef<str>> Ord for ByName<K> {
348 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
349 self.0.as_ref().cmp(other.0.as_ref())
350 }
351}
352
353impl<K: AsRef<str>> PartialOrd for ByName<K> {
354 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
355 Some(self.cmp(other))
356 }
357}
358
359impl<K: AsRef<str> + Clone, V: Clone> MapCache<K, V> {
360 pub fn new() -> Self {
361 Self::default()
362 }
363
364 pub fn get(&self, name: &str) -> Option<V> {
365 self.entries.load().get(name).cloned()
366 }
367
368 pub fn contains_key(&self, name: &str) -> bool {
369 self.entries.load().contains_key(name)
370 }
371
372 pub fn owned_key(&self, name: &str) -> Option<K> {
375 self.entries
376 .load()
377 .get_key_value(name)
378 .map(|(key, _)| key.0.clone())
379 }
380
381 pub fn len(&self) -> usize {
382 self.entries.load().size()
383 }
384
385 pub fn is_empty(&self) -> bool {
386 self.entries.load().is_empty()
387 }
388
389 pub fn clear(&self) {
390 self.entries
391 .store(Arc::new(RedBlackTreeMapSync::new_sync()));
392 }
393
394 pub fn view(&self) -> Entries<K, V> {
399 Entries {
400 held: self.entries.load_full(),
401 }
402 }
403
404 pub fn keys(&self) -> Walk<K, V, K> {
406 Walk::new(self.entries.load_full(), |key, _| key.0.clone())
407 }
408
409 pub fn entries(&self) -> Walk<K, V, (K, V)> {
411 Walk::new(self.entries.load_full(), |key, value| {
412 (key.0.clone(), value.clone())
413 })
414 }
415}
416
417type Snapshot<K, V> = RedBlackTreeMapSync<ByName<K>, V>;
418type Held<K, V> = Arc<Snapshot<K, V>>;
419type Pairs<'a, K, V> = <&'a Snapshot<K, V> as IntoIterator>::IntoIter;
420pub struct Values<'a, K: AsRef<str>, V> {
425 pairs: Pairs<'a, K, V>,
426}
427
428impl<'a, K: AsRef<str>, V> Iterator for Values<'a, K, V> {
429 type Item = (&'a K, &'a V);
430
431 fn next(&mut self) -> Option<Self::Item> {
432 self.pairs.next().map(|(key, value)| (&key.0, value))
433 }
434
435 fn size_hint(&self) -> (usize, Option<usize>) {
436 self.pairs.size_hint()
437 }
438}
439
440impl<K: AsRef<str>, V> DoubleEndedIterator for Values<'_, K, V> {
441 fn next_back(&mut self) -> Option<Self::Item> {
442 self.pairs.next_back().map(|(key, value)| (&key.0, value))
443 }
444}
445
446pub struct Walk<K: AsRef<str> + 'static, V: 'static, T> {
455 pairs: Pairs<'static, K, V>,
456 take: fn(&ByName<K>, &V) -> T,
457 _held: Held<K, V>,
458}
459
460impl<K: AsRef<str>, V, T> Walk<K, V, T> {
461 fn new(held: Held<K, V>, take: fn(&ByName<K>, &V) -> T) -> Self {
462 let pairs =
473 unsafe { std::mem::transmute::<Pairs<'_, K, V>, Pairs<'static, K, V>>(held.iter()) };
474
475 Self {
476 pairs,
477 take,
478 _held: held,
479 }
480 }
481}
482
483impl<K: AsRef<str>, V, T> Iterator for Walk<K, V, T> {
484 type Item = T;
485
486 fn next(&mut self) -> Option<T> {
487 self.pairs
488 .next()
489 .map(|(key, value)| (self.take)(key, value))
490 }
491
492 fn size_hint(&self) -> (usize, Option<usize>) {
493 self.pairs.size_hint()
494 }
495}
496
497impl<K: AsRef<str>, V, T> DoubleEndedIterator for Walk<K, V, T> {
498 fn next_back(&mut self) -> Option<T> {
499 self.pairs
500 .next_back()
501 .map(|(key, value)| (self.take)(key, value))
502 }
503}
504
505impl<K: AsRef<str>, V, T> ExactSizeIterator for Walk<K, V, T> {}
506
507pub struct Entries<K, V> {
509 held: Held<K, V>,
510}
511
512impl<'e, K: AsRef<str>, V> IntoIterator for &'e Entries<K, V> {
513 type Item = (&'e K, &'e V);
514 type IntoIter = Values<'e, K, V>;
515
516 fn into_iter(self) -> Self::IntoIter {
517 Values {
518 pairs: self.held.iter(),
519 }
520 }
521}
522
523impl<K: AsRef<str>, V> Entries<K, V> {
524 pub fn iter(&self) -> impl DoubleEndedIterator<Item = (&K, &V)> {
525 self.held.iter().map(|(key, value)| (&key.0, value))
526 }
527
528 pub fn len(&self) -> usize {
529 self.held.size()
530 }
531
532 pub fn is_empty(&self) -> bool {
533 self.held.is_empty()
534 }
535}
536
537impl<K: AsRef<str> + Clone, V: Clone> MapCache<K, V> {
538 pub fn insert(&self, key: K, value: V) -> Option<V> {
539 let replaced = self
540 .entries
541 .rcu(|current| current.insert(ByName(key.clone()), value.clone()));
542
543 replaced.get(key.as_ref()).cloned()
544 }
545
546 pub fn remove(&self, name: &str) -> Option<(K, V)> {
547 let previous = self.entries.rcu(|current| current.remove(name));
548
549 previous
550 .get_key_value(name)
551 .map(|(key, value)| (key.0.clone(), value.clone()))
552 }
553
554 pub fn insert_settled(&self, key: K, value: V, at: u64) -> bool {
557 let mut settled = held(&self.settled);
558 if !settled.admits(key.as_ref(), at) {
559 return false;
560 }
561
562 settled.took(key.as_ref(), at);
563 self.insert(key, value);
564 true
565 }
566
567 pub fn remove_settled(&self, name: &str, at: u64) -> bool {
570 let mut settled = held(&self.settled);
571 if !settled.admits(name, at) {
572 return false;
573 }
574
575 settled.took(name, at);
576 self.remove(name);
577 let held = self.entries.load();
578 settled.let_go_of_the_long_removed(|name| held.contains_key(name));
579 true
580 }
581
582 pub fn clear_settled(&self, at: u64) -> bool {
586 let mut settled = held(&self.settled);
587 if at < settled.cleared {
588 return false;
589 }
590
591 settled.cleared = at;
592 settled.latest = settled.latest.max(at);
593 settled.keys.retain(|_, held| *held > at);
594
595 let later = &settled.keys;
596 self.entries.rcu(|current| {
597 current
598 .iter()
599 .filter(|(key, _)| later.contains_key(key.0.as_ref()))
600 .fold(RedBlackTreeMapSync::new_sync(), |kept, (key, value)| {
601 kept.insert(key.clone(), value.clone())
602 })
603 });
604 true
605 }
606}
607
608impl<K: AsRef<str> + Debug, V: Debug> Debug for MapCache<K, V> {
609 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
610 let held = self.entries.load();
611
612 f.debug_map()
613 .entries(held.iter().map(|(key, value)| (&key.0, value)))
614 .finish()
615 }
616}
617
618pub struct ReactiveMapCore<K, V> {
619 pub subscribers_any: Arc<Mutex<Vec<(u64, SubscriberAny<K, V>, SubscriptionMeta)>>>,
620 pub subscribers_key: Arc<DashMap<K, Vec<(u64, SubscriberKey<K, V>, SubscriptionMeta)>>>,
621 pub next_id: Arc<AtomicU64>,
622 pub cache: Arc<MapCache<K, V>>,
623}
624
625impl<K, V> Clone for ReactiveMapCore<K, V> {
626 fn clone(&self) -> Self {
627 Self {
628 subscribers_any: self.subscribers_any.clone(),
629 subscribers_key: self.subscribers_key.clone(),
630 next_id: self.next_id.clone(),
631 cache: self.cache.clone(),
632 }
633 }
634}
635
636struct Counted<'a, T>(&'a Mutex<Vec<T>>);
637
638impl<T> Debug for Counted<'_, T> {
639 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
640 match self.0.try_lock() {
641 Ok(list) => write!(f, "{}", list.len()),
642 Err(_) => f.write_str("<locked>"),
643 }
644 }
645}
646
647impl<K: AsRef<str> + Debug + Hash + Eq, V: Debug> Debug for ReactiveMapCore<K, V> {
648 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
649 f.debug_struct("ReactiveMapCore")
650 .field("cache", &self.cache)
651 .field("subscribers_any", &Counted(&self.subscribers_any))
652 .finish()
653 }
654}
655
656impl<K: ReactiveMapKey, V: ReactiveMapValue> Default for ReactiveMapCore<K, V> {
657 fn default() -> Self {
658 Self::new()
659 }
660}
661
662impl<K: ReactiveMapKey, V: ReactiveMapValue> ReactiveMapCore<K, V> {
663 pub fn new() -> Self {
664 Self::with_capacity(0)
665 }
666
667 pub fn with_capacity(_entries: usize) -> Self {
671 Self {
672 subscribers_any: Arc::new(Mutex::new(Vec::new())),
673 subscribers_key: Arc::new(DashMap::new()),
674 next_id: Arc::new(AtomicU64::new(0)),
675 cache: Arc::new(MapCache::new()),
676 }
677 }
678
679 #[track_caller]
680 pub fn subscribe_any<F>(&self, callback: F) -> SignalSubscription
681 where
682 F: Fn(&MapChange<K, V>) + Send + Sync + 'static,
683 {
684 let location = Location::caller();
685 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
686 let meta = SubscriptionMeta {
687 id,
688 location,
689 name: None,
690 };
691 held(&self.subscribers_any).push((id, Arc::new(callback), meta));
692
693 let subs_for_name = self.subscribers_any.clone();
694 let set_name = Arc::new(move |name: &'static str| {
695 label(&mut held(&subs_for_name), id, name);
696 });
697 let subs_for_cleanup = self.subscribers_any.clone();
698 SignalSubscription::new(
699 id,
700 location,
701 set_name,
702 Arc::new(move |id| {
703 forget(&mut held(&subs_for_cleanup), id);
704 }),
705 )
706 }
707
708 #[track_caller]
709 pub fn subscribe_key<F>(&self, key: K, callback: F) -> SignalSubscription
710 where
711 F: Fn(&MapChange<K, V>) + Send + Sync + 'static,
712 {
713 let location = Location::caller();
714 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
715 let meta = SubscriptionMeta {
716 id,
717 location,
718 name: None,
719 };
720 self.subscribers_key
721 .entry(key.clone())
722 .or_default()
723 .push((id, Arc::new(callback), meta));
724
725 let subs_for_name = self.subscribers_key.clone();
726 let key_for_name = key.clone();
727 let set_name = Arc::new(move |name: &'static str| {
728 if let Some(mut list) = subs_for_name.get_mut(&key_for_name) {
729 label(&mut list, id, name);
730 }
731 });
732 let subs_for_cleanup = self.subscribers_key.clone();
733 SignalSubscription::new(
734 id,
735 location,
736 set_name,
737 Arc::new(move |id| {
738 if let Some(mut list) = subs_for_cleanup.get_mut(&key) {
739 forget(&mut list, id);
740 }
741 subs_for_cleanup.remove_if(&key, |_, list| list.is_empty());
742 }),
743 )
744 }
745
746 pub fn notify(&self, change: &MapChange<K, V>) {
753 let keyed: Vec<_> = match change.key() {
754 Some(k) => self
755 .subscribers_key
756 .get(k)
757 .map(|entries| {
758 entries
759 .iter()
760 .map(|(_, cb, meta)| (cb.clone(), *meta))
761 .collect()
762 })
763 .unwrap_or_default(),
764
765 None => self
770 .subscribers_key
771 .iter()
772 .flat_map(|entries| {
773 entries
774 .value()
775 .iter()
776 .map(|(_, cb, meta)| (cb.clone(), *meta))
777 .collect::<Vec<_>>()
778 })
779 .collect(),
780 };
781
782 let any: Vec<_> = held(&self.subscribers_any)
783 .iter()
784 .map(|(_, cb, meta)| (cb.clone(), *meta))
785 .collect();
786
787 for (cb, meta) in keyed {
788 tracing::trace!(
789 target: "amethystate",
790 subscription_id = meta.id,
791 name = meta.name,
792 location = format!("{}:{}", meta.location.file(), meta.location.line()),
793 "map signal emit → key subscription fire",
794 );
795 cb(change);
796 }
797
798 for (cb, meta) in any {
799 tracing::trace!(
800 target: "amethystate",
801 subscription_id = meta.id,
802 name = meta.name,
803 location = format!("{}:{}", meta.location.file(), meta.location.line()),
804 "map signal emit → any subscription fire",
805 );
806 cb(change);
807 }
808 }
809}
810
811#[cfg(test)]
812mod tests {
813 use super::*;
814
815 #[test]
816 fn the_cache_lists_entries_in_the_order_the_store_keys_them() {
817 let cache = MapCache::<String, u8>::new();
818 let names = ["a1b", "a.b", "10", "9", "100", "a\\b", "cpu", "Cpu"];
819
820 for (n, name) in names.iter().enumerate() {
821 cache.insert((*name).to_string(), n as u8);
822 }
823
824 let under = StorePath::segment("m");
825 let mut by_store = names.to_vec();
826 by_store.sort_by_key(|name| under.push(name).key().as_bytes().to_vec());
827
828 assert_eq!(cache.keys().collect::<Vec<_>>(), by_store);
829 }
830
831 #[test]
832 fn keys_that_come_and_go_leave_a_bounded_record_behind() {
833 let cache = MapCache::<String, u8>::new();
834
835 for n in 0..10_000u64 {
836 cache.insert_settled(format!("pid{n}"), 1, 2 * n);
837 cache.remove_settled(&format!("pid{n}"), 2 * n + 1);
838 }
839
840 assert!(held(&cache.settled).keys.len() <= 2 * FORGET_AFTER as usize);
841 }
842
843 #[test]
844 fn a_removal_long_settled_still_keeps_a_late_insert_out() {
845 let cache = MapCache::<String, u8>::new();
846
847 cache.insert_settled("gone".to_string(), 1, 1);
848 cache.remove_settled("gone", 2);
849 for n in 0..3 * FORGET_AFTER {
850 cache.insert_settled(format!("other{n}"), 1, 3 + 2 * n);
851 cache.remove_settled(&format!("other{n}"), 4 + 2 * n);
852 }
853
854 assert_eq!(
855 (
856 cache.insert_settled("gone".to_string(), 7, 1),
857 cache.get("gone")
858 ),
859 (false, None)
860 );
861 }
862
863 #[test]
864 fn a_name_holding_the_separator_is_one_entry() {
865 let cache = MapCache::<String, u8>::new();
866
867 cache.insert("dark.mode".to_string(), 1);
868
869 assert_eq!(cache.len(), 1);
870 assert_eq!(cache.get("dark.mode"), Some(1));
871 assert_eq!(cache.get("dark"), None);
872 assert_eq!(cache.get("dark\\.mode"), None);
873 assert_eq!(cache.keys().collect::<Vec<_>>(), ["dark.mode"]);
874 }
875
876 #[test]
877 fn a_dropped_key_subscription_leaves_no_entry_behind() {
878 let core = ReactiveMapCore::<String, u64>::new();
879
880 for i in 0..64u64 {
881 let key = format!("col{i}");
882 let sub = core.subscribe_key(key, |_| {});
883 drop(sub);
884 }
885
886 assert_eq!(core.subscribers_key.len(), 0);
887 }
888
889 #[test]
890 fn a_key_still_watched_by_someone_keeps_its_entry() {
891 let core = ReactiveMapCore::<String, u64>::new();
892
893 let first = core.subscribe_key("cpu".to_string(), |_| {});
894 let second = core.subscribe_key("cpu".to_string(), |_| {});
895
896 drop(first);
897 assert_eq!(core.subscribers_key.len(), 1);
898
899 drop(second);
900 assert_eq!(core.subscribers_key.len(), 0);
901 }
902}