1use crate::{
4 changed, closed, config::RegistrationPolicy, lock, millis, operation, timeout, Connection,
5 Core, Gate, OwnerPhase,
6};
7use nadisc::managed::{
8 DiscoveryError, ErrorCode, RegistrationEvidence, RegistrationPlan, Result, WriteAdmission,
9};
10use std::{
11 sync::{
12 atomic::{AtomicBool, AtomicU64, Ordering},
13 Arc, Mutex,
14 },
15 time::Duration,
16};
17use tokio::{sync::watch, time::Instant};
18
19#[derive(Clone)]
21pub struct Eligibility {
22 inner: Arc<Mutex<bool>>,
23 changes: watch::Sender<bool>,
24 check: Option<Arc<dyn Fn() -> bool + Send + Sync>>,
25 valid_until: Arc<Mutex<Option<Instant>>>,
26 source: Option<std::sync::Weak<()>>,
27}
28
29pub struct EligibilityController {
31 signal: Eligibility,
32 _lifetime: Arc<()>,
33}
34impl EligibilityController {
35 pub fn update(&self, ready: bool, valid_for: Duration) -> Result<()> {
43 if valid_for < Duration::from_millis(1) || valid_for > Duration::from_secs(30) {
44 self.signal.set(false);
45 return Err(DiscoveryError::new(
46 ErrorCode::InvalidConfig,
47 "eligibility validity is out of bounds",
48 ));
49 }
50 let mut value = lock(&self.signal.inner);
51 *value = ready;
52 *lock(&self.signal.valid_until) = ready.then(|| Instant::now() + valid_for);
53 self.signal.changes.send_replace(ready);
54 Ok(())
55 }
56}
57impl Drop for EligibilityController {
58 fn drop(&mut self) {
64 self.signal.set(false);
65 }
66}
67
68impl Eligibility {
69 pub fn new() -> Self {
75 Self {
76 inner: Arc::new(Mutex::new(false)),
77 changes: watch::channel(false).0,
78 check: None,
79 valid_until: Arc::new(Mutex::new(None)),
80 source: None,
81 }
82 }
83
84 pub fn channel() -> (EligibilityController, Self) {
90 let lifetime = Arc::new(());
91 let mut signal = Self::new();
92 signal.source = Some(Arc::downgrade(&lifetime));
93 (
94 EligibilityController {
95 signal: signal.clone(),
96 _lifetime: lifetime,
97 },
98 signal,
99 )
100 }
101
102 pub fn set(&self, ready: bool) {
109 let mut value = lock(&self.inner);
110 *value = ready;
111 *lock(&self.valid_until) = ready.then(|| Instant::now() + Duration::from_secs(30));
112 self.changes.send_replace(ready);
113 }
114
115 pub fn is_ready(&self) -> bool {
121 self.current(*lock(&self.inner))
122 }
123
124 pub fn from_check(check: impl Fn() -> bool + Send + Sync + 'static) -> Self {
131 Self {
132 inner: Arc::new(Mutex::new(true)),
133 changes: watch::channel(true).0,
134 check: Some(Arc::new(check)),
135 valid_until: Arc::new(Mutex::new(None)),
136 source: None,
137 }
138 }
139
140 fn current(&self, ready: bool) -> bool {
147 ready
148 && self
149 .source
150 .as_ref()
151 .is_none_or(|source| source.strong_count() > 0)
152 && self.check.as_ref().map_or_else(
153 || lock(&self.valid_until).is_some_and(|until| Instant::now() < until),
154 |check| check(),
155 )
156 }
157}
158
159impl Default for Eligibility {
160 fn default() -> Self {
166 Self::new()
167 }
168}
169
170pub struct RegistrationSpec {
172 pub connection_ref: String,
174 pub plan: RegistrationPlan,
176 pub policy: RegistrationPolicy,
178 pub eligibility: Eligibility,
180}
181
182#[derive(Clone, Copy, Debug, Eq, PartialEq)]
184pub enum RegistrationState {
185 Planned,
187 Registering,
189 OutcomeUnknown,
191 PendingHealth,
193 Active,
195 Suspended,
197 Withdrawing,
199 Residual,
201 Removed,
203}
204
205#[derive(Clone, Debug)]
207pub struct RegistrationStatus {
208 pub id: String,
210 pub connection: String,
212 pub state: RegistrationState,
214 pub startup_confirmed: bool,
216 pub readiness_critical: bool,
218 pub observed_at: Option<Instant>,
220 pub valid_until: Option<Instant>,
222 pub error: Option<DiscoveryError>,
224}
225
226impl RegistrationStatus {
227 pub fn is_ready(&self) -> bool {
233 self.state == RegistrationState::Active
234 && self
235 .valid_until
236 .is_some_and(|deadline| Instant::now() < deadline)
237 }
238}
239
240pub(crate) struct Record {
241 pub spec: RegistrationSpec,
242 pub connection: Arc<Connection>,
243 pub status: Mutex<RegistrationStatus>,
244 pub started: AtomicBool,
245 pub confirmed: AtomicBool,
246 health_revocation: AtomicU64,
247 attempted: AtomicBool,
248 uncertain: AtomicBool,
249 unsettled_renewal: AtomicBool,
250}
251
252impl Record {
253 pub(crate) fn new(spec: RegistrationSpec, connection: Arc<Connection>) -> Self {
261 let health_revocation = connection.revocation.load(Ordering::Acquire);
262 Self {
263 status: Mutex::new(RegistrationStatus {
264 id: spec.plan.id.clone(),
265 connection: spec.connection_ref.clone(),
266 state: RegistrationState::Planned,
267 startup_confirmed: false,
268 readiness_critical: spec.policy.readiness_critical,
269 observed_at: None,
270 valid_until: None,
271 error: None,
272 }),
273 spec,
274 connection,
275 started: AtomicBool::new(false),
276 confirmed: AtomicBool::new(false),
277 health_revocation: AtomicU64::new(health_revocation),
278 attempted: AtomicBool::new(false),
279 uncertain: AtomicBool::new(false),
280 unsettled_renewal: AtomicBool::new(false),
281 }
282 }
283
284 pub(crate) fn current_status(&self, gate: &Gate) -> RegistrationStatus {
291 let mut status = lock(&self.status).clone();
292 if self.health_revocation.load(Ordering::Acquire)
294 != self.connection.revocation.load(Ordering::Acquire)
295 {
296 status.valid_until = None;
297 if matches!(
298 status.state,
299 RegistrationState::Active | RegistrationState::PendingHealth
300 ) {
301 status.state = RegistrationState::Suspended;
302 status.error = Some(DiscoveryError::new(
303 ErrorCode::Unauthenticated,
304 "registration evidence belongs to a revoked identity",
305 ));
306 }
307 }
308 if !self.spec.eligibility.is_ready() || gate.failed || gate.close_started {
309 status.valid_until = None;
310 }
311 status
312 }
313}
314
315struct Admission<'a> {
316 core: &'a Core,
317 record: &'a Record,
318}
319impl WriteAdmission for Admission<'_> {
320 fn submit(&self, submit: &mut dyn FnMut()) -> Result<()> {
327 let (core, record) = (self.core, self.record);
328 let gate = lock(&core.gate);
329 let eligible = lock(&record.spec.eligibility.inner);
330 if gate.failed
331 || gate.close_started
332 || !matches!(gate.phase, OwnerPhase::Starting | OwnerPhase::Serving)
333 || !record.spec.eligibility.current(*eligible)
334 || core.registrations_stop.is_cancelled()
335 {
336 return Err(closed());
337 }
338 record.attempted.store(true, Ordering::Release);
340 record.uncertain.store(true, Ordering::Release);
341 submit();
342 Ok(())
343 }
344}
345
346pub(crate) async fn run(core: Arc<Core>, record: Arc<Record>) {
354 let mut eligibility = record.spec.eligibility.changes.subscribe();
355 let mut prepared = record.connection.prepared.subscribe();
356 let admission = Admission {
357 core: &core,
358 record: &record,
359 };
360 let monitor = millis(record.spec.policy.monitor_interval_ms);
361 let interval = record
362 .connection
363 .provider
364 .renewal_interval(&record.spec.plan)
365 .unwrap_or(monitor)
366 .min(monitor);
367 let idempotent_renewal = record
368 .connection
369 .provider
370 .supports_idempotent_renewal(&record.spec.plan);
371 let mut registered = false;
372 let mut renewal_allowed = false;
373 let mut registration_blocked = false;
374 let mut backoff = record.connection.policy.initial_backoff_ms;
375 let mut attempt = 0u64;
376 loop {
377 let mut renewal_failed = false;
378 if core.registrations_stop.is_cancelled() {
379 return;
380 }
381 if !record.spec.eligibility.is_ready() || !*prepared.borrow() {
382 if registered {
383 set_state(&core, &record, RegistrationState::Suspended, None);
385 withdraw(
386 record.clone(),
387 Instant::now() + millis(record.spec.policy.deregister_timeout_ms),
388 millis(record.spec.policy.deregister_timeout_ms),
389 )
390 .await;
391 registered = false;
392 registration_blocked |= lock(&record.status).state != RegistrationState::Removed;
393 changed(&core);
394 }
395 tokio::select! {
396 _ = core.registrations_stop.cancelled() => return,
397 _ = eligibility.changed() => {},
398 _ = prepared.changed() => {},
399 _ = tokio::time::sleep(interval) => {},
400 }
401 continue;
402 }
403 if !registered && !registration_blocked {
404 set_state(&core, &record, RegistrationState::Registering, None);
405 let result = operation(
406 &core,
407 Instant::now() + millis(record.connection.policy.request_timeout_ms),
408 record
409 .connection
410 .provider
411 .register(&record.spec.plan, &admission),
412 )
413 .await;
414 match result {
415 Ok(()) => {
416 record.uncertain.store(false, Ordering::Release);
417 registered = true;
418 renewal_allowed = true;
419 set_state(&core, &record, RegistrationState::PendingHealth, None);
420 }
421 Err(error) => {
422 {
423 let gate = lock(&core.gate);
424 revoke_identity(&core, &gate, &record.connection, &error);
425 }
426 if definite_rejection(&error) {
427 record.uncertain.store(false, Ordering::Release);
428 }
429 registration_blocked = record.uncertain.load(Ordering::Acquire)
430 || matches!(
431 error.code,
432 ErrorCode::InvalidConfig | ErrorCode::Unsupported
433 );
434 set_state(
436 &core,
437 &record,
438 if record.uncertain.load(Ordering::Acquire) {
439 RegistrationState::OutcomeUnknown
440 } else {
441 RegistrationState::Suspended
442 },
443 Some(error),
444 );
445 }
446 }
447 }
448 if registered
449 && renewal_allowed
450 && (!record.uncertain.load(Ordering::Acquire) || idempotent_renewal)
452 && record
453 .connection
454 .provider
455 .renewal_interval(&record.spec.plan)
456 .is_some()
457 {
458 let result = operation(
459 &core,
460 Instant::now() + millis(record.connection.policy.request_timeout_ms),
461 record
462 .connection
463 .provider
464 .renew(&record.spec.plan, true, &admission),
465 )
466 .await;
467 match result {
468 Ok(()) => record.uncertain.store(false, Ordering::Release),
470 Err(error) => {
471 renewal_failed = true;
472 {
473 let gate = lock(&core.gate);
474 revoke_identity(&core, &gate, &record.connection, &error);
475 }
476 if definite_rejection(&error) {
477 record.uncertain.store(false, Ordering::Release);
478 }
479 if record.uncertain.load(Ordering::Acquire) {
480 record.unsettled_renewal.store(true, Ordering::Release);
482 renewal_allowed = idempotent_renewal && !error.revokes_evidence();
483 set_state(
484 &core,
485 &record,
486 RegistrationState::OutcomeUnknown,
487 Some(error),
488 );
489 } else {
490 renewal_allowed = false;
491 let mut status = lock(&record.status);
492 if error.revokes_evidence() {
493 status.valid_until = None;
494 status.state = RegistrationState::Suspended;
495 }
496 status.error = Some(error);
497 }
498 }
499 }
500 }
501 if registered && !record.uncertain.load(Ordering::Acquire) {
502 let revocation = record.connection.revocation.load(Ordering::Acquire);
503 let result = operation(
504 &core,
505 Instant::now() + millis(record.connection.policy.request_timeout_ms),
506 record.connection.provider.inspect(&record.spec.plan),
507 )
508 .await;
509 let gate = lock(&core.gate);
510 if core.registrations_stop.is_cancelled() || gate.close_started {
511 return;
512 }
513 if let Err(error) = &result {
514 if error.revokes_evidence() {
515 renewal_allowed = false;
516 }
517 revoke_identity(&core, &gate, &record.connection, error);
518 }
519 if revocation == record.connection.revocation.load(Ordering::Acquire) {
521 let mut status = lock(&record.status);
522 match result {
523 Ok(RegistrationEvidence::Healthy) if record.spec.eligibility.is_ready() => {
524 registration_blocked = false;
525 renewal_allowed = true;
526 let now = Instant::now();
527 status.state = RegistrationState::Active;
528 status.observed_at = Some(now);
529 status.valid_until =
530 Some(now + millis(record.spec.policy.observation_grace_ms));
531 status.error = None;
532 status.startup_confirmed = true;
533 record.confirmed.store(true, Ordering::Release);
534 record
535 .health_revocation
536 .store(revocation, Ordering::Release);
537 }
538 Ok(RegistrationEvidence::Disabled | RegistrationEvidence::Missing) => {
539 renewal_allowed = false;
541 registration_blocked = true;
542 status.state = RegistrationState::Suspended;
543 status.valid_until = None;
544 status.error = Some(DiscoveryError::new(
545 ErrorCode::NotReady,
546 "registration was disabled or removed",
547 ));
548 }
549 Ok(_) => {
550 renewal_allowed = true;
551 status.state = RegistrationState::PendingHealth;
552 status.valid_until = None;
553 status.error = None;
554 }
555 Err(error) => {
556 if error.revokes_evidence() {
558 status.state = RegistrationState::Suspended;
559 status.valid_until = None;
560 renewal_allowed = false;
561 }
562 status.error = Some(error);
563 }
564 }
565 }
566 }
567 changed(&core);
568 let delay = if registered && !renewal_failed {
569 backoff = record.connection.policy.initial_backoff_ms;
570 attempt = 0;
571 interval
572 } else {
573 let delay = crate::retry_delay(&core, &record.connection, backoff, attempt);
574 attempt = attempt.wrapping_add(1);
575 backoff = backoff
576 .saturating_mul(2)
577 .min(record.connection.policy.max_backoff_ms);
578 delay.max(interval)
580 };
581 let next_attempt = Instant::now() + delay;
582 loop {
583 tokio::select! {
584 _ = core.registrations_stop.cancelled() => return,
585 _ = eligibility.changed() => if !record.spec.eligibility.is_ready() { break; },
587 _ = tokio::time::sleep_until(next_attempt) => break,
588 }
589 }
590 }
591}
592
593fn revoke_identity(core: &Core, gate: &Gate, connection: &Connection, error: &DiscoveryError) {
603 if matches!(
604 error.code,
605 ErrorCode::Unauthenticated | ErrorCode::PermissionDenied | ErrorCode::PermissionFiltered
606 ) {
607 crate::revoke_connection(core, gate, connection);
608 }
609}
610
611fn definite_rejection(error: &DiscoveryError) -> bool {
618 matches!(
619 error.code,
620 ErrorCode::InvalidConfig
621 | ErrorCode::Unauthenticated
622 | ErrorCode::PermissionDenied
623 | ErrorCode::Conflict
624 | ErrorCode::Unsupported
625 )
626}
627
628fn set_state(
638 core: &Core,
639 record: &Record,
640 state: RegistrationState,
641 error: Option<DiscoveryError>,
642) {
643 let mut status = lock(&record.status);
644 status.state = state;
645 status.error = error;
646 if state != RegistrationState::Active {
647 status.valid_until = None;
648 }
649 drop(status);
650 changed(core);
651}
652
653pub(crate) async fn withdraw(record: Arc<Record>, deadline: Instant, share: Duration) {
662 if !record.attempted.load(Ordering::Acquire) {
663 lock(&record.status).state = RegistrationState::Removed;
664 return;
665 }
666 if lock(&record.status).state == RegistrationState::Removed {
667 return;
668 }
669 let uncertain = record.uncertain.load(Ordering::Acquire)
671 || record.unsettled_renewal.load(Ordering::Acquire);
672 lock(&record.status).state = RegistrationState::Withdrawing;
673 let until =
674 deadline.min(Instant::now() + share.min(millis(record.spec.policy.deregister_timeout_ms)));
675 let result = tokio::time::timeout_at(
676 until,
677 record.connection.provider.deregister(&record.spec.plan),
678 )
679 .await
680 .unwrap_or_else(|_| Err(timeout()));
681 let mut status = lock(&record.status);
682 status.valid_until = None;
683 match result {
684 Ok(()) if !uncertain => {
685 status.state = RegistrationState::Removed;
686 status.error = None;
687 record.attempted.store(false, Ordering::Release);
688 }
689 Ok(()) => {
690 status.state = RegistrationState::Residual;
691 status.error = Some(DiscoveryError::new(
692 ErrorCode::OutcomeUnknown,
693 "an earlier write may still complete after cleanup",
694 ));
695 }
696 Err(error) => {
697 record.uncertain.store(true, Ordering::Release);
698 status.state = RegistrationState::Residual;
699 status.error = Some(error);
700 }
701 }
702}
703
704pub(crate) fn finish_withdrawal(record: &Record) {
711 let mut status = lock(&record.status);
712 status.valid_until = None;
713 if !record.attempted.load(Ordering::Acquire) {
714 status.state = RegistrationState::Removed;
715 } else if !matches!(
716 status.state,
717 RegistrationState::Removed | RegistrationState::Residual
718 ) {
719 status.state = RegistrationState::Residual;
720 status.error = Some(DiscoveryError::new(
721 ErrorCode::OutcomeUnknown,
722 "registration cleanup did not complete",
723 ));
724 }
725}