mod table;
use core::hash::Hash;
#[cfg(loom)]
use loom::sync::{Arc, RwLock};
#[cfg(not(loom))]
use parking_lot::RwLock;
#[cfg(not(loom))]
use std::sync::Arc;
use std::{fmt, hash::Hasher, num::NonZeroU64};
use table::OpenTable;
pub struct Resolved<E>(E);
impl<E: Clone> Clone for Resolved<E> {
fn clone(&self) -> Self {
Self(self.0.clone())
}
}
impl<E> core::ops::Deref for Resolved<E> {
type Target = E;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl<E> AsRef<E> for Resolved<E> {
fn as_ref(&self) -> &E {
self
}
}
impl<E: fmt::Debug> fmt::Debug for Resolved<E> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.debug_tuple("Resolved").field(&self.0).finish()
}
}
impl<E: PartialEq> PartialEq for Resolved<E> {
fn eq(&self, other: &Self) -> bool {
self.as_ref() == other.as_ref()
}
}
impl<E: PartialEq> PartialEq<E> for Resolved<E> {
fn eq(&self, other: &E) -> bool {
self.as_ref() == other
}
}
impl<E: Eq> Eq for Resolved<E> {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AddressInUse<A>(pub A);
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClaimError<A> {
AddressInUse(A),
RegistrationIdsExhausted(A),
}
struct RegistrationScopeMarker {
_private: u8,
}
#[derive(Clone)]
pub struct RegistrationScopeId(Arc<RegistrationScopeMarker>);
impl PartialEq for RegistrationScopeId {
fn eq(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.0, &other.0)
}
}
impl Eq for RegistrationScopeId {}
impl core::hash::Hash for RegistrationScopeId {
fn hash<H: Hasher>(&self, state: &mut H) {
Arc::as_ptr(&self.0).hash(state);
}
}
impl fmt::Debug for RegistrationScopeId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_tuple("RegistrationScopeId")
.field(&Arc::as_ptr(&self.0))
.finish()
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct RegistrationId {
scope: RegistrationScopeId,
generation: NonZeroU64,
}
impl RegistrationId {
#[must_use]
pub fn scope_id(&self) -> &RegistrationScopeId {
&self.scope
}
}
impl fmt::Debug for RegistrationId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("RegistrationId")
.field("scope", &self.scope)
.field("generation", &self.generation)
.finish()
}
}
struct Inner<A, E> {
scope: RegistrationScopeId,
entries: RwLock<OpenTable<A, E>>,
}
impl<A, E> Inner<A, E> {
#[cfg(loom)]
fn read_entries(&self) -> loom::sync::RwLockReadGuard<'_, OpenTable<A, E>> {
self.entries.read().unwrap()
}
#[cfg(loom)]
fn write_entries(&self) -> loom::sync::RwLockWriteGuard<'_, OpenTable<A, E>> {
self.entries.write().unwrap()
}
#[cfg(not(loom))]
fn read_entries(&self) -> parking_lot::RwLockReadGuard<'_, OpenTable<A, E>> {
self.entries.read()
}
#[cfg(not(loom))]
fn write_entries(&self) -> parking_lot::RwLockWriteGuard<'_, OpenTable<A, E>> {
self.entries.write()
}
}
pub struct AddressSpace<A, E> {
inner: Arc<Inner<A, E>>,
}
impl<A, E> Clone for AddressSpace<A, E> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
}
}
}
impl<A, E> Default for AddressSpace<A, E> {
fn default() -> Self {
Self::new()
}
}
impl<A, E> AddressSpace<A, E> {
#[must_use]
pub fn new() -> Self {
Self {
inner: Arc::new(Inner {
scope: RegistrationScopeId(Arc::new(RegistrationScopeMarker { _private: 0 })),
entries: RwLock::new(OpenTable::new()),
}),
}
}
#[must_use]
pub fn registration_scope_id(&self) -> RegistrationScopeId {
self.inner.scope.clone()
}
#[must_use]
pub fn len(&self) -> usize {
self.inner.read_entries().len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl<A: Eq + Hash, E> AddressSpace<A, E> {
pub fn shrink_to_fit(&self) {
self.inner.write_entries().shrink_to_fit();
}
}
impl<A, E> AddressSpace<A, E>
where
A: Eq + Hash + Clone,
{
pub fn claim(&self, address: A, endpoint: E) -> Result<Lease<A, E>, AddressInUse<A>> {
match self.try_claim(address, endpoint) {
Ok(lease) => Ok(lease),
Err(ClaimError::AddressInUse(address)) => Err(AddressInUse(address)),
Err(ClaimError::RegistrationIdsExhausted(_)) => {
panic!("registration identities exhausted")
}
}
}
pub fn try_claim(&self, address: A, endpoint: E) -> Result<Lease<A, E>, ClaimError<A>> {
let key = address.clone();
let mut entries = self.inner.write_entries();
if entries.get(&address).is_some() {
return Err(ClaimError::AddressInUse(address));
}
let Some(generation) = entries.take_generation() else {
return Err(ClaimError::RegistrationIdsExhausted(address));
};
entries.insert(key, generation.get(), endpoint);
Ok(Lease {
inner: self.inner.clone(),
address,
generation,
released: false,
})
}
}
impl<A, E> AddressSpace<A, E>
where
A: Eq + Hash,
E: Clone,
{
#[must_use]
pub fn resolve(&self, address: &A) -> Option<Resolved<E>> {
let guard = self.inner.read_entries();
let endpoint = guard.get(address).map(|entry| Arc::clone(&entry.endpoint));
drop(guard);
endpoint.as_deref().cloned().map(Resolved)
}
}
pub struct Lease<A, E>
where
A: Eq + Hash,
{
inner: Arc<Inner<A, E>>,
address: A,
generation: NonZeroU64,
released: bool,
}
impl<A, E> Lease<A, E>
where
A: Eq + Hash,
{
#[must_use]
pub fn address(&self) -> &A {
&self.address
}
#[must_use]
pub fn registration_id(&self) -> RegistrationId {
RegistrationId {
scope: self.inner.scope.clone(),
generation: self.generation,
}
}
pub fn release(mut self) {
self.release_inner();
}
fn release_inner(&mut self) {
if self.released {
return;
}
self.released = true;
let removed = {
let mut entries = self.inner.write_entries();
entries.remove_if(&self.address, self.generation.get())
};
drop(removed);
}
}
impl<A, E> Drop for Lease<A, E>
where
A: Eq + Hash,
{
fn drop(&mut self) {
self.release_inner();
}
}
#[cfg(test)]
mod tests {
use super::{AddressSpace, ClaimError, Resolved};
use std::collections::HashSet;
fn assert_send_sync<T: Send + Sync>() {}
#[test]
fn resolved_inherits_positive_auto_traits_from_endpoint() {
assert_send_sync::<Resolved<String>>();
}
#[test]
fn resolved_preserves_explicit_endpoint_interior_mutability() {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
let state = Arc::new(AtomicUsize::new(1));
let space = AddressSpace::new();
let _lease = space.claim("worker", Arc::clone(&state)).unwrap();
let endpoint = space.resolve(&"worker").unwrap();
endpoint.store(2, Ordering::SeqCst);
assert_eq!(state.load(Ordering::SeqCst), 2);
}
#[test]
fn lease_owns_and_releases_one_registration() {
let space = AddressSpace::new();
let lease = space.claim(7, "first").unwrap();
assert_eq!(space.resolve(&7).as_deref().copied(), Some("first"));
assert!(matches!(
space.try_claim(7, "second"),
Err(ClaimError::AddressInUse(7))
));
drop(lease);
assert_eq!(space.resolve(&7), None);
let replacement = space.claim(7, "second").unwrap();
assert_eq!(space.resolve(&7).as_deref().copied(), Some("second"));
drop(replacement);
}
#[test]
fn registration_identity_distinguishes_replacements() {
let space = AddressSpace::new();
let first = space.claim("worker", 1).unwrap();
let first_id = first.registration_id();
first.release();
let second = space.claim("worker", 2).unwrap();
let second_id = second.registration_id();
assert_ne!(first_id, second_id);
assert_eq!(first_id.scope_id(), second_id.scope_id());
let ids = HashSet::from([first_id, second_id]);
assert_eq!(ids.len(), 2);
}
#[test]
fn registration_identity_is_scoped_to_the_address_space() {
let first = AddressSpace::new();
let first_peer = first.clone();
let second = AddressSpace::new();
assert_eq!(
first.registration_scope_id(),
first_peer.registration_scope_id()
);
assert_ne!(
first.registration_scope_id(),
second.registration_scope_id()
);
let first_id = first.claim((1_u8, 2_u8), ()).unwrap().registration_id();
let second_id = second.claim((1_u8, 2_u8), ()).unwrap().registration_id();
assert_ne!(first_id, second_id);
}
#[test]
fn retaining_an_identity_grants_no_registration_lifetime() {
let space = AddressSpace::new();
let lease = space
.claim(String::from("worker"), String::from("mailbox"))
.unwrap();
let identity = lease.registration_id();
drop(lease);
assert_eq!(space.resolve(&String::from("worker")), None);
let replacement = space
.claim(String::from("worker"), String::from("replacement"))
.unwrap();
assert_ne!(identity, replacement.registration_id());
}
#[test]
fn resolved_snapshot_survives_release_and_address_reuse() {
let space = AddressSpace::new();
let first = space.claim(7, String::from("first")).unwrap();
let snapshot = space.resolve(&7).unwrap();
first.release();
let second = space.claim(7, String::from("second")).unwrap();
assert_eq!(snapshot.as_str(), "first");
assert_eq!(space.resolve(&7).unwrap().as_str(), "second");
drop(second);
}
#[test]
fn exhausted_space_rejects_claims_without_reusing_an_identity() {
let space = AddressSpace::new();
space.inner.write_entries().set_next_generation(u64::MAX);
let last = space.claim(1_u64, "last").unwrap();
let last_id = last.registration_id();
last.release();
assert!(matches!(
space.try_claim(2, "never inserted"),
Err(ClaimError::RegistrationIdsExhausted(2))
));
assert!(space.is_empty());
assert_eq!(last_id.scope_id(), &space.registration_scope_id());
}
}