use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use arc_swap::ArcSwap;
use super::org::OrgId;
use super::org_grant::{CapabilityAuthorityId, OrgCapabilityGrant};
use super::org_grant_registry::ConsumerGrantSnapshot;
use super::org_routing_registry::{
DemandRefused, DemandSet, GrantMovementFence, PrivateAudienceScope, ReplaceRefused,
RoutingFamily, SlotKey, MAX_HANDLES_PER_FAMILY,
};
use super::org_scoped_ingest::CapabilityAudienceScope;
#[allow(dead_code)] pub(crate) const MAX_CAPABILITY_ENTRIES_PER_FAMILY: usize = MAX_HANDLES_PER_FAMILY;
#[derive(Clone, Debug)]
#[allow(dead_code)] pub(crate) struct FamilyDiscoveryCredentials {
pub acting_org: OrgId,
pub owner_audience_handle: [u8; 32],
pub grants: Vec<Arc<OrgCapabilityGrant>>,
}
impl FamilyDiscoveryCredentials {
fn owner_scope(&self) -> Option<PrivateAudienceScope> {
PrivateAudienceScope::new(CapabilityAudienceScope::Owner {
org_id: self.acting_org,
audience_handle: self.owner_audience_handle,
})
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum GrantDemand {
Leased([u8; 32]),
NotDiscovery,
NotLeased,
}
fn classify(
grant: &OrgCapabilityGrant,
capability: &CapabilityAuthorityId,
acting_org: &OrgId,
leased: &ConsumerGrantSnapshot,
) -> Option<GrantDemand> {
if &grant.capability != capability || &grant.grantee_org != acting_org {
return None;
}
if !grant.permits_discover() {
return Some(GrantDemand::NotDiscovery);
}
let Some(binding) = grant.discovery.as_ref() else {
return Some(GrantDemand::NotDiscovery);
};
let installed = leased
.get(&grant.grant_id)
.is_some_and(|record| record.audience_handle() == &binding.audience_handle);
if installed {
Some(GrantDemand::Leased(binding.audience_handle))
} else {
Some(GrantDemand::NotLeased)
}
}
#[allow(dead_code)] pub(crate) fn demand_set_for(
credentials: &FamilyDiscoveryCredentials,
capability: &CapabilityAuthorityId,
leased: &ConsumerGrantSnapshot,
) -> Option<Vec<SlotKey>> {
demand_set_from(credentials, credentials.grants.iter(), capability, leased)
}
fn demand_set_from<'a>(
credentials: &FamilyDiscoveryCredentials,
grants: impl Iterator<Item = &'a Arc<OrgCapabilityGrant>>,
capability: &CapabilityAuthorityId,
leased: &ConsumerGrantSnapshot,
) -> Option<Vec<SlotKey>> {
let mut keys = vec![SlotKey {
scope: credentials.owner_scope()?,
capability: *capability,
}];
for grant in grants {
let Some(GrantDemand::Leased(audience_handle)) =
classify(grant, capability, &credentials.acting_org, leased)
else {
continue;
};
let scope = CapabilityAudienceScope::Grant {
grant_id: grant.grant_id,
audience_handle,
};
if let Some(scope) = PrivateAudienceScope::new(scope) {
keys.push(SlotKey {
scope,
capability: *capability,
});
}
}
keys.sort();
keys.dedup();
Some(keys)
}
#[allow(dead_code)] pub(crate) struct CapabilityRouteHandle {
capability: CapabilityAuthorityId,
derived_at: u64,
demanded: Vec<SlotKey>,
demands: DemandSet,
}
#[allow(dead_code)] impl CapabilityRouteHandle {
pub(crate) fn capability(&self) -> &CapabilityAuthorityId {
&self.capability
}
pub(crate) fn demands(&self) -> &DemandSet {
&self.demands
}
pub(crate) fn demanded(&self) -> &[SlotKey] {
&self.demanded
}
fn retains_grant(&self, grant_id: &[u8; 32], audience_handle: &[u8; 32]) -> bool {
self.demanded.iter().any(|key| {
matches!(
key.scope.scope(),
CapabilityAudienceScope::Grant {
grant_id: retained_id,
audience_handle: retained_handle,
} if retained_id == grant_id && retained_handle == audience_handle
)
})
}
}
#[derive(Default)]
pub(crate) struct CapabilityIndex {
entries: BTreeMap<CapabilityAuthorityId, Arc<CapabilityRouteHandle>>,
}
impl CapabilityIndex {
fn get(&self, capability: &CapabilityAuthorityId) -> Option<&Arc<CapabilityRouteHandle>> {
self.entries.get(capability)
}
#[allow(dead_code)] pub(crate) fn len(&self) -> usize {
self.entries.len()
}
fn with(&self, handle: Arc<CapabilityRouteHandle>) -> Self {
let mut entries = self.entries.clone();
entries.insert(handle.capability, handle);
Self { entries }
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[allow(dead_code)] pub(crate) enum RouteLookup {
Warm,
Cold(ColdReason),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum SnapshotOrder {
Older,
Current,
Newer,
}
fn encode_revision(revision: GrantMovementFence) -> u64 {
match revision {
GrantMovementFence::Publication(generation) => generation,
GrantMovementFence::Terminal => u64::MAX,
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[allow(dead_code)] pub(crate) enum ColdReason {
NoOwnerScope,
SnapshotSuperseded,
Refused(DemandRefused),
}
#[allow(dead_code)] pub(crate) struct OrgRoutingState {
family: RoutingFamily,
credentials: FamilyDiscoveryCredentials,
grants_by_capability: BTreeMap<CapabilityAuthorityId, Vec<Arc<OrgCapabilityGrant>>>,
index: ArcSwap<CapabilityIndex>,
mutate: parking_lot::Mutex<()>,
snapshot_high_water: AtomicU64,
mutate_acquisitions: AtomicU64,
}
#[allow(dead_code)] impl OrgRoutingState {
pub(crate) fn new(family: RoutingFamily, credentials: FamilyDiscoveryCredentials) -> Self {
let mut grants_by_capability: BTreeMap<
CapabilityAuthorityId,
Vec<Arc<OrgCapabilityGrant>>,
> = BTreeMap::new();
for grant in &credentials.grants {
grants_by_capability
.entry(grant.capability)
.or_default()
.push(grant.clone());
}
Self {
family,
credentials,
grants_by_capability,
index: ArcSwap::from_pointee(CapabilityIndex::default()),
snapshot_high_water: AtomicU64::new(0),
mutate: parking_lot::Mutex::new(()),
mutate_acquisitions: AtomicU64::new(0),
}
}
fn grants_for(
&self,
capability: &CapabilityAuthorityId,
) -> impl Iterator<Item = &Arc<OrgCapabilityGrant>> {
self.grants_by_capability
.get(capability)
.into_iter()
.flatten()
}
fn demand_set(
&self,
capability: &CapabilityAuthorityId,
leased: &ConsumerGrantSnapshot,
) -> Option<Vec<SlotKey>> {
demand_set_from(
&self.credentials,
self.grants_for(capability),
capability,
leased,
)
}
fn leases_current(
&self,
entry: &CapabilityRouteHandle,
capability: &CapabilityAuthorityId,
leased: &ConsumerGrantSnapshot,
) -> bool {
for key in &entry.demanded {
let CapabilityAudienceScope::Grant {
grant_id,
audience_handle,
} = key.scope.scope()
else {
continue;
};
let still_leased = leased
.get(grant_id)
.is_some_and(|record| record.audience_handle() == audience_handle);
if !still_leased {
return false;
}
}
for grant in self.grants_for(capability) {
let Some(GrantDemand::Leased(audience_handle)) =
classify(grant, capability, &self.credentials.acting_org, leased)
else {
continue;
};
if !entry.retains_grant(&grant.grant_id, &audience_handle) {
return false;
}
}
true
}
pub(crate) fn warm(
&self,
capability: &CapabilityAuthorityId,
) -> Option<Arc<CapabilityRouteHandle>> {
self.index.load().get(capability).cloned()
}
pub(crate) fn route_handle(
&self,
capability: &CapabilityAuthorityId,
leased: &ConsumerGrantSnapshot,
) -> RouteLookup {
match self.note_snapshot(leased) {
SnapshotOrder::Older => return RouteLookup::Cold(ColdReason::SnapshotSuperseded),
SnapshotOrder::Current | SnapshotOrder::Newer => {}
}
if let Some(entry) = self.warm(capability) {
if self.leases_current(&entry, capability, leased) {
return RouteLookup::Warm;
}
}
self.acquire(capability, leased)
}
fn note_snapshot(&self, leased: &ConsumerGrantSnapshot) -> SnapshotOrder {
let revision = encode_revision(leased.revision());
let previous = self
.snapshot_high_water
.fetch_max(revision, Ordering::AcqRel);
match revision.cmp(&previous) {
std::cmp::Ordering::Less => SnapshotOrder::Older,
std::cmp::Ordering::Equal => SnapshotOrder::Current,
std::cmp::Ordering::Greater => SnapshotOrder::Newer,
}
}
fn acquire(
&self,
capability: &CapabilityAuthorityId,
leased: &ConsumerGrantSnapshot,
) -> RouteLookup {
let mutate = self.mutate.lock();
self.mutate_acquisitions.fetch_add(1, Ordering::AcqRel);
if encode_revision(leased.revision()) < self.snapshot_high_water.load(Ordering::Acquire) {
return RouteLookup::Cold(ColdReason::SnapshotSuperseded);
}
let current = self.index.load().get(capability).cloned();
match current {
Some(current) => {
if self.leases_current(¤t, capability, leased) {
return RouteLookup::Warm;
}
self.acquire_under_mutate(&mutate, capability, leased, Some(¤t))
}
None => self.acquire_under_mutate(&mutate, capability, leased, None),
}
}
fn acquire_under_mutate(
&self,
_mutate: &parking_lot::MutexGuard<'_, ()>,
capability: &CapabilityAuthorityId,
leased: &ConsumerGrantSnapshot,
superseded: Option<&Arc<CapabilityRouteHandle>>,
) -> RouteLookup {
let Some(keys) = self.demand_set(capability, leased) else {
return RouteLookup::Cold(ColdReason::NoOwnerScope);
};
let revision = encode_revision(leased.revision());
if superseded.is_some_and(|stale| stale.derived_at == revision) {
return RouteLookup::Cold(ColdReason::SnapshotSuperseded);
}
let acquired = match superseded {
Some(stale) => stale
.demands()
.replace(keys.clone())
.map_err(|refusal| match refusal {
ReplaceRefused::Demand(refusal) => refusal,
ReplaceRefused::Superseded => DemandRefused::NodeAtCapacity,
}),
None => self.family.demand_set(keys.clone()),
};
match acquired {
Ok(demands) => {
let handle = Arc::new(CapabilityRouteHandle {
capability: *capability,
derived_at: revision,
demanded: keys,
demands,
});
self.index.store(Arc::new(self.index.load().with(handle)));
self.snapshot_high_water
.fetch_max(revision, Ordering::AcqRel);
RouteLookup::Warm
}
Err(refusal) => RouteLookup::Cold(ColdReason::Refused(refusal)),
}
}
pub(crate) fn entries(&self) -> usize {
self.index.load().len()
}
pub(crate) fn handles(&self) -> usize {
self.family.handles()
}
pub(crate) fn mutate_acquisitions(&self) -> u64 {
self.mutate_acquisitions.load(Ordering::Acquire)
}
#[cfg(test)]
pub(crate) fn mutate_lock_for_test(&self) -> &parking_lot::Mutex<()> {
&self.mutate
}
}
#[cfg(test)]
#[path = "org_routing_state_tests.rs"]
mod tests;