use crate::{
changed, closed, config::RegistrationPolicy, lock, millis, operation, timeout, Connection,
Core, Gate, OwnerPhase,
};
use nadisc::managed::{
DiscoveryError, ErrorCode, RegistrationEvidence, RegistrationPlan, Result, WriteAdmission,
};
use std::{
sync::{
atomic::{AtomicBool, AtomicU64, Ordering},
Arc, Mutex,
},
time::Duration,
};
use tokio::{sync::watch, time::Instant};
#[derive(Clone)]
pub struct Eligibility {
inner: Arc<Mutex<bool>>,
changes: watch::Sender<bool>,
check: Option<Arc<dyn Fn() -> bool + Send + Sync>>,
valid_until: Arc<Mutex<Option<Instant>>>,
source: Option<std::sync::Weak<()>>,
}
pub struct EligibilityController {
signal: Eligibility,
_lifetime: Arc<()>,
}
impl EligibilityController {
pub fn update(&self, ready: bool, valid_for: Duration) -> Result<()> {
if valid_for < Duration::from_millis(1) || valid_for > Duration::from_secs(30) {
self.signal.set(false);
return Err(DiscoveryError::new(
ErrorCode::InvalidConfig,
"eligibility validity is out of bounds",
));
}
let mut value = lock(&self.signal.inner);
*value = ready;
*lock(&self.signal.valid_until) = ready.then(|| Instant::now() + valid_for);
self.signal.changes.send_replace(ready);
Ok(())
}
}
impl Drop for EligibilityController {
fn drop(&mut self) {
self.signal.set(false);
}
}
impl Eligibility {
pub fn new() -> Self {
Self {
inner: Arc::new(Mutex::new(false)),
changes: watch::channel(false).0,
check: None,
valid_until: Arc::new(Mutex::new(None)),
source: None,
}
}
pub fn channel() -> (EligibilityController, Self) {
let lifetime = Arc::new(());
let mut signal = Self::new();
signal.source = Some(Arc::downgrade(&lifetime));
(
EligibilityController {
signal: signal.clone(),
_lifetime: lifetime,
},
signal,
)
}
pub fn set(&self, ready: bool) {
let mut value = lock(&self.inner);
*value = ready;
*lock(&self.valid_until) = ready.then(|| Instant::now() + Duration::from_secs(30));
self.changes.send_replace(ready);
}
pub fn is_ready(&self) -> bool {
self.current(*lock(&self.inner))
}
pub fn from_check(check: impl Fn() -> bool + Send + Sync + 'static) -> Self {
Self {
inner: Arc::new(Mutex::new(true)),
changes: watch::channel(true).0,
check: Some(Arc::new(check)),
valid_until: Arc::new(Mutex::new(None)),
source: None,
}
}
fn current(&self, ready: bool) -> bool {
ready
&& self
.source
.as_ref()
.is_none_or(|source| source.strong_count() > 0)
&& self.check.as_ref().map_or_else(
|| lock(&self.valid_until).is_some_and(|until| Instant::now() < until),
|check| check(),
)
}
}
impl Default for Eligibility {
fn default() -> Self {
Self::new()
}
}
pub struct RegistrationSpec {
pub connection_ref: String,
pub plan: RegistrationPlan,
pub policy: RegistrationPolicy,
pub eligibility: Eligibility,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RegistrationState {
Planned,
Registering,
OutcomeUnknown,
PendingHealth,
Active,
Suspended,
Withdrawing,
Residual,
Removed,
}
#[derive(Clone, Debug)]
pub struct RegistrationStatus {
pub id: String,
pub connection: String,
pub state: RegistrationState,
pub startup_confirmed: bool,
pub readiness_critical: bool,
pub observed_at: Option<Instant>,
pub valid_until: Option<Instant>,
pub error: Option<DiscoveryError>,
}
impl RegistrationStatus {
pub fn is_ready(&self) -> bool {
self.state == RegistrationState::Active
&& self
.valid_until
.is_some_and(|deadline| Instant::now() < deadline)
}
}
pub(crate) struct Record {
pub spec: RegistrationSpec,
pub connection: Arc<Connection>,
pub status: Mutex<RegistrationStatus>,
pub started: AtomicBool,
pub confirmed: AtomicBool,
health_revocation: AtomicU64,
attempted: AtomicBool,
uncertain: AtomicBool,
unsettled_renewal: AtomicBool,
}
impl Record {
pub(crate) fn new(spec: RegistrationSpec, connection: Arc<Connection>) -> Self {
let health_revocation = connection.revocation.load(Ordering::Acquire);
Self {
status: Mutex::new(RegistrationStatus {
id: spec.plan.id.clone(),
connection: spec.connection_ref.clone(),
state: RegistrationState::Planned,
startup_confirmed: false,
readiness_critical: spec.policy.readiness_critical,
observed_at: None,
valid_until: None,
error: None,
}),
spec,
connection,
started: AtomicBool::new(false),
confirmed: AtomicBool::new(false),
health_revocation: AtomicU64::new(health_revocation),
attempted: AtomicBool::new(false),
uncertain: AtomicBool::new(false),
unsettled_renewal: AtomicBool::new(false),
}
}
pub(crate) fn current_status(&self, gate: &Gate) -> RegistrationStatus {
let mut status = lock(&self.status).clone();
if self.health_revocation.load(Ordering::Acquire)
!= self.connection.revocation.load(Ordering::Acquire)
{
status.valid_until = None;
if matches!(
status.state,
RegistrationState::Active | RegistrationState::PendingHealth
) {
status.state = RegistrationState::Suspended;
status.error = Some(DiscoveryError::new(
ErrorCode::Unauthenticated,
"registration evidence belongs to a revoked identity",
));
}
}
if !self.spec.eligibility.is_ready() || gate.failed || gate.close_started {
status.valid_until = None;
}
status
}
}
struct Admission<'a> {
core: &'a Core,
record: &'a Record,
}
impl WriteAdmission for Admission<'_> {
fn submit(&self, submit: &mut dyn FnMut()) -> Result<()> {
let (core, record) = (self.core, self.record);
let gate = lock(&core.gate);
let eligible = lock(&record.spec.eligibility.inner);
if gate.failed
|| gate.close_started
|| !matches!(gate.phase, OwnerPhase::Starting | OwnerPhase::Serving)
|| !record.spec.eligibility.current(*eligible)
|| core.registrations_stop.is_cancelled()
{
return Err(closed());
}
record.attempted.store(true, Ordering::Release);
record.uncertain.store(true, Ordering::Release);
submit();
Ok(())
}
}
pub(crate) async fn run(core: Arc<Core>, record: Arc<Record>) {
let mut eligibility = record.spec.eligibility.changes.subscribe();
let mut prepared = record.connection.prepared.subscribe();
let admission = Admission {
core: &core,
record: &record,
};
let monitor = millis(record.spec.policy.monitor_interval_ms);
let interval = record
.connection
.provider
.renewal_interval(&record.spec.plan)
.unwrap_or(monitor)
.min(monitor);
let idempotent_renewal = record
.connection
.provider
.supports_idempotent_renewal(&record.spec.plan);
let mut registered = false;
let mut renewal_allowed = false;
let mut registration_blocked = false;
let mut backoff = record.connection.policy.initial_backoff_ms;
let mut attempt = 0u64;
loop {
let mut renewal_failed = false;
if core.registrations_stop.is_cancelled() {
return;
}
if !record.spec.eligibility.is_ready() || !*prepared.borrow() {
if registered {
set_state(&core, &record, RegistrationState::Suspended, None);
withdraw(
record.clone(),
Instant::now() + millis(record.spec.policy.deregister_timeout_ms),
millis(record.spec.policy.deregister_timeout_ms),
)
.await;
registered = false;
registration_blocked |= lock(&record.status).state != RegistrationState::Removed;
changed(&core);
}
tokio::select! {
_ = core.registrations_stop.cancelled() => return,
_ = eligibility.changed() => {},
_ = prepared.changed() => {},
_ = tokio::time::sleep(interval) => {},
}
continue;
}
if !registered && !registration_blocked {
set_state(&core, &record, RegistrationState::Registering, None);
let result = operation(
&core,
Instant::now() + millis(record.connection.policy.request_timeout_ms),
record
.connection
.provider
.register(&record.spec.plan, &admission),
)
.await;
match result {
Ok(()) => {
record.uncertain.store(false, Ordering::Release);
registered = true;
renewal_allowed = true;
set_state(&core, &record, RegistrationState::PendingHealth, None);
}
Err(error) => {
{
let gate = lock(&core.gate);
revoke_identity(&core, &gate, &record.connection, &error);
}
if definite_rejection(&error) {
record.uncertain.store(false, Ordering::Release);
}
registration_blocked = record.uncertain.load(Ordering::Acquire)
|| matches!(
error.code,
ErrorCode::InvalidConfig | ErrorCode::Unsupported
);
set_state(
&core,
&record,
if record.uncertain.load(Ordering::Acquire) {
RegistrationState::OutcomeUnknown
} else {
RegistrationState::Suspended
},
Some(error),
);
}
}
}
if registered
&& renewal_allowed
&& (!record.uncertain.load(Ordering::Acquire) || idempotent_renewal)
&& record
.connection
.provider
.renewal_interval(&record.spec.plan)
.is_some()
{
let result = operation(
&core,
Instant::now() + millis(record.connection.policy.request_timeout_ms),
record
.connection
.provider
.renew(&record.spec.plan, true, &admission),
)
.await;
match result {
Ok(()) => record.uncertain.store(false, Ordering::Release),
Err(error) => {
renewal_failed = true;
{
let gate = lock(&core.gate);
revoke_identity(&core, &gate, &record.connection, &error);
}
if definite_rejection(&error) {
record.uncertain.store(false, Ordering::Release);
}
if record.uncertain.load(Ordering::Acquire) {
record.unsettled_renewal.store(true, Ordering::Release);
renewal_allowed = idempotent_renewal && !error.revokes_evidence();
set_state(
&core,
&record,
RegistrationState::OutcomeUnknown,
Some(error),
);
} else {
renewal_allowed = false;
let mut status = lock(&record.status);
if error.revokes_evidence() {
status.valid_until = None;
status.state = RegistrationState::Suspended;
}
status.error = Some(error);
}
}
}
}
if registered && !record.uncertain.load(Ordering::Acquire) {
let revocation = record.connection.revocation.load(Ordering::Acquire);
let result = operation(
&core,
Instant::now() + millis(record.connection.policy.request_timeout_ms),
record.connection.provider.inspect(&record.spec.plan),
)
.await;
let gate = lock(&core.gate);
if core.registrations_stop.is_cancelled() || gate.close_started {
return;
}
if let Err(error) = &result {
if error.revokes_evidence() {
renewal_allowed = false;
}
revoke_identity(&core, &gate, &record.connection, error);
}
if revocation == record.connection.revocation.load(Ordering::Acquire) {
let mut status = lock(&record.status);
match result {
Ok(RegistrationEvidence::Healthy) if record.spec.eligibility.is_ready() => {
registration_blocked = false;
renewal_allowed = true;
let now = Instant::now();
status.state = RegistrationState::Active;
status.observed_at = Some(now);
status.valid_until =
Some(now + millis(record.spec.policy.observation_grace_ms));
status.error = None;
status.startup_confirmed = true;
record.confirmed.store(true, Ordering::Release);
record
.health_revocation
.store(revocation, Ordering::Release);
}
Ok(RegistrationEvidence::Disabled | RegistrationEvidence::Missing) => {
renewal_allowed = false;
registration_blocked = true;
status.state = RegistrationState::Suspended;
status.valid_until = None;
status.error = Some(DiscoveryError::new(
ErrorCode::NotReady,
"registration was disabled or removed",
));
}
Ok(_) => {
renewal_allowed = true;
status.state = RegistrationState::PendingHealth;
status.valid_until = None;
status.error = None;
}
Err(error) => {
if error.revokes_evidence() {
status.state = RegistrationState::Suspended;
status.valid_until = None;
renewal_allowed = false;
}
status.error = Some(error);
}
}
}
}
changed(&core);
let delay = if registered && !renewal_failed {
backoff = record.connection.policy.initial_backoff_ms;
attempt = 0;
interval
} else {
let delay = crate::retry_delay(&core, &record.connection, backoff, attempt);
attempt = attempt.wrapping_add(1);
backoff = backoff
.saturating_mul(2)
.min(record.connection.policy.max_backoff_ms);
delay.max(interval)
};
let next_attempt = Instant::now() + delay;
loop {
tokio::select! {
_ = core.registrations_stop.cancelled() => return,
_ = eligibility.changed() => if !record.spec.eligibility.is_ready() { break; },
_ = tokio::time::sleep_until(next_attempt) => break,
}
}
}
}
fn revoke_identity(core: &Core, gate: &Gate, connection: &Connection, error: &DiscoveryError) {
if matches!(
error.code,
ErrorCode::Unauthenticated | ErrorCode::PermissionDenied | ErrorCode::PermissionFiltered
) {
crate::revoke_connection(core, gate, connection);
}
}
fn definite_rejection(error: &DiscoveryError) -> bool {
matches!(
error.code,
ErrorCode::InvalidConfig
| ErrorCode::Unauthenticated
| ErrorCode::PermissionDenied
| ErrorCode::Conflict
| ErrorCode::Unsupported
)
}
fn set_state(
core: &Core,
record: &Record,
state: RegistrationState,
error: Option<DiscoveryError>,
) {
let mut status = lock(&record.status);
status.state = state;
status.error = error;
if state != RegistrationState::Active {
status.valid_until = None;
}
drop(status);
changed(core);
}
pub(crate) async fn withdraw(record: Arc<Record>, deadline: Instant, share: Duration) {
if !record.attempted.load(Ordering::Acquire) {
lock(&record.status).state = RegistrationState::Removed;
return;
}
if lock(&record.status).state == RegistrationState::Removed {
return;
}
let uncertain = record.uncertain.load(Ordering::Acquire)
|| record.unsettled_renewal.load(Ordering::Acquire);
lock(&record.status).state = RegistrationState::Withdrawing;
let until =
deadline.min(Instant::now() + share.min(millis(record.spec.policy.deregister_timeout_ms)));
let result = tokio::time::timeout_at(
until,
record.connection.provider.deregister(&record.spec.plan),
)
.await
.unwrap_or_else(|_| Err(timeout()));
let mut status = lock(&record.status);
status.valid_until = None;
match result {
Ok(()) if !uncertain => {
status.state = RegistrationState::Removed;
status.error = None;
record.attempted.store(false, Ordering::Release);
}
Ok(()) => {
status.state = RegistrationState::Residual;
status.error = Some(DiscoveryError::new(
ErrorCode::OutcomeUnknown,
"an earlier write may still complete after cleanup",
));
}
Err(error) => {
record.uncertain.store(true, Ordering::Release);
status.state = RegistrationState::Residual;
status.error = Some(error);
}
}
}
pub(crate) fn finish_withdrawal(record: &Record) {
let mut status = lock(&record.status);
status.valid_until = None;
if !record.attempted.load(Ordering::Acquire) {
status.state = RegistrationState::Removed;
} else if !matches!(
status.state,
RegistrationState::Removed | RegistrationState::Residual
) {
status.state = RegistrationState::Residual;
status.error = Some(DiscoveryError::new(
ErrorCode::OutcomeUnknown,
"registration cleanup did not complete",
));
}
}