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;
#[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<E> {
let endpoint = {
let guard = self.inner.read_entries();
guard.get(address).map(|entry| Arc::clone(&entry.endpoint))
};
endpoint.as_deref().cloned()
}
}
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};
use std::collections::HashSet;
#[test]
fn lease_owns_and_releases_one_registration() {
let space = AddressSpace::new();
let lease = space.claim(7, "first").unwrap();
assert_eq!(space.resolve(&7), 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), 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 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());
}
}