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 {
58 Self {
59 limits: Limits::default(),
60 connections: BTreeMap::new(),
61 default_connection: None,
62 }
63 }
64}
65
66impl OwnerBuilder {
67 pub fn limits(mut self, limits: Limits) -> Self {
71 self.limits = limits;
72 self
73 }
74
75 pub fn default_connection(mut self, name: impl Into<String>) -> Self {
79 self.default_connection = Some(name.into());
80 self
81 }
82
83 pub fn connection(mut self, spec: ConnectionSpec) -> Result<Self> {
87 validate_name(&spec.name)?;
88 spec.policy.validate()?;
89 if self.connections.contains_key(&spec.name) {
90 return Err(invalid("duplicate connection name"));
91 }
92 self.connections.insert(spec.name.clone(), spec);
93 Ok(self)
94 }
95
96 pub fn build(self) -> Result<DiscoveryOwner> {
100 self.limits.validate()?;
101 if self.connections.len() > self.limits.max_connections {
102 return Err(invalid("connection capacity exceeded"));
103 }
104 let default_connection = match self.default_connection {
105 Some(name) if self.connections.contains_key(&name) => Some(name),
106 Some(_) => return Err(invalid("unknown default connection")),
107 None => self
108 .connections
109 .contains_key("primary")
110 .then(|| "primary".into()),
111 };
112 let connections = self
113 .connections
114 .into_iter()
115 .map(|(name, spec)| {
116 let connection = Arc::new(Connection {
117 name: name.clone(),
118 provider: spec.provider,
119 policy: spec.policy,
120 slots: Mutex::new(BTreeMap::new()),
121 prepared: watch::channel(false).0,
122 error: Mutex::new(None),
123 epoch: AtomicU64::new(1),
124 revocation: AtomicU64::new(0),
125 queries: AtomicU64::new(0),
126 query_errors: AtomicU64::new(0),
127 query_micros: AtomicU64::new(0),
128 });
129 (name, connection)
130 })
131 .collect();
132 let core = Arc::new(Core {
133 id: uuid::Uuid::new_v4().to_string(),
134 gate: Mutex::new(Gate {
135 phase: OwnerPhase::Prepared,
136 registrations: BTreeMap::new(),
137 frozen: false,
138 close_started: false,
139 failed: false,
140 shutdown_deadline: None,
141 }),
142 connections,
143 default_connection,
144 operations: Semaphore::new(self.limits.max_parallel_operations),
145 watches: Semaphore::new(self.limits.max_watch_streams),
146 quota: Arc::new(Mutex::new(Quota::default())),
147 bridges: Arc::new(Semaphore::new(
148 self.limits
149 .max_service_watches_per_connection
150 .saturating_mul(self.limits.max_connections),
151 )),
152 transports: Arc::new(Semaphore::new(self.limits.max_transport_tasks)),
153 limits: self.limits,
154 tasks: Mutex::new(Vec::new()),
155 registration_tasks: Mutex::new(Vec::new()),
156 withdrawal_task: Mutex::new(None),
157 close_task: Mutex::new(None),
158 stop: CancellationToken::new(),
159 registrations_stop: CancellationToken::new(),
160 changes: watch::channel(0).0,
161 close: watch::channel(None).0,
162 withdrawn: watch::channel(None).0,
163 });
164 Ok(DiscoveryOwner { core })
165 }
166}
167
168pub struct DiscoveryOwner {
170 core: Arc<Core>,
171}
172
173impl DiscoveryOwner {
174 pub async fn prepare(&self, deadline: Instant) -> Result<()> {
178 if Instant::now() >= deadline {
179 return Err(timeout());
180 }
181 {
182 let mut gate = lock(&self.core.gate);
183 match gate.phase {
184 OwnerPhase::Prepared => {
185 gate.phase = OwnerPhase::Starting;
186 for connection in self.core.connections.values() {
187 let (core, connection) = (self.core.clone(), connection.clone());
188 spawn_owned(
189 &self.core,
190 async move { prepare_loop(core, connection).await },
191 );
192 }
193 }
194 OwnerPhase::Starting => {}
195 _ => return Err(invalid("owner cannot prepare in its current phase")),
196 }
197 }
198 let mut changes = self.core.changes.subscribe();
199 loop {
200 if self.core.stop.is_cancelled() {
201 return Err(closed());
202 }
203 if Instant::now() >= deadline {
204 return Err(timeout());
205 }
206 let ready = self
207 .core
208 .connections
209 .values()
210 .filter(|connection| connection.policy.startup_required)
211 .all(|connection| *connection.prepared.borrow());
212 if ready {
213 return Ok(());
214 }
215 tokio::time::timeout_at(deadline, changes.changed())
216 .await
217 .map_err(|_| timeout())?
218 .map_err(|_| closed())?;
219 }
220 }
221
222 pub fn client(&self, qualifier: Option<&str>) -> Result<DiscoveryHandle> {
226 let gate = lock(&self.core.gate);
227 if gate.close_started || gate.failed {
228 return Err(closed());
229 }
230 let name = qualifier
231 .or(self.core.default_connection.as_deref())
232 .ok_or_else(|| invalid("no default discovery connection"))?;
233 let connection = self
234 .core
235 .connections
236 .get(name)
237 .cloned()
238 .ok_or_else(|| invalid("unknown discovery connection"))?;
239 Ok(DiscoveryHandle {
240 core: Arc::downgrade(&self.core),
241 connection,
242 })
243 }
244
245 pub fn freeze_registrations(&self, specs: Vec<RegistrationSpec>) -> Result<()> {
249 let mut gate = lock(&self.core.gate);
250 if gate.frozen || gate.phase != OwnerPhase::Starting {
251 return Err(invalid(
252 "registration plans are already frozen or owner is not starting",
253 ));
254 }
255 if specs.len() > self.core.limits.max_registrations {
256 return Err(invalid("registration capacity exceeded"));
257 }
258 let mut records = BTreeMap::new();
259 let mut keys = BTreeSet::new();
260 for mut spec in specs {
261 validate_name(&spec.plan.id)?;
262 let connection = self
263 .core
264 .connections
265 .get(&spec.connection_ref)
266 .cloned()
267 .ok_or_else(|| invalid("unknown registration connection"))?;
268 spec.policy.validate(&connection.policy)?;
269 if !connection.provider.capabilities().registration {
270 return Err(invalid("registration provider is read only"));
271 }
272 if spec
273 .plan
274 .instance
275 .metadata
276 .contains_key("nasa.discovery.owner")
277 {
278 return Err(invalid("registration ownership metadata is reserved"));
279 }
280 spec.plan
281 .instance
282 .metadata
283 .insert("nasa.discovery.owner".into(), self.core.id.clone());
284 let instance = &spec.plan.instance;
285 if instance.ip.trim() != instance.ip
286 || instance.ip.is_empty()
287 || instance.port == 0
288 || instance.ip == "0.0.0.0"
289 || instance.ip == "::"
290 || instance.ip == "[::]"
291 || !instance.weight.is_finite()
292 || instance.weight < 0.0
293 || spec.plan.service.is_empty()
294 {
295 return Err(invalid("registration endpoint or weight is invalid"));
296 }
297 connection.provider.validate_registration(&spec.plan)?;
298 if !keys.insert(connection.provider.registration_key(&spec.plan))
299 || records.contains_key(&spec.plan.id)
300 {
301 return Err(invalid(
302 "duplicate external registration key or local registration name",
303 ));
304 }
305 records.insert(
306 spec.plan.id.clone(),
307 Arc::new(Record::new(spec, connection)),
308 );
309 }
310 gate.registrations = records;
311 gate.frozen = true;
312 Ok(())
313 }
314
315 pub async fn publish(&self, deadline: Instant) -> Result<()> {
319 if Instant::now() >= deadline {
320 return Err(timeout());
321 }
322 {
323 let gate = lock(&self.core.gate);
324 if !gate.frozen
325 || gate.phase != OwnerPhase::Starting
326 || gate.failed
327 || gate.close_started
328 {
329 return Err(invalid("owner is not ready to publish frozen plans"));
330 }
331 for record in gate.registrations.values() {
332 if !record
333 .started
334 .swap(true, std::sync::atomic::Ordering::AcqRel)
335 {
336 let (core, record) = (self.core.clone(), record.clone());
337 crate::spawn_registration(&self.core, async move {
338 crate::registration::run(core, record).await
339 });
340 }
341 }
342 }
343 let mut changes = self.core.changes.subscribe();
344 loop {
345 {
346 let mut gate = lock(&self.core.gate);
347 if gate.phase != OwnerPhase::Starting || gate.failed || gate.close_started {
348 return Err(closed());
349 }
350 if Instant::now() >= deadline {
352 return Err(timeout());
353 }
354 let ready = gate
355 .registrations
356 .values()
357 .filter(|record| record.spec.policy.startup_required)
358 .all(|record| record.current_status(&gate).is_ready());
359 if ready {
360 gate.phase = OwnerPhase::Serving;
361 return Ok(());
362 }
363 }
364 tokio::time::timeout_at(deadline, changes.changed())
365 .await
366 .map_err(|_| timeout())?
367 .map_err(|_| closed())?;
368 }
369 }
370
371 pub fn begin_quiesce(&self, deadline: Instant) -> QuiesceTicket {
375 begin_quiesce(self.core.clone(), deadline)
376 }
377
378 pub fn begin_close(&self, deadline: Instant) -> CloseTicket {
382 begin_close(self.core.clone(), deadline)
383 }
384
385 pub fn phase(&self) -> OwnerPhase {
389 lock(&self.core.gate).phase
390 }
391
392 pub fn registrations(&self) -> Vec<RegistrationStatus> {
396 let gate = lock(&self.core.gate);
397 gate.registrations
398 .values()
399 .map(|record| record.current_status(&gate))
400 .collect()
401 }
402
403 pub fn changes(&self) -> watch::Receiver<u64> {
407 self.core.changes.subscribe()
408 }
409
410 pub fn observer(&self) -> crate::RuntimeObserver {
414 crate::RuntimeObserver {
415 core: Arc::downgrade(&self.core),
416 }
417 }
418
419 pub fn service_key(&self, connection: &str, service: &str) -> Result<String> {
423 self.core
424 .connections
425 .get(connection)
426 .ok_or_else(|| invalid("unknown discovery connection"))?
427 .provider
428 .service_key(service)
429 }
430}
431
432impl Drop for DiscoveryOwner {
433 fn drop(&mut self) {
437 if !lock(&self.core.gate).close_started {
438 if tokio::runtime::Handle::try_current().is_ok() {
439 begin_close(self.core.clone(), Instant::now() + millis(3000));
440 } else {
441 let mut gate = lock(&self.core.gate);
442 gate.phase = OwnerPhase::Closed;
443 gate.close_started = true;
444 self.core.stop.cancel();
445 self.core.registrations_stop.cancel();
446 for task in lock(&self.core.tasks).drain(..) {
447 task.abort();
448 }
449 for task in lock(&self.core.registration_tasks).drain(..) {
450 task.abort();
451 }
452 }
453 }
454 }
455}
456
457#[derive(Clone, Debug)]
459pub struct CloseReport {
460 pub registrations: Vec<RegistrationStatus>,
462 pub provider_errors: BTreeMap<String, DiscoveryError>,
464 pub deadline_elapsed: bool,
466}
467
468#[derive(Clone)]
470pub struct CloseTicket {
471 receiver: watch::Receiver<Option<CloseReport>>,
472}
473impl CloseTicket {
474 pub async fn wait(mut self) -> Result<CloseReport> {
478 loop {
479 if let Some(report) = self.receiver.borrow().clone() {
480 return Ok(report);
481 }
482 self.receiver.changed().await.map_err(|_| closed())?;
483 }
484 }
485}
486
487#[derive(Clone)]
489pub struct QuiesceTicket {
490 receiver: watch::Receiver<Option<Vec<RegistrationStatus>>>,
491}
492impl QuiesceTicket {
493 pub async fn wait(mut self) -> Result<Vec<RegistrationStatus>> {
497 loop {
498 if let Some(statuses) = self.receiver.borrow().clone() {
499 return Ok(statuses);
500 }
501 self.receiver.changed().await.map_err(|_| closed())?;
502 }
503 }
504}
505
506fn begin_quiesce(core: Arc<Core>, deadline: Instant) -> QuiesceTicket {
510 let mut gate = lock(&core.gate);
511 let ticket = QuiesceTicket {
512 receiver: core.withdrawn.subscribe(),
513 };
514 if gate.shutdown_deadline.is_none() {
515 gate.shutdown_deadline = Some(deadline);
516 gate.phase = OwnerPhase::Quiescing;
517 core.registrations_stop.cancel();
518 let shared = core.clone();
519 *lock(&core.withdrawal_task) = Some(tokio::spawn(async move {
520 use futures_util::FutureExt;
521 let _ = std::panic::AssertUnwindSafe(cleanup_registrations(shared.clone(), deadline))
522 .catch_unwind()
523 .await;
524 finish_withdrawal(&shared);
525 }));
526 }
527 ticket
528}
529
530fn begin_close(core: Arc<Core>, deadline: Instant) -> CloseTicket {
534 begin_quiesce(core.clone(), deadline);
535 let mut gate = lock(&core.gate);
536 let ticket = CloseTicket {
537 receiver: core.close.subscribe(),
538 };
539 if !gate.close_started {
540 gate.close_started = true;
542 gate.phase = OwnerPhase::Quiescing;
543 core.registrations_stop.cancel();
544 core.stop.cancel();
545 let deadline = gate.shutdown_deadline.unwrap_or(deadline).min(deadline);
546 let shared = core.clone();
547 *lock(&core.close_task) = Some(tokio::spawn(async move {
548 use futures_util::FutureExt;
549 if std::panic::AssertUnwindSafe(cleanup(shared.clone(), deadline))
550 .catch_unwind()
551 .await
552 .is_err()
553 {
554 finish_withdrawal(&shared);
555 lock(&shared.gate).phase = OwnerPhase::Closed;
556 shared.close.send_replace(Some(CloseReport {
557 registrations: lock(&shared.gate)
558 .registrations
559 .values()
560 .map(|r| lock(&r.status).clone())
561 .collect(),
562 provider_errors: shared
563 .connections
564 .keys()
565 .map(|name| {
566 (
567 name.clone(),
568 DiscoveryError::new(
569 nadisc::managed::ErrorCode::Unavailable,
570 "provider cleanup was interrupted",
571 ),
572 )
573 })
574 .collect(),
575 deadline_elapsed: Instant::now() >= deadline,
576 }));
577 }
578 }));
579 }
580 ticket
581}
582
583async fn prepare_loop(core: Arc<Core>, connection: Arc<Connection>) {
587 let mut backoff = connection.policy.initial_backoff_ms;
588 let mut attempt = 0u64;
589 loop {
590 let result = operation(
591 &core,
592 Instant::now() + millis(connection.policy.request_timeout_ms),
593 connection.provider.prepare(),
594 )
595 .await;
596 if core.stop.is_cancelled() {
597 break;
598 }
599 match result {
600 Ok(()) => {
601 connection.prepared.send_replace(true);
602 *lock(&connection.error) = None;
603 changed(&core);
604 return;
605 }
606 Err(error) => {
607 if matches!(
608 error.code,
609 ErrorCode::Unauthenticated
610 | ErrorCode::PermissionDenied
611 | ErrorCode::PermissionFiltered
612 ) {
613 let gate = lock(&core.gate);
614 crate::revoke_connection(&core, &gate, &connection);
615 }
616 *lock(&connection.error) = Some(error);
617 changed(&core);
618 }
619 }
620 tokio::select! { _ = core.stop.cancelled() => break, _ = tokio::time::sleep(crate::retry_delay(&core, &connection, backoff, attempt)) => {} }
621 attempt = attempt.wrapping_add(1);
622 backoff = backoff
623 .saturating_mul(2)
624 .min(connection.policy.max_backoff_ms);
625 }
626}
627
628async fn cleanup_registrations(core: Arc<Core>, deadline: Instant) {
632 use futures_util::{stream, StreamExt};
633 let tasks = std::mem::take(&mut *lock(&core.registration_tasks));
634 for task in &tasks {
635 task.abort();
636 }
637 for task in tasks {
638 let _ = tokio::time::timeout_at(deadline, task).await;
639 }
640 let records: Vec<_> = lock(&core.gate).registrations.values().cloned().collect();
641 let count = records.len().max(1);
642 let share = deadline.saturating_duration_since(Instant::now())
643 / count.div_ceil(core.limits.max_parallel_operations) as u32;
644 stream::iter(records.iter().cloned())
645 .for_each_concurrent(core.limits.max_parallel_operations, |record| async move {
646 use futures_util::FutureExt;
647 let _ = std::panic::AssertUnwindSafe(crate::registration::withdraw(
648 record, deadline, share,
649 ))
650 .catch_unwind()
651 .await;
652 })
653 .await;
654}
655
656fn finish_withdrawal(core: &Core) {
660 let gate = lock(&core.gate);
661 for record in gate.registrations.values() {
662 crate::registration::finish_withdrawal(record);
663 }
664 core.withdrawn.send_replace(Some(
665 gate.registrations
666 .values()
667 .map(|record| lock(&record.status).clone())
668 .collect(),
669 ));
670 changed(core);
671}
672
673async fn cleanup(core: Arc<Core>, deadline: Instant) {
677 use futures_util::{stream, StreamExt};
678 let withdrawal = lock(&core.withdrawal_task).take();
679 if let Some(mut task) = withdrawal {
680 if tokio::time::timeout_at(deadline, &mut task).await.is_err() {
681 task.abort();
682 let _ = task.await;
683 }
684 }
685 finish_withdrawal(&core);
687 let tasks = std::mem::take(&mut *lock(&core.tasks));
688 for task in &tasks {
689 task.abort();
690 }
691 for task in tasks {
692 let _ = tokio::time::timeout_at(deadline, task).await;
693 }
694 let records: Vec<_> = lock(&core.gate).registrations.values().cloned().collect();
695 let providers = stream::iter(core.connections.values().cloned())
696 .map(|connection| async move {
697 use futures_util::FutureExt;
698 let result = tokio::time::timeout_at(
699 deadline,
700 std::panic::AssertUnwindSafe(connection.provider.close()).catch_unwind(),
701 )
702 .await
703 .map(|result| {
704 result.unwrap_or_else(|_| {
705 Err(DiscoveryError::new(
706 nadisc::managed::ErrorCode::Unavailable,
707 "provider cleanup was interrupted",
708 ))
709 })
710 })
711 .unwrap_or_else(|_| Err(timeout()));
712 (connection.name.clone(), result)
713 })
714 .buffer_unordered(core.limits.max_parallel_operations)
715 .collect::<Vec<_>>()
716 .await;
717 let provider_errors = providers
718 .into_iter()
719 .filter_map(|(name, result)| result.err().map(|error| (name, error)))
720 .collect();
721 let registrations = records
722 .iter()
723 .map(|record| lock(&record.status).clone())
724 .collect();
725 lock(&core.gate).phase = OwnerPhase::Closed;
726 core.close.send_replace(Some(CloseReport {
728 registrations,
729 provider_errors,
730 deadline_elapsed: Instant::now() >= deadline,
731 }));
732 changed(&core);
733}