1use crate::{
4 closed, lock, millis, operation, spawn_owned, timeout, Connection, Core, OwnerPhase, Quota,
5};
6use nadisc::managed::{DiscoveryError, ErrorCode, EvidenceKind, Result};
7use nadisc::{
8 is_traffic_instance, DiscoveryClient, Instance, ServiceWatch, ServiceWatchGuard, WatchOptions,
9};
10use std::{
11 collections::BTreeSet,
12 sync::{atomic::Ordering, Arc, Mutex, Weak},
13};
14use tokio::{sync::watch, time::Instant};
15use tokio_util::sync::CancellationToken;
16
17pub(crate) struct Slot {
18 pub service: String,
19 pub state: Mutex<State>,
20 pub changes: watch::Sender<u64>,
21}
22
23pub(crate) struct State {
24 data: Option<Arc<Data>>,
25 pub valid_until: Option<Instant>,
26 pub observed_at: Option<Instant>,
27 pub error: Option<DiscoveryError>,
28 sequence: u64,
29 epoch: u64,
30 revocation: u64,
31 touched: Instant,
32}
33
34struct Data {
35 instances: Vec<Instance>,
36 identities: Vec<String>,
37 evidence: EvidenceKind,
38 bytes: usize,
39 quota: Arc<Mutex<Quota>>,
40}
41
42impl Drop for Data {
43 fn drop(&mut self) {
49 let mut quota = lock(&self.quota);
50 quota.instances -= self.instances.len();
51 quota.bytes -= self.bytes;
52 }
53}
54
55#[derive(Clone)]
57pub struct Snapshot {
58 data: Arc<Data>,
59 pub owner_id: String,
61 pub connection: String,
63 pub service: String,
65 pub epoch: u64,
67 pub sequence: u64,
69 pub observed_at: Instant,
71 pub valid_until: Option<Instant>,
73}
74
75impl Snapshot {
76 pub fn instances(&self) -> &[Instance] {
82 &self.data.instances
83 }
84
85 pub fn identities(&self) -> &[String] {
91 &self.data.identities
92 }
93
94 pub fn evidence(&self) -> EvidenceKind {
100 self.data.evidence
101 }
102}
103
104#[derive(Clone)]
106pub struct DiscoveryHandle {
107 pub(crate) core: Weak<Core>,
108 pub(crate) connection: Arc<Connection>,
109}
110
111pub struct Observation {
113 core: Weak<Core>,
114 connection: Arc<Connection>,
115 slot: Arc<Slot>,
116 changes: watch::Receiver<u64>,
117}
118
119impl Clone for Observation {
120 fn clone(&self) -> Self {
126 Self {
127 core: self.core.clone(),
128 connection: self.connection.clone(),
129 slot: self.slot.clone(),
130 changes: self.changes.clone(),
131 }
132 }
133}
134
135impl DiscoveryHandle {
136 pub async fn observe(&self, service: &str, deadline: Instant) -> Result<Observation> {
144 if Instant::now() >= deadline {
146 return Err(timeout());
147 }
148 let core = self.core.upgrade().ok_or_else(closed)?;
149 if service.is_empty() || service.len() > 512 || service.trim() != service {
150 return Err(crate::config::invalid("invalid discovery service"));
151 }
152 let mut observation = {
153 let gate = lock(&core.gate);
154 if gate.failed
155 || matches!(gate.phase, OwnerPhase::Prepared | OwnerPhase::Closed)
156 || gate.close_started
157 {
158 return Err(closed());
159 }
160 let mut slots = lock(&self.connection.slots);
161 if Instant::now() >= deadline {
163 return Err(timeout());
164 }
165 let slot = if let Some(slot) = slots.get(service) {
166 let mut state = lock(&slot.state);
167 if Instant::now() >= deadline {
168 return Err(timeout());
169 }
170 state.touched = Instant::now();
171 slot.clone()
172 } else {
173 if slots.len() >= core.limits.max_service_watches_per_connection {
174 return Err(DiscoveryError::new(
175 ErrorCode::CapacityExceeded,
176 "service observation capacity reached",
177 ));
178 }
179 let slot = Arc::new(Slot {
180 service: service.into(),
181 state: Mutex::new(State {
182 data: None,
183 valid_until: None,
184 observed_at: None,
185 error: None,
186 sequence: 0,
187 epoch: 0,
188 revocation: 0,
189 touched: Instant::now(),
190 }),
191 changes: watch::channel(0).0,
192 });
193 slots.insert(service.into(), slot.clone());
194 let (shared, connection, target) =
195 (core.clone(), self.connection.clone(), slot.clone());
196 spawn_owned(&core, async move {
197 observe_loop(shared, connection, target).await
198 });
199 slot
200 };
201 Observation {
202 core: self.core.clone(),
203 connection: self.connection.clone(),
204 changes: slot.changes.subscribe(),
205 slot,
206 }
207 };
208 loop {
209 if Instant::now() >= deadline {
211 return Err(timeout());
212 }
213 if lock(&observation.slot.state).sequence != 0 {
214 let snapshot = observation.snapshot(false);
215 if Instant::now() >= deadline {
217 return Err(timeout());
218 }
219 snapshot?;
220 return Ok(observation);
221 }
222 tokio::select! {
223 _ = core.stop.cancelled() => return Err(closed()),
224 result = tokio::time::timeout_at(deadline, observation.changes.changed()) => {
225 result.map_err(|_| timeout())?.map_err(|_| closed())?;
226 }
227 }
228 }
229 }
230
231 pub fn name(&self) -> &str {
237 &self.connection.name
238 }
239}
240
241impl Observation {
242 pub fn snapshot(&self, strict: bool) -> Result<Snapshot> {
249 let core = self.core.upgrade().ok_or_else(closed)?;
250 let gate = lock(&core.gate);
251 if gate.failed || gate.close_started || gate.phase == OwnerPhase::Closed {
252 return Err(closed());
253 }
254 let state = lock(&self.slot.state);
255 check(&self.connection, &state, strict)?;
256 Ok(Snapshot {
257 data: state.data.as_ref().ok_or_else(closed)?.clone(),
258 owner_id: core.id.clone(),
259 connection: self.connection.name.clone(),
260 service: self.slot.service.clone(),
261 epoch: state.epoch,
262 sequence: state.sequence,
263 observed_at: state.observed_at.ok_or_else(closed)?,
264 valid_until: state.valid_until,
265 })
266 }
267
268 pub async fn changed(&mut self, deadline: Instant) -> Result<()> {
275 let core = self.core.upgrade().ok_or_else(closed)?;
276 let expiry = lock(&self.slot.state).valid_until;
277 let wake_at = expiry.map_or(deadline, |expiry| expiry.min(deadline));
278 tokio::select! {
279 _ = core.stop.cancelled() => Err(closed()),
280 result = self.changes.changed() => result.map_err(|_| closed()),
281 _ = tokio::time::sleep_until(wake_at) => if expiry.is_some_and(|expiry| expiry <= deadline) { Ok(()) } else { Err(timeout()) },
283 }
284 }
285
286 pub fn submit_if_current<T>(
294 &self,
295 instance: &Instance,
296 submit: impl FnOnce() -> T,
297 ) -> Result<T> {
298 self.submit(instance, false, submit)
299 }
300
301 pub fn submit_if_current_strict<T>(
309 &self,
310 instance: &Instance,
311 submit: impl FnOnce() -> T,
312 ) -> Result<T> {
313 self.submit(instance, true, submit)
314 }
315
316 fn submit<T>(
325 &self,
326 instance: &Instance,
327 strict: bool,
328 submit: impl FnOnce() -> T,
329 ) -> Result<T> {
330 let core = self.core.upgrade().ok_or_else(closed)?;
331 let gate = lock(&core.gate);
332 if gate.failed || gate.close_started || gate.phase == OwnerPhase::Closed {
333 return Err(closed());
334 }
335 let state = lock(&self.slot.state);
336 check(&self.connection, &state, strict)?;
337 if !is_traffic_instance(instance)
338 || !state
339 .data
340 .as_ref()
341 .is_some_and(|data| data.instances.contains(instance))
342 {
343 return Err(DiscoveryError::new(
344 ErrorCode::NoInstances,
345 "candidate no longer has routing eligibility",
346 ));
347 }
348 Ok(submit())
349 }
350}
351
352pub(crate) fn check(connection: &Connection, state: &State, strict: bool) -> Result<()> {
361 if state.epoch != connection.epoch.load(Ordering::Acquire)
362 || state.revocation != connection.revocation.load(Ordering::Acquire)
363 {
364 return Err(closed());
365 }
366 if strict && state.error.is_some() {
367 return Err(state.error.clone().unwrap());
368 }
369 if state.data.is_none()
370 || state.observed_at.is_none()
371 || state
372 .valid_until
373 .is_some_and(|deadline| Instant::now() >= deadline)
374 || (state.valid_until.is_none()
375 && state
376 .data
377 .as_ref()
378 .is_none_or(|data| data.evidence != EvidenceKind::Static))
379 {
380 return Err(state.error.clone().unwrap_or_else(|| {
381 DiscoveryError::new(
382 ErrorCode::Unavailable,
383 "discovery evidence is unavailable or expired",
384 )
385 }));
386 }
387 Ok(())
388}
389
390fn record_failure(connection: &Connection, state: &mut State, error: &DiscoveryError) {
399 let grace = if error.revokes_evidence() {
400 0
401 } else {
402 connection
403 .policy
404 .stale_if_error_ms
405 .min(connection.policy.max_age_ms)
406 };
407 let failure_deadline = state
408 .observed_at
409 .map_or(Instant::now(), |at| at + millis(grace));
410 state.valid_until = Some(
412 state
413 .valid_until
414 .map_or(failure_deadline, |current| current.min(failure_deadline)),
415 );
416 state.error = Some(error.clone());
417}
418
419async fn observe_loop(core: Arc<Core>, connection: Arc<Connection>, slot: Arc<Slot>) {
428 let mut backoff = connection.policy.initial_backoff_ms;
429 loop {
430 let epoch = connection.epoch.load(Ordering::Acquire);
431 let revocation = connection.revocation.load(Ordering::Acquire);
432 let deadline = Instant::now() + millis(connection.policy.request_timeout_ms);
433 let started = Instant::now();
434 let result = operation(&core, deadline, connection.provider.observe(&slot.service)).await;
435 let failed = result.is_err();
436 connection.queries.fetch_add(1, Ordering::Relaxed);
437 connection.query_micros.fetch_add(
438 started.elapsed().as_micros().min(u64::MAX as u128) as u64,
439 Ordering::Relaxed,
440 );
441 if result.is_err() {
442 connection.query_errors.fetch_add(1, Ordering::Relaxed);
443 }
444 {
445 let gate = lock(&core.gate);
446 if gate.failed || gate.close_started || core.stop.is_cancelled() {
447 break;
448 }
449 if epoch == connection.epoch.load(Ordering::Acquire)
450 && revocation == connection.revocation.load(Ordering::Acquire)
451 {
452 let mut state = lock(&slot.state);
453 let result = result.and_then(|snapshot| {
454 if let Some(data) = state.data.as_ref().filter(|data| {
455 data.evidence == snapshot.evidence && data.instances == snapshot.instances
456 }) {
457 Ok(data.clone())
458 } else {
459 if state.data.as_ref().is_some_and(|data| {
461 data.instances
462 .iter()
463 .filter(|instance| is_traffic_instance(instance))
464 .any(|instance| !snapshot.instances.contains(instance))
465 }) {
466 state.valid_until = Some(Instant::now());
467 }
468 allocate(
469 &core,
470 &connection,
471 snapshot.instances,
472 snapshot.evidence,
473 &mut state.data,
474 )
475 }
476 });
477 state.sequence = state.sequence.wrapping_add(1);
478 state.epoch = epoch;
479 state.revocation = revocation;
480 match result {
481 Ok(data) => {
482 let now = Instant::now();
483 state.valid_until = (data.evidence != EvidenceKind::Static)
484 .then(|| now + millis(connection.policy.max_age_ms));
485 state.data = Some(data);
486 state.observed_at = Some(now);
487 state.error = None;
488 }
489 Err(error) => {
490 record_failure(&connection, &mut state, &error);
491 if matches!(
492 error.code,
493 ErrorCode::Unauthenticated
494 | ErrorCode::PermissionDenied
495 | ErrorCode::PermissionFiltered
496 ) {
497 drop(state);
498 crate::revoke_connection(&core, &gate, &connection);
499 }
500 }
501 }
502 slot.changes
503 .send_modify(|value| *value = value.wrapping_add(1));
504 }
505 }
506 let delay = if failed {
507 let delay = crate::retry_delay(
508 &core,
509 &connection,
510 backoff,
511 connection.queries.load(Ordering::Relaxed),
512 );
513 backoff = backoff
514 .saturating_mul(2)
515 .min(connection.policy.max_backoff_ms);
516 delay
517 } else {
518 backoff = connection.policy.initial_backoff_ms;
519 millis(connection.policy.poll_interval_ms)
520 };
521 let wait_for_change = async {
522 if failed {
523 return std::future::pending().await;
524 }
525 if let Some(timeout) = connection.provider.change_wait_timeout() {
526 let streams = connection.provider.change_wait_streams().max(1);
528 let _permit = core
529 .watches
530 .acquire_many(streams)
531 .await
532 .map_err(|_| closed())?;
533 tokio::time::timeout(timeout, connection.provider.wait_for_change(&slot.service))
534 .await
535 .unwrap_or_else(|_| Err(crate::timeout()))
536 } else {
537 std::future::pending().await
538 }
539 };
540 tokio::select! {
541 _ = core.stop.cancelled() => break,
542 _ = tokio::time::sleep(delay) => {},
543 result = wait_for_change => {
544 if let Err(error) = result {
545 {
546 let gate = lock(&core.gate);
547 let mut state = lock(&slot.state);
548 record_failure(&connection, &mut state, &error);
549 drop(state);
550 if matches!(error.code, ErrorCode::Unauthenticated | ErrorCode::PermissionDenied | ErrorCode::PermissionFiltered) {
551 crate::revoke_connection(&core, &gate, &connection);
552 }
553 slot.changes.send_modify(|value| *value = value.wrapping_add(1));
554 }
555 tokio::select! {
557 _ = core.stop.cancelled() => break,
558 _ = tokio::time::sleep(millis(connection.policy.initial_backoff_ms)) => {},
559 }
560 } else {
561 tokio::select! {
563 _ = core.stop.cancelled() => break,
564 _ = tokio::time::sleep(millis(connection.policy.poll_interval_ms.min(100))) => {},
565 }
566 }
567 }
568 }
569 let _gate = lock(&core.gate);
570 let mut slots = lock(&connection.slots);
571 if slot.changes.receiver_count() == 0
572 && lock(&slot.state).touched.elapsed() >= millis(connection.policy.idle_timeout_ms)
573 {
574 slots.remove(&slot.service);
575 break;
576 }
577 }
578}
579
580fn allocate(
591 core: &Core,
592 connection: &Connection,
593 instances: Vec<Instance>,
594 evidence: EvidenceKind,
595 previous: &mut Option<Arc<Data>>,
596) -> Result<Arc<Data>> {
597 let mut keys = BTreeSet::new();
598 let mut identities = Vec::with_capacity(instances.len());
599 let mut bytes = instances
600 .capacity()
601 .saturating_mul(std::mem::size_of::<Instance>());
602 for instance in &instances {
603 let identity = connection.provider.instance_key(instance);
604 if instance.ip.is_empty()
605 || instance.port == 0
606 || !instance.weight.is_finite()
607 || identity.is_empty()
608 || identity.len() > 2048
609 || !keys.insert(identity.clone())
610 {
611 return Err(DiscoveryError::new(
612 ErrorCode::InvalidSnapshot,
613 "invalid or duplicate instance in complete observation",
614 ));
615 }
616 bytes = bytes.saturating_add(identity.capacity());
617 identities.push(identity);
618 bytes = bytes
619 .saturating_add(instance.ip.capacity())
620 .saturating_add(instance.cluster_name.as_ref().map_or(0, String::capacity))
621 .saturating_add(
622 instance
623 .metadata
624 .capacity()
625 .saturating_mul(std::mem::size_of::<(String, String)>() + 16),
626 );
627 for (key, value) in &instance.metadata {
628 bytes = bytes
629 .saturating_add(key.capacity())
630 .saturating_add(value.capacity());
631 }
632 }
633 bytes = bytes.saturating_add(
634 identities
635 .capacity()
636 .saturating_mul(std::mem::size_of::<String>()),
637 );
638 let mut quota = lock(&core.quota);
639 let reusable = previous.as_mut().and_then(Arc::get_mut);
641 let old_instances = reusable.as_ref().map_or(0, |data| data.instances.len());
642 let old_bytes = reusable.as_ref().map_or(0, |data| data.bytes);
643 let next_instances = (quota.instances - old_instances).saturating_add(instances.len());
644 let next_bytes = (quota.bytes - old_bytes).saturating_add(bytes);
645 if instances.len() > core.limits.max_instances_per_service
646 || next_instances > core.limits.max_total_instances
647 || next_bytes > core.limits.max_snapshot_bytes
648 {
649 return Err(DiscoveryError::new(
650 ErrorCode::CapacityExceeded,
651 "snapshot capacity reached",
652 ));
653 }
654 quota.instances = next_instances;
656 quota.bytes = next_bytes;
657 if let Some(data) = reusable {
658 data.instances = instances;
659 data.identities = identities;
660 data.evidence = evidence;
661 data.bytes = bytes;
662 return Ok(previous.as_ref().unwrap().clone());
663 }
664 Ok(Arc::new(Data {
665 instances,
666 identities,
667 evidence,
668 bytes,
669 quota: core.quota.clone(),
670 }))
671}
672
673#[async_trait::async_trait]
674impl DiscoveryClient for DiscoveryHandle {
675 fn requires_admission(&self) -> bool {
681 true
682 }
683
684 fn submit_if_current(
693 &self,
694 service: &str,
695 instance: &Instance,
696 submit: &mut dyn FnMut(),
697 ) -> anyhow::Result<()> {
698 let slot = lock(&self.connection.slots)
699 .get(service)
700 .cloned()
701 .ok_or_else(closed)?;
702 let observation = Observation {
703 core: self.core.clone(),
704 connection: self.connection.clone(),
705 changes: slot.changes.subscribe(),
706 slot,
707 };
708 Ok(observation.submit_if_current(instance, submit)?)
709 }
710
711 async fn list_services(&self) -> anyhow::Result<Vec<String>> {
717 let core = self.core.upgrade().ok_or_else(closed)?;
718 Ok(operation(
719 &core,
720 Instant::now() + millis(self.connection.policy.request_timeout_ms),
721 self.connection.provider.services(),
722 )
723 .await?)
724 }
725
726 async fn discover(&self, service: &str) -> anyhow::Result<Vec<Instance>> {
733 Ok(self
734 .discover_all(service)
735 .await?
736 .into_iter()
737 .filter(is_traffic_instance)
738 .collect())
739 }
740
741 async fn discover_all(&self, service: &str) -> anyhow::Result<Vec<Instance>> {
748 let observation = self
749 .observe(
750 service,
751 Instant::now() + millis(self.connection.policy.request_timeout_ms),
752 )
753 .await?;
754 Ok(observation.snapshot(false)?.instances().to_vec())
755 }
756
757 async fn watch_with_options(
765 &self,
766 service: &str,
767 options: WatchOptions,
768 ) -> anyhow::Result<ServiceWatch> {
769 options.validate()?;
770 if options.poll_interval < millis(self.connection.policy.poll_interval_ms) {
771 return Err(crate::config::invalid(
772 "watch interval is faster than frozen observation policy",
773 )
774 .into());
775 }
776 let core = self.core.upgrade().ok_or_else(closed)?;
777 let bridge = core.bridges.clone().try_acquire_owned().map_err(|_| {
778 DiscoveryError::new(
779 ErrorCode::CapacityExceeded,
780 "compatibility watch capacity reached",
781 )
782 })?;
783 let mut observation = self
784 .observe(
785 service,
786 Instant::now() + millis(self.connection.policy.request_timeout_ms),
787 )
788 .await?;
789 let initial = observation
790 .snapshot(false)?
791 .instances()
792 .iter()
793 .filter(|i| is_traffic_instance(i))
794 .cloned()
795 .collect();
796 let (sender, receiver) = watch::channel(initial);
797 let cancel = CancellationToken::new();
798 let stopped = cancel.clone();
799 spawn_owned(&core, async move {
800 let _bridge = bridge;
801 loop {
802 tokio::select! {
803 _ = stopped.cancelled() => break,
804 result = observation.changed(Instant::now() + millis(86_400_000)) => if result.is_err() { break; },
805 }
806 let Ok(snapshot) = observation.snapshot(false) else {
807 break;
808 };
809 if sender
810 .send(
811 snapshot
812 .instances()
813 .iter()
814 .filter(|i| is_traffic_instance(i))
815 .cloned()
816 .collect(),
817 )
818 .is_err()
819 {
820 break;
821 }
822 }
823 });
824 Ok(ServiceWatch {
825 guard: Box::new(BridgeGuard(cancel)),
826 receiver,
827 })
828 }
829}
830
831struct BridgeGuard(CancellationToken);
832impl Drop for BridgeGuard {
833 fn drop(&mut self) {
839 self.0.cancel();
840 }
841}
842#[async_trait::async_trait]
843impl ServiceWatchGuard for BridgeGuard {
844 async fn unsubscribe(self: Box<Self>) -> anyhow::Result<()> {
850 self.0.cancel();
851 Ok(())
852 }
853}