1use crate::{
4 changed, closed,
5 config::{invalid, validate_name, ConnectionPolicy, Limits},
6 lock, millis, operation,
7 registration::{Record, RegistrationSpec, RegistrationStatus},
8 spawn_owned, timeout, Connection, Core, DiscoveryHandle, Gate, Quota,
9};
10use nadisc::managed::{DiscoveryError, ErrorCode, ManagedProvider, Result};
11use std::{
12 collections::{BTreeMap, BTreeSet},
13 sync::{atomic::AtomicU64, Arc, Mutex},
14};
15use tokio::{
16 sync::{watch, Semaphore},
17 time::Instant,
18};
19use tokio_util::sync::CancellationToken;
20
21#[derive(Clone, Copy, Debug, Eq, PartialEq)]
23pub enum OwnerPhase {
24 Prepared,
26 Starting,
28 Serving,
30 Quiescing,
32 Closed,
34}
35
36pub struct ConnectionSpec {
38 pub name: String,
40 pub provider: Arc<dyn ManagedProvider>,
42 pub policy: ConnectionPolicy,
44}
45
46pub struct OwnerBuilder {
48 limits: Limits,
49 connections: BTreeMap<String, ConnectionSpec>,
50 default_connection: Option<String>,
51}
52
53impl Default for OwnerBuilder {
54 fn default() -> Self {
60 Self {
61 limits: Limits::default(),
62 connections: BTreeMap::new(),
63 default_connection: None,
64 }
65 }
66}
67
68impl OwnerBuilder {
69 pub fn limits(mut self, limits: Limits) -> Self {
76 self.limits = limits;
77 self
78 }
79
80 pub fn default_connection(mut self, name: impl Into<String>) -> Self {
87 self.default_connection = Some(name.into());
88 self
89 }
90
91 pub fn connection(mut self, spec: ConnectionSpec) -> Result<Self> {
98 validate_name(&spec.name)?;
99 spec.policy.validate()?;
100 if self.connections.contains_key(&spec.name) {
101 return Err(invalid("duplicate connection name"));
102 }
103 self.connections.insert(spec.name.clone(), spec);
104 Ok(self)
105 }
106
107 pub fn build(self) -> Result<DiscoveryOwner> {
113 self.limits.validate()?;
114 if self.connections.len() > self.limits.max_connections {
115 return Err(invalid("connection capacity exceeded"));
116 }
117 let default_connection = match self.default_connection {
118 Some(name) if self.connections.contains_key(&name) => Some(name),
119 Some(_) => return Err(invalid("unknown default connection")),
120 None => self
121 .connections
122 .contains_key("primary")
123 .then(|| "primary".into()),
124 };
125 let connections = self
126 .connections
127 .into_iter()
128 .map(|(name, spec)| {
129 let connection = Arc::new(Connection {
130 name: name.clone(),
131 provider: spec.provider,
132 policy: spec.policy,
133 slots: Mutex::new(BTreeMap::new()),
134 prepared: watch::channel(false).0,
135 error: Mutex::new(None),
136 epoch: AtomicU64::new(1),
137 revocation: AtomicU64::new(0),
138 queries: AtomicU64::new(0),
139 query_errors: AtomicU64::new(0),
140 query_micros: AtomicU64::new(0),
141 });
142 (name, connection)
143 })
144 .collect();
145 let core = Arc::new(Core {
146 id: uuid::Uuid::new_v4().to_string(),
147 gate: Mutex::new(Gate {
148 phase: OwnerPhase::Prepared,
149 registrations: BTreeMap::new(),
150 frozen: false,
151 close_started: false,
152 failed: false,
153 shutdown_deadline: None,
154 }),
155 connections,
156 default_connection,
157 operations: Semaphore::new(self.limits.max_parallel_operations),
158 watches: Semaphore::new(self.limits.max_watch_streams),
159 quota: Arc::new(Mutex::new(Quota::default())),
160 bridges: Arc::new(Semaphore::new(
161 self.limits
162 .max_service_watches_per_connection
163 .saturating_mul(self.limits.max_connections),
164 )),
165 transports: Arc::new(Semaphore::new(self.limits.max_transport_tasks)),
166 limits: self.limits,
167 tasks: Mutex::new(Vec::new()),
168 registration_tasks: Mutex::new(Vec::new()),
169 withdrawal_task: Mutex::new(None),
170 close_task: Mutex::new(None),
171 stop: CancellationToken::new(),
172 registrations_stop: CancellationToken::new(),
173 changes: watch::channel(0).0,
174 close: watch::channel(None).0,
175 withdrawn: watch::channel(None).0,
176 });
177 Ok(DiscoveryOwner {
178 core,
179 })
180 }
181}
182
183pub struct DiscoveryOwner {
185 core: Arc<Core>,
186}
187
188impl DiscoveryOwner {
189 pub async fn prepare(&self, deadline: Instant) -> Result<()> {
196 if Instant::now() >= deadline {
197 return Err(timeout());
198 }
199 {
200 let mut gate = lock(&self.core.gate);
201 match gate.phase {
202 OwnerPhase::Prepared => {
203 gate.phase = OwnerPhase::Starting;
204 for connection in self.core.connections.values() {
205 let (core, connection) = (self.core.clone(), connection.clone());
206 spawn_owned(
207 &self.core,
208 async move { prepare_loop(core, connection).await },
209 );
210 }
211 }
212 OwnerPhase::Starting => {}
213 _ => return Err(invalid("owner cannot prepare in its current phase")),
214 }
215 }
216 let mut changes = self.core.changes.subscribe();
217 loop {
218 if self.core.stop.is_cancelled() {
219 return Err(closed());
220 }
221 if Instant::now() >= deadline {
222 return Err(timeout());
223 }
224 let ready = self
225 .core
226 .connections
227 .values()
228 .filter(|connection| connection.policy.startup_required)
229 .all(|connection| *connection.prepared.borrow());
230 if ready {
231 return Ok(());
232 }
233 tokio::time::timeout_at(deadline, changes.changed())
234 .await
235 .map_err(|_| timeout())?
236 .map_err(|_| closed())?;
237 }
238 }
239
240 pub fn client(&self, qualifier: Option<&str>) -> Result<DiscoveryHandle> {
247 let gate = lock(&self.core.gate);
248 if gate.close_started || gate.failed {
249 return Err(closed());
250 }
251 let name = qualifier
252 .or(self.core.default_connection.as_deref())
253 .ok_or_else(|| invalid("no default discovery connection"))?;
254 let connection = self
255 .core
256 .connections
257 .get(name)
258 .cloned()
259 .ok_or_else(|| invalid("unknown discovery connection"))?;
260 Ok(DiscoveryHandle {
261 core: Arc::downgrade(&self.core),
262 connection,
263 })
264 }
265
266 pub fn freeze_registrations(&self, specs: Vec<RegistrationSpec>) -> Result<()> {
273 let mut gate = lock(&self.core.gate);
274 if gate.frozen || gate.phase != OwnerPhase::Starting {
275 return Err(invalid(
276 "registration plans are already frozen or owner is not starting",
277 ));
278 }
279 if specs.len() > self.core.limits.max_registrations {
280 return Err(invalid("registration capacity exceeded"));
281 }
282 let mut records = BTreeMap::new();
283 let mut keys = BTreeSet::new();
284 for mut spec in specs {
285 validate_name(&spec.plan.id)?;
286 let connection = self
287 .core
288 .connections
289 .get(&spec.connection_ref)
290 .cloned()
291 .ok_or_else(|| invalid("unknown registration connection"))?;
292 spec.policy.validate(&connection.policy)?;
293 if !connection.provider.capabilities().registration {
294 return Err(invalid("registration provider is read only"));
295 }
296 if spec
297 .plan
298 .instance
299 .metadata
300 .contains_key("nasa.discovery.owner")
301 {
302 return Err(invalid("registration ownership metadata is reserved"));
303 }
304 spec.plan
305 .instance
306 .metadata
307 .insert("nasa.discovery.owner".into(), self.core.id.clone());
308 let instance = &spec.plan.instance;
309 if instance.ip.trim() != instance.ip
310 || instance.ip.is_empty()
311 || instance.port == 0
312 || instance.ip == "0.0.0.0"
313 || instance.ip == "::"
314 || instance.ip == "[::]"
315 || !instance.weight.is_finite()
316 || instance.weight < 0.0
317 || spec.plan.service.is_empty()
318 {
319 return Err(invalid("registration endpoint or weight is invalid"));
320 }
321 connection.provider.validate_registration(&spec.plan)?;
322 if !keys.insert(connection.provider.registration_key(&spec.plan))
323 || records.contains_key(&spec.plan.id)
324 {
325 return Err(invalid(
326 "duplicate external registration key or local registration name",
327 ));
328 }
329 records.insert(
330 spec.plan.id.clone(),
331 Arc::new(Record::new(spec, connection)),
332 );
333 }
334 gate.registrations = records;
335 gate.frozen = true;
336 Ok(())
337 }
338
339 pub async fn publish(&self, deadline: Instant) -> Result<()> {
346 if Instant::now() >= deadline {
347 return Err(timeout());
348 }
349 {
350 let gate = lock(&self.core.gate);
351 if !gate.frozen
352 || gate.phase != OwnerPhase::Starting
353 || gate.failed
354 || gate.close_started
355 {
356 return Err(invalid("owner is not ready to publish frozen plans"));
357 }
358 for record in gate.registrations.values() {
359 if !record
360 .started
361 .swap(true, std::sync::atomic::Ordering::AcqRel)
362 {
363 let (core, record) = (self.core.clone(), record.clone());
364 crate::spawn_registration(&self.core, async move {
365 crate::registration::run(core, record).await
366 });
367 }
368 }
369 }
370 let mut changes = self.core.changes.subscribe();
371 loop {
372 {
373 let mut gate = lock(&self.core.gate);
374 if gate.phase != OwnerPhase::Starting || gate.failed || gate.close_started {
375 return Err(closed());
376 }
377 if Instant::now() >= deadline {
379 return Err(timeout());
380 }
381 let ready = gate
382 .registrations
383 .values()
384 .filter(|record| record.spec.policy.startup_required)
385 .all(|record| record.current_status(&gate).is_ready());
386 if ready {
387 gate.phase = OwnerPhase::Serving;
388 return Ok(());
389 }
390 }
391 tokio::time::timeout_at(deadline, changes.changed())
392 .await
393 .map_err(|_| timeout())?
394 .map_err(|_| closed())?;
395 }
396 }
397
398 pub fn begin_quiesce(&self, deadline: Instant) -> QuiesceTicket {
405 begin_quiesce(self.core.clone(), deadline)
406 }
407
408 pub fn begin_close(&self, deadline: Instant) -> CloseTicket {
415 begin_close(self.core.clone(), deadline)
416 }
417
418 pub fn phase(&self) -> OwnerPhase {
424 lock(&self.core.gate).phase
425 }
426
427 pub fn registrations(&self) -> Vec<RegistrationStatus> {
433 let gate = lock(&self.core.gate);
434 gate.registrations
435 .values()
436 .map(|record| record.current_status(&gate))
437 .collect()
438 }
439
440 pub fn changes(&self) -> watch::Receiver<u64> {
446 self.core.changes.subscribe()
447 }
448
449 pub fn observer(&self) -> crate::RuntimeObserver {
455 crate::RuntimeObserver {
456 core: Arc::downgrade(&self.core),
457 }
458 }
459
460 pub fn service_key(&self, connection: &str, service: &str) -> Result<String> {
468 self.core
469 .connections
470 .get(connection)
471 .ok_or_else(|| invalid("unknown discovery connection"))?
472 .provider
473 .service_key(service)
474 }
475}
476
477impl Drop for DiscoveryOwner {
478 fn drop(&mut self) {
484 if !lock(&self.core.gate).close_started {
485 if tokio::runtime::Handle::try_current().is_ok() {
486 begin_close(self.core.clone(), Instant::now() + millis(3000));
487 } else {
488 let mut gate = lock(&self.core.gate);
489 gate.phase = OwnerPhase::Closed;
490 gate.close_started = true;
491 self.core.stop.cancel();
492 self.core.registrations_stop.cancel();
493 for task in lock(&self.core.tasks).drain(..) {
494 task.abort();
495 }
496 for task in lock(&self.core.registration_tasks).drain(..) {
497 task.abort();
498 }
499 }
500 }
501 }
502}
503
504#[derive(Clone, Debug)]
506pub struct CloseReport {
507 pub registrations: Vec<RegistrationStatus>,
509 pub provider_errors: BTreeMap<String, DiscoveryError>,
511 pub deadline_elapsed: bool,
513}
514
515#[derive(Clone)]
517pub struct CloseTicket {
518 receiver: watch::Receiver<Option<CloseReport>>,
519}
520impl CloseTicket {
521 pub async fn wait(mut self) -> Result<CloseReport> {
527 loop {
528 if let Some(report) = self.receiver.borrow().clone() {
529 return Ok(report);
530 }
531 self.receiver.changed().await.map_err(|_| closed())?;
532 }
533 }
534}
535
536#[derive(Clone)]
538pub struct QuiesceTicket {
539 receiver: watch::Receiver<Option<Vec<RegistrationStatus>>>,
540}
541impl QuiesceTicket {
542 pub async fn wait(mut self) -> Result<Vec<RegistrationStatus>> {
548 loop {
549 if let Some(statuses) = self.receiver.borrow().clone() {
550 return Ok(statuses);
551 }
552 self.receiver.changed().await.map_err(|_| closed())?;
553 }
554 }
555}
556
557fn begin_quiesce(core: Arc<Core>, deadline: Instant) -> QuiesceTicket {
565 let mut gate = lock(&core.gate);
566 let ticket = QuiesceTicket {
567 receiver: core.withdrawn.subscribe(),
568 };
569 if gate.shutdown_deadline.is_none() {
570 gate.shutdown_deadline = Some(deadline);
571 gate.phase = OwnerPhase::Quiescing;
572 core.registrations_stop.cancel();
573 let shared = core.clone();
574 *lock(&core.withdrawal_task) = Some(tokio::spawn(async move {
575 use futures_util::FutureExt;
576 let _ = std::panic::AssertUnwindSafe(cleanup_registrations(shared.clone(), deadline))
577 .catch_unwind()
578 .await;
579 finish_withdrawal(&shared);
580 }));
581 }
582 ticket
583}
584
585fn begin_close(core: Arc<Core>, deadline: Instant) -> CloseTicket {
593 begin_quiesce(core.clone(), deadline);
594 let mut gate = lock(&core.gate);
595 let ticket = CloseTicket {
596 receiver: core.close.subscribe(),
597 };
598 if !gate.close_started {
599 gate.close_started = true;
601 gate.phase = OwnerPhase::Quiescing;
602 core.registrations_stop.cancel();
603 core.stop.cancel();
604 let deadline = gate.shutdown_deadline.unwrap_or(deadline).min(deadline);
605 let shared = core.clone();
606 *lock(&core.close_task) = Some(tokio::spawn(async move {
607 use futures_util::FutureExt;
608 if std::panic::AssertUnwindSafe(cleanup(shared.clone(), deadline))
609 .catch_unwind()
610 .await
611 .is_err()
612 {
613 finish_withdrawal(&shared);
614 lock(&shared.gate).phase = OwnerPhase::Closed;
615 shared.close.send_replace(Some(CloseReport {
616 registrations: lock(&shared.gate)
617 .registrations
618 .values()
619 .map(|r| lock(&r.status).clone())
620 .collect(),
621 provider_errors: shared
622 .connections
623 .keys()
624 .map(|name| {
625 (
626 name.clone(),
627 DiscoveryError::new(
628 nadisc::managed::ErrorCode::Unavailable,
629 "provider cleanup was interrupted",
630 ),
631 )
632 })
633 .collect(),
634 deadline_elapsed: Instant::now() >= deadline,
635 }));
636 }
637 }));
638 }
639 ticket
640}
641
642async fn prepare_loop(core: Arc<Core>, connection: Arc<Connection>) {
650 let mut backoff = connection.policy.initial_backoff_ms;
651 let mut attempt = 0u64;
652 loop {
653 let result = operation(
654 &core,
655 Instant::now() + millis(connection.policy.request_timeout_ms),
656 connection.provider.prepare(),
657 )
658 .await;
659 if core.stop.is_cancelled() {
660 break;
661 }
662 match result {
663 Ok(()) => {
664 connection.prepared.send_replace(true);
665 *lock(&connection.error) = None;
666 changed(&core);
667 return;
668 }
669 Err(error) => {
670 if matches!(
671 error.code,
672 ErrorCode::Unauthenticated
673 | ErrorCode::PermissionDenied
674 | ErrorCode::PermissionFiltered
675 ) {
676 let gate = lock(&core.gate);
677 crate::revoke_connection(&core, &gate, &connection);
678 }
679 *lock(&connection.error) = Some(error);
680 changed(&core);
681 }
682 }
683 tokio::select! { _ = core.stop.cancelled() => break, _ = tokio::time::sleep(crate::retry_delay(&core, &connection, backoff, attempt)) => {} }
684 attempt = attempt.wrapping_add(1);
685 backoff = backoff
686 .saturating_mul(2)
687 .min(connection.policy.max_backoff_ms);
688 }
689}
690
691async fn cleanup_registrations(core: Arc<Core>, deadline: Instant) {
699 use futures_util::{stream, StreamExt};
700 let tasks = std::mem::take(&mut *lock(&core.registration_tasks));
701 for task in &tasks {
702 task.abort();
703 }
704 for task in tasks {
705 let _ = tokio::time::timeout_at(deadline, task).await;
706 }
707 let records: Vec<_> = lock(&core.gate).registrations.values().cloned().collect();
708 let count = records.len().max(1);
709 let share = deadline.saturating_duration_since(Instant::now())
710 / count.div_ceil(core.limits.max_parallel_operations) as u32;
711 stream::iter(records.iter().cloned())
712 .for_each_concurrent(core.limits.max_parallel_operations, |record| async move {
713 use futures_util::FutureExt;
714 let _ = std::panic::AssertUnwindSafe(crate::registration::withdraw(
715 record, deadline, share,
716 ))
717 .catch_unwind()
718 .await;
719 })
720 .await;
721}
722
723fn finish_withdrawal(core: &Core) {
730 let gate = lock(&core.gate);
731 for record in gate.registrations.values() {
732 crate::registration::finish_withdrawal(record);
733 }
734 core.withdrawn.send_replace(Some(
735 gate.registrations
736 .values()
737 .map(|record| lock(&record.status).clone())
738 .collect(),
739 ));
740 changed(core);
741}
742
743async fn cleanup(core: Arc<Core>, deadline: Instant) {
751 use futures_util::{stream, StreamExt};
752 let withdrawal = lock(&core.withdrawal_task).take();
753 if let Some(mut task) = withdrawal {
754 if tokio::time::timeout_at(deadline, &mut task).await.is_err() {
755 task.abort();
756 let _ = task.await;
757 }
758 }
759 finish_withdrawal(&core);
761 let tasks = std::mem::take(&mut *lock(&core.tasks));
762 for task in &tasks {
763 task.abort();
764 }
765 for task in tasks {
766 let _ = tokio::time::timeout_at(deadline, task).await;
767 }
768 let records: Vec<_> = lock(&core.gate).registrations.values().cloned().collect();
769 let providers = stream::iter(core.connections.values().cloned())
770 .map(|connection| async move {
771 use futures_util::FutureExt;
772 let result = tokio::time::timeout_at(
773 deadline,
774 std::panic::AssertUnwindSafe(connection.provider.close()).catch_unwind(),
775 )
776 .await
777 .map(|result| {
778 result.unwrap_or_else(|_| {
779 Err(DiscoveryError::new(
780 nadisc::managed::ErrorCode::Unavailable,
781 "provider cleanup was interrupted",
782 ))
783 })
784 })
785 .unwrap_or_else(|_| Err(timeout()));
786 (connection.name.clone(), result)
787 })
788 .buffer_unordered(core.limits.max_parallel_operations)
789 .collect::<Vec<_>>()
790 .await;
791 let provider_errors = providers
792 .into_iter()
793 .filter_map(|(name, result)| result.err().map(|error| (name, error)))
794 .collect();
795 let registrations = records
796 .iter()
797 .map(|record| lock(&record.status).clone())
798 .collect();
799 lock(&core.gate).phase = OwnerPhase::Closed;
800 core.close.send_replace(Some(CloseReport {
802 registrations,
803 provider_errors,
804 deadline_elapsed: Instant::now() >= deadline,
805 }));
806 changed(&core);
807}