123 lines
3.9 KiB
Rust
123 lines
3.9 KiB
Rust
//! Generic refcounted-binding helper used by the connection's
|
|
//! bound-session set.
|
|
//!
|
|
//! Multiple [`crate::ToolServer`] instances can share one
|
|
//! [`crate::HubConnection`] when they target the same `(url, principal)`.
|
|
//! Each instance independently asks for a session binding; the substrate
|
|
//! must `register_session` once per session (not once per consumer) and
|
|
//! `unregister_session` only when the LAST consumer drops its borrow.
|
|
//! [`RefCountedSet`] tracks the per-key borrow count behind a
|
|
//! [`dashmap::DashMap`] so increments and decrements never serialise on
|
|
//! a single mutex.
|
|
|
|
use std::hash::Hash;
|
|
|
|
use dashmap::DashMap;
|
|
|
|
/// Refcounted set keyed by `K`. Each [`Self::increment`] returns the
|
|
/// new count; the corresponding [`Self::decrement`] returns the count
|
|
/// AFTER the decrement (so callers fire teardown when the result is
|
|
/// `Some(0)`).
|
|
#[derive(Debug, Default)]
|
|
pub struct RefCountedSet<K: Eq + Hash> {
|
|
counts: DashMap<K, u64>,
|
|
}
|
|
|
|
impl<K: Eq + Hash> RefCountedSet<K> {
|
|
/// Empty set.
|
|
pub fn new() -> Self {
|
|
Self {
|
|
counts: DashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Increment `key`'s refcount. Returns `(prev_count, new_count)`
|
|
/// so callers can detect the 0→1 edge (when the protocol-level
|
|
/// register call must fire).
|
|
pub fn increment(&self, key: K) -> (u64, u64)
|
|
where
|
|
K: Clone,
|
|
{
|
|
let mut entry = self.counts.entry(key).or_insert(0);
|
|
let prev = *entry;
|
|
*entry = prev.saturating_add(1);
|
|
(prev, *entry)
|
|
}
|
|
|
|
/// Decrement `key`'s refcount. Returns the post-decrement count;
|
|
/// `Some(0)` means the entry was removed and callers should fire
|
|
/// the protocol-level unregister. `None` means the key was not
|
|
/// present (idempotent drop).
|
|
pub fn decrement(&self, key: &K) -> Option<u64> {
|
|
let mut current = None;
|
|
self.counts.remove_if_mut(key, |_, value| {
|
|
*value = value.saturating_sub(1);
|
|
current = Some(*value);
|
|
*value == 0
|
|
});
|
|
current
|
|
}
|
|
|
|
/// Snapshot the live keys. Allocates a fresh `Vec` — only used by
|
|
/// the reconnect-replay path which fires once per disconnect.
|
|
pub fn snapshot_keys(&self) -> Vec<K>
|
|
where
|
|
K: Clone,
|
|
{
|
|
self.counts.iter().map(|kv| kv.key().clone()).collect()
|
|
}
|
|
|
|
/// `true` when no key has a non-zero refcount.
|
|
pub fn is_empty(&self) -> bool {
|
|
self.counts.is_empty()
|
|
}
|
|
|
|
/// Number of distinct live keys.
|
|
pub fn len(&self) -> usize {
|
|
self.counts.len()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn increment_returns_new_count() {
|
|
let set = RefCountedSet::<&'static str>::new();
|
|
assert_eq!(set.increment("a"), (0, 1));
|
|
assert_eq!(set.increment("a"), (1, 2));
|
|
assert_eq!(set.increment("b"), (0, 1));
|
|
assert_eq!(set.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn decrement_removes_at_zero() {
|
|
let set = RefCountedSet::<&'static str>::new();
|
|
set.increment("a");
|
|
set.increment("a");
|
|
assert_eq!(set.decrement(&"a"), Some(1));
|
|
assert!(!set.is_empty());
|
|
assert_eq!(set.decrement(&"a"), Some(0));
|
|
assert!(set.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn decrement_unknown_returns_none() {
|
|
let set = RefCountedSet::<&'static str>::new();
|
|
assert!(set.decrement(&"missing").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn increment_saturates_at_u64_max() {
|
|
let set = RefCountedSet::<&'static str>::new();
|
|
// Pre-load the entry to MAX-1 via direct DashMap access. The
|
|
// public API only ever reaches this region via overflow,
|
|
// which is impossible in practice; this test pins the
|
|
// saturating_add defensive line so it can't silently regress
|
|
// to wrapping_add.
|
|
set.counts.insert("max", u64::MAX - 1);
|
|
assert_eq!(set.increment("max"), (u64::MAX - 1, u64::MAX));
|
|
assert_eq!(set.increment("max"), (u64::MAX, u64::MAX));
|
|
}
|
|
}
|