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) {
47 let mut quota = lock(&self.quota);
48 quota.instances -= self.instances.len();
49 quota.bytes -= self.bytes;
50 }
51}
52
53#[derive(Clone)]
55pub struct Snapshot {
56 data: Arc<Data>,
57 pub owner_id: String,
59 pub connection: String,
61 pub service: String,
63 pub epoch: u64,
65 pub sequence: u64,
67 pub observed_at: Instant,
69 pub valid_until: Option<Instant>,
71}
72
73impl Snapshot {
74 pub fn instances(&self) -> &[Instance] {
78 &self.data.instances
79 }
80
81 pub fn identities(&self) -> &[String] {
85 &self.data.identities
86 }
87
88 pub fn evidence(&self) -> EvidenceKind {
92 self.data.evidence
93 }
94}
95
96#[derive(Clone)]
98pub struct DiscoveryHandle {
99 pub(crate) core: Weak<Core>,
100 pub(crate) connection: Arc<Connection>,
101}
102
103pub struct Observation {
105 core: Weak<Core>,
106 connection: Arc<Connection>,
107 slot: Arc<Slot>,
108 changes: watch::Receiver<u64>,
109}
110
111impl Clone for Observation {
112 fn clone(&self) -> Self {
116 Self {
117 core: self.core.clone(),
118 connection: self.connection.clone(),
119 slot: self.slot.clone(),
120 changes: self.changes.clone(),
121 }
122 }
123}
124
125impl DiscoveryHandle {
126 pub async fn observe(&self, service: &str, deadline: Instant) -> Result<Observation> {
130 if Instant::now() >= deadline {
132 return Err(timeout());
133 }
134 let core = self.core.upgrade().ok_or_else(closed)?;
135 if service.is_empty() || service.len() > 512 || service.trim() != service {
136 return Err(crate::config::invalid("invalid discovery service"));
137 }
138 let mut observation = {
139 let gate = lock(&core.gate);
140 if gate.failed
141 || matches!(gate.phase, OwnerPhase::Prepared | OwnerPhase::Closed)
142 || gate.close_started
143 {
144 return Err(closed());
145 }
146 let mut slots = lock(&self.connection.slots);
147 if Instant::now() >= deadline {
149 return Err(timeout());
150 }
151 let slot = if let Some(slot) = slots.get(service) {
152 let mut state = lock(&slot.state);
153 if Instant::now() >= deadline {
154 return Err(timeout());
155 }
156 state.touched = Instant::now();
157 slot.clone()
158 } else {
159 if slots.len() >= core.limits.max_service_watches_per_connection {
160 return Err(DiscoveryError::new(
161 ErrorCode::CapacityExceeded,
162 "service observation capacity reached",
163 ));
164 }
165 let slot = Arc::new(Slot {
166 service: service.into(),
167 state: Mutex::new(State {
168 data: None,
169 valid_until: None,
170 observed_at: None,
171 error: None,
172 sequence: 0,
173 epoch: 0,
174 revocation: 0,
175 touched: Instant::now(),
176 }),
177 changes: watch::channel(0).0,
178 });
179 slots.insert(service.into(), slot.clone());
180 let (shared, connection, target) =
181 (core.clone(), self.connection.clone(), slot.clone());
182 spawn_owned(&core, async move {
183 observe_loop(shared, connection, target).await
184 });
185 slot
186 };
187 Observation {
188 core: self.core.clone(),
189 connection: self.connection.clone(),
190 changes: slot.changes.subscribe(),
191 slot,
192 }
193 };
194 loop {
195 if Instant::now() >= deadline {
197 return Err(timeout());
198 }
199 if lock(&observation.slot.state).sequence != 0 {
200 let snapshot = observation.snapshot(false);
201 if Instant::now() >= deadline {
203 return Err(timeout());
204 }
205 snapshot?;
206 return Ok(observation);
207 }
208 tokio::select! {
209 _ = core.stop.cancelled() => return Err(closed()),
210 result = tokio::time::timeout_at(deadline, observation.changes.changed()) => {
211 result.map_err(|_| timeout())?.map_err(|_| closed())?;
212 }
213 }
214 }
215 }
216
217 pub fn name(&self) -> &str {
221 &self.connection.name
222 }
223}
224
225impl Observation {
226 pub fn snapshot(&self, strict: bool) -> Result<Snapshot> {
230 let core = self.core.upgrade().ok_or_else(closed)?;
231 let gate = lock(&core.gate);
232 if gate.failed || gate.close_started || gate.phase == OwnerPhase::Closed {
233 return Err(closed());
234 }
235 let state = lock(&self.slot.state);
236 check(&self.connection, &state, strict)?;
237 Ok(Snapshot {
238 data: state.data.as_ref().ok_or_else(closed)?.clone(),
239 owner_id: core.id.clone(),
240 connection: self.connection.name.clone(),
241 service: self.slot.service.clone(),
242 epoch: state.epoch,
243 sequence: state.sequence,
244 observed_at: state.observed_at.ok_or_else(closed)?,
245 valid_until: state.valid_until,
246 })
247 }
248
249 pub async fn changed(&mut self, deadline: Instant) -> Result<()> {
253 let core = self.core.upgrade().ok_or_else(closed)?;
254 let expiry = lock(&self.slot.state).valid_until;
255 let wake_at = expiry.map_or(deadline, |expiry| expiry.min(deadline));
256 tokio::select! {
257 _ = core.stop.cancelled() => Err(closed()),
258 result = self.changes.changed() => result.map_err(|_| closed()),
259 _ = tokio::time::sleep_until(wake_at) => if expiry.is_some_and(|expiry| expiry <= deadline) { Ok(()) } else { Err(timeout()) },
261 }
262 }
263
264 pub fn submit_if_current<T>(
268 &self,
269 instance: &Instance,
270 submit: impl FnOnce() -> T,
271 ) -> Result<T> {
272 self.submit(instance, false, submit)
273 }
274
275 pub fn submit_if_current_strict<T>(
279 &self,
280 instance: &Instance,
281 submit: impl FnOnce() -> T,
282 ) -> Result<T> {
283 self.submit(instance, true, submit)
284 }
285
286 fn submit<T>(
290 &self,
291 instance: &Instance,
292 strict: bool,
293 submit: impl FnOnce() -> T,
294 ) -> Result<T> {
295 let core = self.core.upgrade().ok_or_else(closed)?;
296 let gate = lock(&core.gate);
297 if gate.failed || gate.close_started || gate.phase == OwnerPhase::Closed {
298 return Err(closed());
299 }
300 let state = lock(&self.slot.state);
301 check(&self.connection, &state, strict)?;
302 if !is_traffic_instance(instance)
303 || !state
304 .data
305 .as_ref()
306 .is_some_and(|data| data.instances.contains(instance))
307 {
308 return Err(DiscoveryError::new(
309 ErrorCode::NoInstances,
310 "candidate no longer has routing eligibility",
311 ));
312 }
313 Ok(submit())
314 }
315}
316
317pub(crate) fn check(connection: &Connection, state: &State, strict: bool) -> Result<()> {
321 if state.epoch != connection.epoch.load(Ordering::Acquire)
322 || state.revocation != connection.revocation.load(Ordering::Acquire)
323 {
324 return Err(closed());
325 }
326 if strict && state.error.is_some() {
327 return Err(state.error.clone().unwrap());
328 }
329 if state.data.is_none()
330 || state.observed_at.is_none()
331 || state
332 .valid_until
333 .is_some_and(|deadline| Instant::now() >= deadline)
334 || (state.valid_until.is_none()
335 && state
336 .data
337 .as_ref()
338 .is_none_or(|data| data.evidence != EvidenceKind::Static))
339 {
340 return Err(state.error.clone().unwrap_or_else(|| {
341 DiscoveryError::new(
342 ErrorCode::Unavailable,
343 "discovery evidence is unavailable or expired",
344 )
345 }));
346 }
347 Ok(())
348}
349
350fn record_failure(connection: &Connection, state: &mut State, error: &DiscoveryError) {
354 let grace = if error.revokes_evidence() {
355 0
356 } else {
357 connection
358 .policy
359 .stale_if_error_ms
360 .min(connection.policy.max_age_ms)
361 };
362 let failure_deadline = state
363 .observed_at
364 .map_or(Instant::now(), |at| at + millis(grace));
365 state.valid_until = Some(
367 state
368 .valid_until
369 .map_or(failure_deadline, |current| current.min(failure_deadline)),
370 );
371 state.error = Some(error.clone());
372}
373
374async fn observe_loop(core: Arc<Core>, connection: Arc<Connection>, slot: Arc<Slot>) {
378 let mut backoff = connection.policy.initial_backoff_ms;
379 loop {
380 let epoch = connection.epoch.load(Ordering::Acquire);
381 let revocation = connection.revocation.load(Ordering::Acquire);
382 let deadline = Instant::now() + millis(connection.policy.request_timeout_ms);
383 let started = Instant::now();
384 let result = operation(&core, deadline, connection.provider.observe(&slot.service)).await;
385 let failed = result.is_err();
386 connection.queries.fetch_add(1, Ordering::Relaxed);
387 connection.query_micros.fetch_add(
388 started.elapsed().as_micros().min(u64::MAX as u128) as u64,
389 Ordering::Relaxed,
390 );
391 if result.is_err() {
392 connection.query_errors.fetch_add(1, Ordering::Relaxed);
393 }
394 {
395 let gate = lock(&core.gate);
396 if gate.failed || gate.close_started || core.stop.is_cancelled() {
397 break;
398 }
399 if epoch == connection.epoch.load(Ordering::Acquire)
400 && revocation == connection.revocation.load(Ordering::Acquire)
401 {
402 let mut state = lock(&slot.state);
403 let result = result.and_then(|snapshot| {
404 if let Some(data) = state.data.as_ref().filter(|data| {
405 data.evidence == snapshot.evidence && data.instances == snapshot.instances
406 }) {
407 Ok(data.clone())
408 } else {
409 if state.data.as_ref().is_some_and(|data| {
411 data.instances
412 .iter()
413 .filter(|instance| is_traffic_instance(instance))
414 .any(|instance| !snapshot.instances.contains(instance))
415 }) {
416 state.valid_until = Some(Instant::now());
417 }
418 allocate(
419 &core,
420 &connection,
421 snapshot.instances,
422 snapshot.evidence,
423 &mut state.data,
424 )
425 }
426 });
427 state.sequence = state.sequence.wrapping_add(1);
428 state.epoch = epoch;
429 state.revocation = revocation;
430 match result {
431 Ok(data) => {
432 let now = Instant::now();
433 state.valid_until = (data.evidence != EvidenceKind::Static)
434 .then(|| now + millis(connection.policy.max_age_ms));
435 state.data = Some(data);
436 state.observed_at = Some(now);
437 state.error = None;
438 }
439 Err(error) => {
440 record_failure(&connection, &mut state, &error);
441 if matches!(
442 error.code,
443 ErrorCode::Unauthenticated
444 | ErrorCode::PermissionDenied
445 | ErrorCode::PermissionFiltered
446 ) {
447 drop(state);
448 crate::revoke_connection(&core, &gate, &connection);
449 }
450 }
451 }
452 slot.changes
453 .send_modify(|value| *value = value.wrapping_add(1));
454 }
455 }
456 let delay = if failed {
457 let delay = crate::retry_delay(
458 &core,
459 &connection,
460 backoff,
461 connection.queries.load(Ordering::Relaxed),
462 );
463 backoff = backoff
464 .saturating_mul(2)
465 .min(connection.policy.max_backoff_ms);
466 delay
467 } else {
468 backoff = connection.policy.initial_backoff_ms;
469 millis(connection.policy.poll_interval_ms)
470 };
471 let wait_for_change = async {
472 if failed {
473 return std::future::pending().await;
474 }
475 if let Some(timeout) = connection.provider.change_wait_timeout() {
476 let streams = connection.provider.change_wait_streams().max(1);
478 let _permit = core
479 .watches
480 .acquire_many(streams)
481 .await
482 .map_err(|_| closed())?;
483 tokio::time::timeout(timeout, connection.provider.wait_for_change(&slot.service))
484 .await
485 .unwrap_or_else(|_| Err(crate::timeout()))
486 } else {
487 std::future::pending().await
488 }
489 };
490 tokio::select! {
491 _ = core.stop.cancelled() => break,
492 _ = tokio::time::sleep(delay) => {},
493 result = wait_for_change => {
494 if let Err(error) = result {
495 {
496 let gate = lock(&core.gate);
497 let mut state = lock(&slot.state);
498 record_failure(&connection, &mut state, &error);
499 drop(state);
500 if matches!(error.code, ErrorCode::Unauthenticated | ErrorCode::PermissionDenied | ErrorCode::PermissionFiltered) {
501 crate::revoke_connection(&core, &gate, &connection);
502 }
503 slot.changes.send_modify(|value| *value = value.wrapping_add(1));
504 }
505 tokio::select! {
507 _ = core.stop.cancelled() => break,
508 _ = tokio::time::sleep(millis(connection.policy.initial_backoff_ms)) => {},
509 }
510 } else {
511 tokio::select! {
513 _ = core.stop.cancelled() => break,
514 _ = tokio::time::sleep(millis(connection.policy.poll_interval_ms.min(100))) => {},
515 }
516 }
517 }
518 }
519 let _gate = lock(&core.gate);
520 let mut slots = lock(&connection.slots);
521 if slot.changes.receiver_count() == 0
522 && lock(&slot.state).touched.elapsed() >= millis(connection.policy.idle_timeout_ms)
523 {
524 slots.remove(&slot.service);
525 break;
526 }
527 }
528}
529
530fn allocate(
534 core: &Core,
535 connection: &Connection,
536 instances: Vec<Instance>,
537 evidence: EvidenceKind,
538 previous: &mut Option<Arc<Data>>,
539) -> Result<Arc<Data>> {
540 let mut keys = BTreeSet::new();
541 let mut identities = Vec::with_capacity(instances.len());
542 let mut bytes = instances
543 .capacity()
544 .saturating_mul(std::mem::size_of::<Instance>());
545 for instance in &instances {
546 let identity = connection.provider.instance_key(instance);
547 if instance.ip.is_empty()
548 || instance.port == 0
549 || !instance.weight.is_finite()
550 || identity.is_empty()
551 || identity.len() > 2048
552 || !keys.insert(identity.clone())
553 {
554 return Err(DiscoveryError::new(
555 ErrorCode::InvalidSnapshot,
556 "invalid or duplicate instance in complete observation",
557 ));
558 }
559 bytes = bytes.saturating_add(identity.capacity());
560 identities.push(identity);
561 bytes = bytes
562 .saturating_add(instance.ip.capacity())
563 .saturating_add(instance.cluster_name.as_ref().map_or(0, String::capacity))
564 .saturating_add(
565 instance
566 .metadata
567 .capacity()
568 .saturating_mul(std::mem::size_of::<(String, String)>() + 16),
569 );
570 for (key, value) in &instance.metadata {
571 bytes = bytes
572 .saturating_add(key.capacity())
573 .saturating_add(value.capacity());
574 }
575 }
576 bytes = bytes.saturating_add(
577 identities
578 .capacity()
579 .saturating_mul(std::mem::size_of::<String>()),
580 );
581 let mut quota = lock(&core.quota);
582 let reusable = previous.as_mut().and_then(Arc::get_mut);
584 let old_instances = reusable.as_ref().map_or(0, |data| data.instances.len());
585 let old_bytes = reusable.as_ref().map_or(0, |data| data.bytes);
586 let next_instances = (quota.instances - old_instances).saturating_add(instances.len());
587 let next_bytes = (quota.bytes - old_bytes).saturating_add(bytes);
588 if instances.len() > core.limits.max_instances_per_service
589 || next_instances > core.limits.max_total_instances
590 || next_bytes > core.limits.max_snapshot_bytes
591 {
592 return Err(DiscoveryError::new(
593 ErrorCode::CapacityExceeded,
594 "snapshot capacity reached",
595 ));
596 }
597 quota.instances = next_instances;
599 quota.bytes = next_bytes;
600 if let Some(data) = reusable {
601 data.instances = instances;
602 data.identities = identities;
603 data.evidence = evidence;
604 data.bytes = bytes;
605 return Ok(previous.as_ref().unwrap().clone());
606 }
607 Ok(Arc::new(Data {
608 instances,
609 identities,
610 evidence,
611 bytes,
612 quota: core.quota.clone(),
613 }))
614}
615
616#[async_trait::async_trait]
617impl DiscoveryClient for DiscoveryHandle {
618 fn requires_admission(&self) -> bool {
622 true
623 }
624
625 fn submit_if_current(
629 &self,
630 service: &str,
631 instance: &Instance,
632 submit: &mut dyn FnMut(),
633 ) -> anyhow::Result<()> {
634 let slot = lock(&self.connection.slots)
635 .get(service)
636 .cloned()
637 .ok_or_else(closed)?;
638 let observation = Observation {
639 core: self.core.clone(),
640 connection: self.connection.clone(),
641 changes: slot.changes.subscribe(),
642 slot,
643 };
644 Ok(observation.submit_if_current(instance, submit)?)
645 }
646
647 async fn list_services(&self) -> anyhow::Result<Vec<String>> {
651 let core = self.core.upgrade().ok_or_else(closed)?;
652 Ok(operation(
653 &core,
654 Instant::now() + millis(self.connection.policy.request_timeout_ms),
655 self.connection.provider.services(),
656 )
657 .await?)
658 }
659
660 async fn discover(&self, service: &str) -> anyhow::Result<Vec<Instance>> {
664 Ok(self
665 .discover_all(service)
666 .await?
667 .into_iter()
668 .filter(is_traffic_instance)
669 .collect())
670 }
671
672 async fn discover_all(&self, service: &str) -> anyhow::Result<Vec<Instance>> {
676 let observation = self
677 .observe(
678 service,
679 Instant::now() + millis(self.connection.policy.request_timeout_ms),
680 )
681 .await?;
682 Ok(observation.snapshot(false)?.instances().to_vec())
683 }
684
685 async fn watch_with_options(
689 &self,
690 service: &str,
691 options: WatchOptions,
692 ) -> anyhow::Result<ServiceWatch> {
693 options.validate()?;
694 if options.poll_interval < millis(self.connection.policy.poll_interval_ms) {
695 return Err(crate::config::invalid(
696 "watch interval is faster than frozen observation policy",
697 )
698 .into());
699 }
700 let core = self.core.upgrade().ok_or_else(closed)?;
701 let bridge = core.bridges.clone().try_acquire_owned().map_err(|_| {
702 DiscoveryError::new(
703 ErrorCode::CapacityExceeded,
704 "compatibility watch capacity reached",
705 )
706 })?;
707 let mut observation = self
708 .observe(
709 service,
710 Instant::now() + millis(self.connection.policy.request_timeout_ms),
711 )
712 .await?;
713 let initial = observation
714 .snapshot(false)?
715 .instances()
716 .iter()
717 .filter(|i| is_traffic_instance(i))
718 .cloned()
719 .collect();
720 let (sender, receiver) = watch::channel(initial);
721 let cancel = CancellationToken::new();
722 let stopped = cancel.clone();
723 spawn_owned(&core, async move {
724 let _bridge = bridge;
725 loop {
726 tokio::select! {
727 _ = stopped.cancelled() => break,
728 result = observation.changed(Instant::now() + millis(86_400_000)) => if result.is_err() { break; },
729 }
730 let Ok(snapshot) = observation.snapshot(false) else {
731 break;
732 };
733 if sender
734 .send(
735 snapshot
736 .instances()
737 .iter()
738 .filter(|i| is_traffic_instance(i))
739 .cloned()
740 .collect(),
741 )
742 .is_err()
743 {
744 break;
745 }
746 }
747 });
748 Ok(ServiceWatch {
749 guard: Box::new(BridgeGuard(cancel)),
750 receiver,
751 })
752 }
753}
754
755struct BridgeGuard(CancellationToken);
756impl Drop for BridgeGuard {
757 fn drop(&mut self) {
761 self.0.cancel();
762 }
763}
764#[async_trait::async_trait]
765impl ServiceWatchGuard for BridgeGuard {
766 async fn unsubscribe(self: Box<Self>) -> anyhow::Result<()> {
770 self.0.cancel();
771 Ok(())
772 }
773}