use crate::{
closed, lock, millis, operation, spawn_owned, timeout, Connection, Core, OwnerPhase, Quota,
};
use nadisc::managed::{DiscoveryError, ErrorCode, EvidenceKind, Result};
use nadisc::{
is_traffic_instance, DiscoveryClient, Instance, ServiceWatch, ServiceWatchGuard, WatchOptions,
};
use std::{
collections::BTreeSet,
sync::{atomic::Ordering, Arc, Mutex, Weak},
};
use tokio::{sync::watch, time::Instant};
use tokio_util::sync::CancellationToken;
pub(crate) struct Slot {
pub service: String,
pub state: Mutex<State>,
pub changes: watch::Sender<u64>,
}
pub(crate) struct State {
data: Option<Arc<Data>>,
pub valid_until: Option<Instant>,
pub observed_at: Option<Instant>,
pub error: Option<DiscoveryError>,
sequence: u64,
epoch: u64,
revocation: u64,
touched: Instant,
}
struct Data {
instances: Vec<Instance>,
identities: Vec<String>,
evidence: EvidenceKind,
bytes: usize,
quota: Arc<Mutex<Quota>>,
}
impl Drop for Data {
fn drop(&mut self) {
let mut quota = lock(&self.quota);
quota.instances -= self.instances.len();
quota.bytes -= self.bytes;
}
}
#[derive(Clone)]
pub struct Snapshot {
data: Arc<Data>,
pub owner_id: String,
pub connection: String,
pub service: String,
pub epoch: u64,
pub sequence: u64,
pub observed_at: Instant,
pub valid_until: Option<Instant>,
}
impl Snapshot {
pub fn instances(&self) -> &[Instance] {
&self.data.instances
}
pub fn identities(&self) -> &[String] {
&self.data.identities
}
pub fn evidence(&self) -> EvidenceKind {
self.data.evidence
}
}
#[derive(Clone)]
pub struct DiscoveryHandle {
pub(crate) core: Weak<Core>,
pub(crate) connection: Arc<Connection>,
}
pub struct Observation {
core: Weak<Core>,
connection: Arc<Connection>,
slot: Arc<Slot>,
changes: watch::Receiver<u64>,
}
impl Clone for Observation {
fn clone(&self) -> Self {
Self {
core: self.core.clone(),
connection: self.connection.clone(),
slot: self.slot.clone(),
changes: self.changes.clone(),
}
}
}
impl DiscoveryHandle {
pub async fn observe(&self, service: &str, deadline: Instant) -> Result<Observation> {
if Instant::now() >= deadline {
return Err(timeout());
}
let core = self.core.upgrade().ok_or_else(closed)?;
if service.is_empty() || service.len() > 512 || service.trim() != service {
return Err(crate::config::invalid("invalid discovery service"));
}
let mut observation = {
let gate = lock(&core.gate);
if gate.failed
|| matches!(gate.phase, OwnerPhase::Prepared | OwnerPhase::Closed)
|| gate.close_started
{
return Err(closed());
}
let mut slots = lock(&self.connection.slots);
if Instant::now() >= deadline {
return Err(timeout());
}
let slot = if let Some(slot) = slots.get(service) {
let mut state = lock(&slot.state);
if Instant::now() >= deadline {
return Err(timeout());
}
state.touched = Instant::now();
slot.clone()
} else {
if slots.len() >= core.limits.max_service_watches_per_connection {
return Err(DiscoveryError::new(
ErrorCode::CapacityExceeded,
"service observation capacity reached",
));
}
let slot = Arc::new(Slot {
service: service.into(),
state: Mutex::new(State {
data: None,
valid_until: None,
observed_at: None,
error: None,
sequence: 0,
epoch: 0,
revocation: 0,
touched: Instant::now(),
}),
changes: watch::channel(0).0,
});
slots.insert(service.into(), slot.clone());
let (shared, connection, target) =
(core.clone(), self.connection.clone(), slot.clone());
spawn_owned(&core, async move {
observe_loop(shared, connection, target).await
});
slot
};
Observation {
core: self.core.clone(),
connection: self.connection.clone(),
changes: slot.changes.subscribe(),
slot,
}
};
loop {
if Instant::now() >= deadline {
return Err(timeout());
}
if lock(&observation.slot.state).sequence != 0 {
let snapshot = observation.snapshot(false);
if Instant::now() >= deadline {
return Err(timeout());
}
snapshot?;
return Ok(observation);
}
tokio::select! {
_ = core.stop.cancelled() => return Err(closed()),
result = tokio::time::timeout_at(deadline, observation.changes.changed()) => {
result.map_err(|_| timeout())?.map_err(|_| closed())?;
}
}
}
}
pub fn name(&self) -> &str {
&self.connection.name
}
}
impl Observation {
pub fn snapshot(&self, strict: bool) -> Result<Snapshot> {
let core = self.core.upgrade().ok_or_else(closed)?;
let gate = lock(&core.gate);
if gate.failed || gate.close_started || gate.phase == OwnerPhase::Closed {
return Err(closed());
}
let state = lock(&self.slot.state);
check(&self.connection, &state, strict)?;
Ok(Snapshot {
data: state.data.as_ref().ok_or_else(closed)?.clone(),
owner_id: core.id.clone(),
connection: self.connection.name.clone(),
service: self.slot.service.clone(),
epoch: state.epoch,
sequence: state.sequence,
observed_at: state.observed_at.ok_or_else(closed)?,
valid_until: state.valid_until,
})
}
pub async fn changed(&mut self, deadline: Instant) -> Result<()> {
let core = self.core.upgrade().ok_or_else(closed)?;
let expiry = lock(&self.slot.state).valid_until;
let wake_at = expiry.map_or(deadline, |expiry| expiry.min(deadline));
tokio::select! {
_ = core.stop.cancelled() => Err(closed()),
result = self.changes.changed() => result.map_err(|_| closed()),
_ = tokio::time::sleep_until(wake_at) => if expiry.is_some_and(|expiry| expiry <= deadline) { Ok(()) } else { Err(timeout()) },
}
}
pub fn submit_if_current<T>(
&self,
instance: &Instance,
submit: impl FnOnce() -> T,
) -> Result<T> {
self.submit(instance, false, submit)
}
pub fn submit_if_current_strict<T>(
&self,
instance: &Instance,
submit: impl FnOnce() -> T,
) -> Result<T> {
self.submit(instance, true, submit)
}
fn submit<T>(
&self,
instance: &Instance,
strict: bool,
submit: impl FnOnce() -> T,
) -> Result<T> {
let core = self.core.upgrade().ok_or_else(closed)?;
let gate = lock(&core.gate);
if gate.failed || gate.close_started || gate.phase == OwnerPhase::Closed {
return Err(closed());
}
let state = lock(&self.slot.state);
check(&self.connection, &state, strict)?;
if !is_traffic_instance(instance)
|| !state
.data
.as_ref()
.is_some_and(|data| data.instances.contains(instance))
{
return Err(DiscoveryError::new(
ErrorCode::NoInstances,
"candidate no longer has routing eligibility",
));
}
Ok(submit())
}
}
pub(crate) fn check(connection: &Connection, state: &State, strict: bool) -> Result<()> {
if state.epoch != connection.epoch.load(Ordering::Acquire)
|| state.revocation != connection.revocation.load(Ordering::Acquire)
{
return Err(closed());
}
if strict && state.error.is_some() {
return Err(state.error.clone().unwrap());
}
if state.data.is_none()
|| state.observed_at.is_none()
|| state
.valid_until
.is_some_and(|deadline| Instant::now() >= deadline)
|| (state.valid_until.is_none()
&& state
.data
.as_ref()
.is_none_or(|data| data.evidence != EvidenceKind::Static))
{
return Err(state.error.clone().unwrap_or_else(|| {
DiscoveryError::new(
ErrorCode::Unavailable,
"discovery evidence is unavailable or expired",
)
}));
}
Ok(())
}
fn record_failure(connection: &Connection, state: &mut State, error: &DiscoveryError) {
let grace = if error.revokes_evidence() {
0
} else {
connection
.policy
.stale_if_error_ms
.min(connection.policy.max_age_ms)
};
let failure_deadline = state
.observed_at
.map_or(Instant::now(), |at| at + millis(grace));
state.valid_until = Some(
state
.valid_until
.map_or(failure_deadline, |current| current.min(failure_deadline)),
);
state.error = Some(error.clone());
}
async fn observe_loop(core: Arc<Core>, connection: Arc<Connection>, slot: Arc<Slot>) {
let mut backoff = connection.policy.initial_backoff_ms;
loop {
let epoch = connection.epoch.load(Ordering::Acquire);
let revocation = connection.revocation.load(Ordering::Acquire);
let deadline = Instant::now() + millis(connection.policy.request_timeout_ms);
let started = Instant::now();
let result = operation(&core, deadline, connection.provider.observe(&slot.service)).await;
let failed = result.is_err();
connection.queries.fetch_add(1, Ordering::Relaxed);
connection.query_micros.fetch_add(
started.elapsed().as_micros().min(u64::MAX as u128) as u64,
Ordering::Relaxed,
);
if result.is_err() {
connection.query_errors.fetch_add(1, Ordering::Relaxed);
}
{
let gate = lock(&core.gate);
if gate.failed || gate.close_started || core.stop.is_cancelled() {
break;
}
if epoch == connection.epoch.load(Ordering::Acquire)
&& revocation == connection.revocation.load(Ordering::Acquire)
{
let mut state = lock(&slot.state);
let result = result.and_then(|snapshot| {
if let Some(data) = state.data.as_ref().filter(|data| {
data.evidence == snapshot.evidence && data.instances == snapshot.instances
}) {
Ok(data.clone())
} else {
if state.data.as_ref().is_some_and(|data| {
data.instances
.iter()
.filter(|instance| is_traffic_instance(instance))
.any(|instance| !snapshot.instances.contains(instance))
}) {
state.valid_until = Some(Instant::now());
}
allocate(
&core,
&connection,
snapshot.instances,
snapshot.evidence,
&mut state.data,
)
}
});
state.sequence = state.sequence.wrapping_add(1);
state.epoch = epoch;
state.revocation = revocation;
match result {
Ok(data) => {
let now = Instant::now();
state.valid_until = (data.evidence != EvidenceKind::Static)
.then(|| now + millis(connection.policy.max_age_ms));
state.data = Some(data);
state.observed_at = Some(now);
state.error = None;
}
Err(error) => {
record_failure(&connection, &mut state, &error);
if matches!(
error.code,
ErrorCode::Unauthenticated
| ErrorCode::PermissionDenied
| ErrorCode::PermissionFiltered
) {
drop(state);
crate::revoke_connection(&core, &gate, &connection);
}
}
}
slot.changes
.send_modify(|value| *value = value.wrapping_add(1));
}
}
let delay = if failed {
let delay = crate::retry_delay(
&core,
&connection,
backoff,
connection.queries.load(Ordering::Relaxed),
);
backoff = backoff
.saturating_mul(2)
.min(connection.policy.max_backoff_ms);
delay
} else {
backoff = connection.policy.initial_backoff_ms;
millis(connection.policy.poll_interval_ms)
};
let wait_for_change = async {
if failed {
return std::future::pending().await;
}
if let Some(timeout) = connection.provider.change_wait_timeout() {
let streams = connection.provider.change_wait_streams().max(1);
let _permit = core
.watches
.acquire_many(streams)
.await
.map_err(|_| closed())?;
tokio::time::timeout(timeout, connection.provider.wait_for_change(&slot.service))
.await
.unwrap_or_else(|_| Err(crate::timeout()))
} else {
std::future::pending().await
}
};
tokio::select! {
_ = core.stop.cancelled() => break,
_ = tokio::time::sleep(delay) => {},
result = wait_for_change => {
if let Err(error) = result {
{
let gate = lock(&core.gate);
let mut state = lock(&slot.state);
record_failure(&connection, &mut state, &error);
drop(state);
if matches!(error.code, ErrorCode::Unauthenticated | ErrorCode::PermissionDenied | ErrorCode::PermissionFiltered) {
crate::revoke_connection(&core, &gate, &connection);
}
slot.changes.send_modify(|value| *value = value.wrapping_add(1));
}
tokio::select! {
_ = core.stop.cancelled() => break,
_ = tokio::time::sleep(millis(connection.policy.initial_backoff_ms)) => {},
}
} else {
tokio::select! {
_ = core.stop.cancelled() => break,
_ = tokio::time::sleep(millis(connection.policy.poll_interval_ms.min(100))) => {},
}
}
}
}
let _gate = lock(&core.gate);
let mut slots = lock(&connection.slots);
if slot.changes.receiver_count() == 0
&& lock(&slot.state).touched.elapsed() >= millis(connection.policy.idle_timeout_ms)
{
slots.remove(&slot.service);
break;
}
}
}
fn allocate(
core: &Core,
connection: &Connection,
instances: Vec<Instance>,
evidence: EvidenceKind,
previous: &mut Option<Arc<Data>>,
) -> Result<Arc<Data>> {
let mut keys = BTreeSet::new();
let mut identities = Vec::with_capacity(instances.len());
let mut bytes = instances
.capacity()
.saturating_mul(std::mem::size_of::<Instance>());
for instance in &instances {
let identity = connection.provider.instance_key(instance);
if instance.ip.is_empty()
|| instance.port == 0
|| !instance.weight.is_finite()
|| identity.is_empty()
|| identity.len() > 2048
|| !keys.insert(identity.clone())
{
return Err(DiscoveryError::new(
ErrorCode::InvalidSnapshot,
"invalid or duplicate instance in complete observation",
));
}
bytes = bytes.saturating_add(identity.capacity());
identities.push(identity);
bytes = bytes
.saturating_add(instance.ip.capacity())
.saturating_add(instance.cluster_name.as_ref().map_or(0, String::capacity))
.saturating_add(
instance
.metadata
.capacity()
.saturating_mul(std::mem::size_of::<(String, String)>() + 16),
);
for (key, value) in &instance.metadata {
bytes = bytes
.saturating_add(key.capacity())
.saturating_add(value.capacity());
}
}
bytes = bytes.saturating_add(
identities
.capacity()
.saturating_mul(std::mem::size_of::<String>()),
);
let mut quota = lock(&core.quota);
let reusable = previous.as_mut().and_then(Arc::get_mut);
let old_instances = reusable.as_ref().map_or(0, |data| data.instances.len());
let old_bytes = reusable.as_ref().map_or(0, |data| data.bytes);
let next_instances = (quota.instances - old_instances).saturating_add(instances.len());
let next_bytes = (quota.bytes - old_bytes).saturating_add(bytes);
if instances.len() > core.limits.max_instances_per_service
|| next_instances > core.limits.max_total_instances
|| next_bytes > core.limits.max_snapshot_bytes
{
return Err(DiscoveryError::new(
ErrorCode::CapacityExceeded,
"snapshot capacity reached",
));
}
quota.instances = next_instances;
quota.bytes = next_bytes;
if let Some(data) = reusable {
data.instances = instances;
data.identities = identities;
data.evidence = evidence;
data.bytes = bytes;
return Ok(previous.as_ref().unwrap().clone());
}
Ok(Arc::new(Data {
instances,
identities,
evidence,
bytes,
quota: core.quota.clone(),
}))
}
#[async_trait::async_trait]
impl DiscoveryClient for DiscoveryHandle {
fn requires_admission(&self) -> bool {
true
}
fn submit_if_current(
&self,
service: &str,
instance: &Instance,
submit: &mut dyn FnMut(),
) -> anyhow::Result<()> {
let slot = lock(&self.connection.slots)
.get(service)
.cloned()
.ok_or_else(closed)?;
let observation = Observation {
core: self.core.clone(),
connection: self.connection.clone(),
changes: slot.changes.subscribe(),
slot,
};
Ok(observation.submit_if_current(instance, submit)?)
}
async fn list_services(&self) -> anyhow::Result<Vec<String>> {
let core = self.core.upgrade().ok_or_else(closed)?;
Ok(operation(
&core,
Instant::now() + millis(self.connection.policy.request_timeout_ms),
self.connection.provider.services(),
)
.await?)
}
async fn discover(&self, service: &str) -> anyhow::Result<Vec<Instance>> {
Ok(self
.discover_all(service)
.await?
.into_iter()
.filter(is_traffic_instance)
.collect())
}
async fn discover_all(&self, service: &str) -> anyhow::Result<Vec<Instance>> {
let observation = self
.observe(
service,
Instant::now() + millis(self.connection.policy.request_timeout_ms),
)
.await?;
Ok(observation.snapshot(false)?.instances().to_vec())
}
async fn watch_with_options(
&self,
service: &str,
options: WatchOptions,
) -> anyhow::Result<ServiceWatch> {
options.validate()?;
if options.poll_interval < millis(self.connection.policy.poll_interval_ms) {
return Err(crate::config::invalid(
"watch interval is faster than frozen observation policy",
)
.into());
}
let core = self.core.upgrade().ok_or_else(closed)?;
let bridge = core.bridges.clone().try_acquire_owned().map_err(|_| {
DiscoveryError::new(
ErrorCode::CapacityExceeded,
"compatibility watch capacity reached",
)
})?;
let mut observation = self
.observe(
service,
Instant::now() + millis(self.connection.policy.request_timeout_ms),
)
.await?;
let initial = observation
.snapshot(false)?
.instances()
.iter()
.filter(|i| is_traffic_instance(i))
.cloned()
.collect();
let (sender, receiver) = watch::channel(initial);
let cancel = CancellationToken::new();
let stopped = cancel.clone();
spawn_owned(&core, async move {
let _bridge = bridge;
loop {
tokio::select! {
_ = stopped.cancelled() => break,
result = observation.changed(Instant::now() + millis(86_400_000)) => if result.is_err() { break; },
}
let Ok(snapshot) = observation.snapshot(false) else {
break;
};
if sender
.send(
snapshot
.instances()
.iter()
.filter(|i| is_traffic_instance(i))
.cloned()
.collect(),
)
.is_err()
{
break;
}
}
});
Ok(ServiceWatch {
guard: Box::new(BridgeGuard(cancel)),
receiver,
})
}
}
struct BridgeGuard(CancellationToken);
impl Drop for BridgeGuard {
fn drop(&mut self) {
self.0.cancel();
}
}
#[async_trait::async_trait]
impl ServiceWatchGuard for BridgeGuard {
async fn unsubscribe(self: Box<Self>) -> anyhow::Result<()> {
self.0.cancel();
Ok(())
}
}