use crate::{
changed, closed,
config::{invalid, validate_name, ConnectionPolicy, Limits},
lock, millis, operation,
registration::{Record, RegistrationSpec, RegistrationStatus},
spawn_owned, timeout, Connection, Core, DiscoveryHandle, Gate, Quota,
};
use nadisc::managed::{DiscoveryError, ErrorCode, ManagedProvider, Result};
use std::{
collections::{BTreeMap, BTreeSet},
sync::{atomic::AtomicU64, Arc, Mutex},
};
use tokio::{
sync::{watch, Semaphore},
time::Instant,
};
use tokio_util::sync::CancellationToken;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OwnerPhase {
Prepared,
Starting,
Serving,
Quiescing,
Closed,
}
pub struct ConnectionSpec {
pub name: String,
pub provider: Arc<dyn ManagedProvider>,
pub policy: ConnectionPolicy,
}
pub struct OwnerBuilder {
limits: Limits,
connections: BTreeMap<String, ConnectionSpec>,
default_connection: Option<String>,
}
impl Default for OwnerBuilder {
fn default() -> Self {
Self {
limits: Limits::default(),
connections: BTreeMap::new(),
default_connection: None,
}
}
}
impl OwnerBuilder {
pub fn limits(mut self, limits: Limits) -> Self {
self.limits = limits;
self
}
pub fn default_connection(mut self, name: impl Into<String>) -> Self {
self.default_connection = Some(name.into());
self
}
pub fn connection(mut self, spec: ConnectionSpec) -> Result<Self> {
validate_name(&spec.name)?;
spec.policy.validate()?;
if self.connections.contains_key(&spec.name) {
return Err(invalid("duplicate connection name"));
}
self.connections.insert(spec.name.clone(), spec);
Ok(self)
}
pub fn build(self) -> Result<DiscoveryOwner> {
self.limits.validate()?;
if self.connections.len() > self.limits.max_connections {
return Err(invalid("connection capacity exceeded"));
}
let default_connection = match self.default_connection {
Some(name) if self.connections.contains_key(&name) => Some(name),
Some(_) => return Err(invalid("unknown default connection")),
None => self
.connections
.contains_key("primary")
.then(|| "primary".into()),
};
let connections = self
.connections
.into_iter()
.map(|(name, spec)| {
let connection = Arc::new(Connection {
name: name.clone(),
provider: spec.provider,
policy: spec.policy,
slots: Mutex::new(BTreeMap::new()),
prepared: watch::channel(false).0,
error: Mutex::new(None),
epoch: AtomicU64::new(1),
revocation: AtomicU64::new(0),
queries: AtomicU64::new(0),
query_errors: AtomicU64::new(0),
query_micros: AtomicU64::new(0),
});
(name, connection)
})
.collect();
let core = Arc::new(Core {
id: uuid::Uuid::new_v4().to_string(),
gate: Mutex::new(Gate {
phase: OwnerPhase::Prepared,
registrations: BTreeMap::new(),
frozen: false,
close_started: false,
failed: false,
shutdown_deadline: None,
}),
connections,
default_connection,
operations: Semaphore::new(self.limits.max_parallel_operations),
watches: Semaphore::new(self.limits.max_watch_streams),
quota: Arc::new(Mutex::new(Quota::default())),
bridges: Arc::new(Semaphore::new(
self.limits
.max_service_watches_per_connection
.saturating_mul(self.limits.max_connections),
)),
transports: Arc::new(Semaphore::new(self.limits.max_transport_tasks)),
limits: self.limits,
tasks: Mutex::new(Vec::new()),
registration_tasks: Mutex::new(Vec::new()),
withdrawal_task: Mutex::new(None),
close_task: Mutex::new(None),
stop: CancellationToken::new(),
registrations_stop: CancellationToken::new(),
changes: watch::channel(0).0,
close: watch::channel(None).0,
withdrawn: watch::channel(None).0,
});
Ok(DiscoveryOwner { core })
}
}
pub struct DiscoveryOwner {
core: Arc<Core>,
}
impl DiscoveryOwner {
pub async fn prepare(&self, deadline: Instant) -> Result<()> {
if Instant::now() >= deadline {
return Err(timeout());
}
{
let mut gate = lock(&self.core.gate);
match gate.phase {
OwnerPhase::Prepared => {
gate.phase = OwnerPhase::Starting;
for connection in self.core.connections.values() {
let (core, connection) = (self.core.clone(), connection.clone());
spawn_owned(
&self.core,
async move { prepare_loop(core, connection).await },
);
}
}
OwnerPhase::Starting => {}
_ => return Err(invalid("owner cannot prepare in its current phase")),
}
}
let mut changes = self.core.changes.subscribe();
loop {
if self.core.stop.is_cancelled() {
return Err(closed());
}
if Instant::now() >= deadline {
return Err(timeout());
}
let ready = self
.core
.connections
.values()
.filter(|connection| connection.policy.startup_required)
.all(|connection| *connection.prepared.borrow());
if ready {
return Ok(());
}
tokio::time::timeout_at(deadline, changes.changed())
.await
.map_err(|_| timeout())?
.map_err(|_| closed())?;
}
}
pub fn client(&self, qualifier: Option<&str>) -> Result<DiscoveryHandle> {
let gate = lock(&self.core.gate);
if gate.close_started || gate.failed {
return Err(closed());
}
let name = qualifier
.or(self.core.default_connection.as_deref())
.ok_or_else(|| invalid("no default discovery connection"))?;
let connection = self
.core
.connections
.get(name)
.cloned()
.ok_or_else(|| invalid("unknown discovery connection"))?;
Ok(DiscoveryHandle {
core: Arc::downgrade(&self.core),
connection,
})
}
pub fn freeze_registrations(&self, specs: Vec<RegistrationSpec>) -> Result<()> {
let mut gate = lock(&self.core.gate);
if gate.frozen || gate.phase != OwnerPhase::Starting {
return Err(invalid(
"registration plans are already frozen or owner is not starting",
));
}
if specs.len() > self.core.limits.max_registrations {
return Err(invalid("registration capacity exceeded"));
}
let mut records = BTreeMap::new();
let mut keys = BTreeSet::new();
for mut spec in specs {
validate_name(&spec.plan.id)?;
let connection = self
.core
.connections
.get(&spec.connection_ref)
.cloned()
.ok_or_else(|| invalid("unknown registration connection"))?;
spec.policy.validate(&connection.policy)?;
if !connection.provider.capabilities().registration {
return Err(invalid("registration provider is read only"));
}
if spec
.plan
.instance
.metadata
.contains_key("nasa.discovery.owner")
{
return Err(invalid("registration ownership metadata is reserved"));
}
spec.plan
.instance
.metadata
.insert("nasa.discovery.owner".into(), self.core.id.clone());
let instance = &spec.plan.instance;
if instance.ip.trim() != instance.ip
|| instance.ip.is_empty()
|| instance.port == 0
|| instance.ip == "0.0.0.0"
|| instance.ip == "::"
|| instance.ip == "[::]"
|| !instance.weight.is_finite()
|| instance.weight < 0.0
|| spec.plan.service.is_empty()
{
return Err(invalid("registration endpoint or weight is invalid"));
}
connection.provider.validate_registration(&spec.plan)?;
if !keys.insert(connection.provider.registration_key(&spec.plan))
|| records.contains_key(&spec.plan.id)
{
return Err(invalid(
"duplicate external registration key or local registration name",
));
}
records.insert(
spec.plan.id.clone(),
Arc::new(Record::new(spec, connection)),
);
}
gate.registrations = records;
gate.frozen = true;
Ok(())
}
pub async fn publish(&self, deadline: Instant) -> Result<()> {
if Instant::now() >= deadline {
return Err(timeout());
}
{
let gate = lock(&self.core.gate);
if !gate.frozen
|| gate.phase != OwnerPhase::Starting
|| gate.failed
|| gate.close_started
{
return Err(invalid("owner is not ready to publish frozen plans"));
}
for record in gate.registrations.values() {
if !record
.started
.swap(true, std::sync::atomic::Ordering::AcqRel)
{
let (core, record) = (self.core.clone(), record.clone());
crate::spawn_registration(&self.core, async move {
crate::registration::run(core, record).await
});
}
}
}
let mut changes = self.core.changes.subscribe();
loop {
{
let mut gate = lock(&self.core.gate);
if gate.phase != OwnerPhase::Starting || gate.failed || gate.close_started {
return Err(closed());
}
if Instant::now() >= deadline {
return Err(timeout());
}
let ready = gate
.registrations
.values()
.filter(|record| record.spec.policy.startup_required)
.all(|record| record.current_status(&gate).is_ready());
if ready {
gate.phase = OwnerPhase::Serving;
return Ok(());
}
}
tokio::time::timeout_at(deadline, changes.changed())
.await
.map_err(|_| timeout())?
.map_err(|_| closed())?;
}
}
pub fn begin_quiesce(&self, deadline: Instant) -> QuiesceTicket {
begin_quiesce(self.core.clone(), deadline)
}
pub fn begin_close(&self, deadline: Instant) -> CloseTicket {
begin_close(self.core.clone(), deadline)
}
pub fn phase(&self) -> OwnerPhase {
lock(&self.core.gate).phase
}
pub fn registrations(&self) -> Vec<RegistrationStatus> {
let gate = lock(&self.core.gate);
gate.registrations
.values()
.map(|record| record.current_status(&gate))
.collect()
}
pub fn changes(&self) -> watch::Receiver<u64> {
self.core.changes.subscribe()
}
pub fn observer(&self) -> crate::RuntimeObserver {
crate::RuntimeObserver {
core: Arc::downgrade(&self.core),
}
}
pub fn service_key(&self, connection: &str, service: &str) -> Result<String> {
self.core
.connections
.get(connection)
.ok_or_else(|| invalid("unknown discovery connection"))?
.provider
.service_key(service)
}
}
impl Drop for DiscoveryOwner {
fn drop(&mut self) {
if !lock(&self.core.gate).close_started {
if tokio::runtime::Handle::try_current().is_ok() {
begin_close(self.core.clone(), Instant::now() + millis(3000));
} else {
let mut gate = lock(&self.core.gate);
gate.phase = OwnerPhase::Closed;
gate.close_started = true;
self.core.stop.cancel();
self.core.registrations_stop.cancel();
for task in lock(&self.core.tasks).drain(..) {
task.abort();
}
for task in lock(&self.core.registration_tasks).drain(..) {
task.abort();
}
}
}
}
}
#[derive(Clone, Debug)]
pub struct CloseReport {
pub registrations: Vec<RegistrationStatus>,
pub provider_errors: BTreeMap<String, DiscoveryError>,
pub deadline_elapsed: bool,
}
#[derive(Clone)]
pub struct CloseTicket {
receiver: watch::Receiver<Option<CloseReport>>,
}
impl CloseTicket {
pub async fn wait(mut self) -> Result<CloseReport> {
loop {
if let Some(report) = self.receiver.borrow().clone() {
return Ok(report);
}
self.receiver.changed().await.map_err(|_| closed())?;
}
}
}
#[derive(Clone)]
pub struct QuiesceTicket {
receiver: watch::Receiver<Option<Vec<RegistrationStatus>>>,
}
impl QuiesceTicket {
pub async fn wait(mut self) -> Result<Vec<RegistrationStatus>> {
loop {
if let Some(statuses) = self.receiver.borrow().clone() {
return Ok(statuses);
}
self.receiver.changed().await.map_err(|_| closed())?;
}
}
}
fn begin_quiesce(core: Arc<Core>, deadline: Instant) -> QuiesceTicket {
let mut gate = lock(&core.gate);
let ticket = QuiesceTicket {
receiver: core.withdrawn.subscribe(),
};
if gate.shutdown_deadline.is_none() {
gate.shutdown_deadline = Some(deadline);
gate.phase = OwnerPhase::Quiescing;
core.registrations_stop.cancel();
let shared = core.clone();
*lock(&core.withdrawal_task) = Some(tokio::spawn(async move {
use futures_util::FutureExt;
let _ = std::panic::AssertUnwindSafe(cleanup_registrations(shared.clone(), deadline))
.catch_unwind()
.await;
finish_withdrawal(&shared);
}));
}
ticket
}
fn begin_close(core: Arc<Core>, deadline: Instant) -> CloseTicket {
begin_quiesce(core.clone(), deadline);
let mut gate = lock(&core.gate);
let ticket = CloseTicket {
receiver: core.close.subscribe(),
};
if !gate.close_started {
gate.close_started = true;
gate.phase = OwnerPhase::Quiescing;
core.registrations_stop.cancel();
core.stop.cancel();
let deadline = gate.shutdown_deadline.unwrap_or(deadline).min(deadline);
let shared = core.clone();
*lock(&core.close_task) = Some(tokio::spawn(async move {
use futures_util::FutureExt;
if std::panic::AssertUnwindSafe(cleanup(shared.clone(), deadline))
.catch_unwind()
.await
.is_err()
{
finish_withdrawal(&shared);
lock(&shared.gate).phase = OwnerPhase::Closed;
shared.close.send_replace(Some(CloseReport {
registrations: lock(&shared.gate)
.registrations
.values()
.map(|r| lock(&r.status).clone())
.collect(),
provider_errors: shared
.connections
.keys()
.map(|name| {
(
name.clone(),
DiscoveryError::new(
nadisc::managed::ErrorCode::Unavailable,
"provider cleanup was interrupted",
),
)
})
.collect(),
deadline_elapsed: Instant::now() >= deadline,
}));
}
}));
}
ticket
}
async fn prepare_loop(core: Arc<Core>, connection: Arc<Connection>) {
let mut backoff = connection.policy.initial_backoff_ms;
let mut attempt = 0u64;
loop {
let result = operation(
&core,
Instant::now() + millis(connection.policy.request_timeout_ms),
connection.provider.prepare(),
)
.await;
if core.stop.is_cancelled() {
break;
}
match result {
Ok(()) => {
connection.prepared.send_replace(true);
*lock(&connection.error) = None;
changed(&core);
return;
}
Err(error) => {
if matches!(
error.code,
ErrorCode::Unauthenticated
| ErrorCode::PermissionDenied
| ErrorCode::PermissionFiltered
) {
let gate = lock(&core.gate);
crate::revoke_connection(&core, &gate, &connection);
}
*lock(&connection.error) = Some(error);
changed(&core);
}
}
tokio::select! { _ = core.stop.cancelled() => break, _ = tokio::time::sleep(crate::retry_delay(&core, &connection, backoff, attempt)) => {} }
attempt = attempt.wrapping_add(1);
backoff = backoff
.saturating_mul(2)
.min(connection.policy.max_backoff_ms);
}
}
async fn cleanup_registrations(core: Arc<Core>, deadline: Instant) {
use futures_util::{stream, StreamExt};
let tasks = std::mem::take(&mut *lock(&core.registration_tasks));
for task in &tasks {
task.abort();
}
for task in tasks {
let _ = tokio::time::timeout_at(deadline, task).await;
}
let records: Vec<_> = lock(&core.gate).registrations.values().cloned().collect();
let count = records.len().max(1);
let share = deadline.saturating_duration_since(Instant::now())
/ count.div_ceil(core.limits.max_parallel_operations) as u32;
stream::iter(records.iter().cloned())
.for_each_concurrent(core.limits.max_parallel_operations, |record| async move {
use futures_util::FutureExt;
let _ = std::panic::AssertUnwindSafe(crate::registration::withdraw(
record, deadline, share,
))
.catch_unwind()
.await;
})
.await;
}
fn finish_withdrawal(core: &Core) {
let gate = lock(&core.gate);
for record in gate.registrations.values() {
crate::registration::finish_withdrawal(record);
}
core.withdrawn.send_replace(Some(
gate.registrations
.values()
.map(|record| lock(&record.status).clone())
.collect(),
));
changed(core);
}
async fn cleanup(core: Arc<Core>, deadline: Instant) {
use futures_util::{stream, StreamExt};
let withdrawal = lock(&core.withdrawal_task).take();
if let Some(mut task) = withdrawal {
if tokio::time::timeout_at(deadline, &mut task).await.is_err() {
task.abort();
let _ = task.await;
}
}
finish_withdrawal(&core);
let tasks = std::mem::take(&mut *lock(&core.tasks));
for task in &tasks {
task.abort();
}
for task in tasks {
let _ = tokio::time::timeout_at(deadline, task).await;
}
let records: Vec<_> = lock(&core.gate).registrations.values().cloned().collect();
let providers = stream::iter(core.connections.values().cloned())
.map(|connection| async move {
use futures_util::FutureExt;
let result = tokio::time::timeout_at(
deadline,
std::panic::AssertUnwindSafe(connection.provider.close()).catch_unwind(),
)
.await
.map(|result| {
result.unwrap_or_else(|_| {
Err(DiscoveryError::new(
nadisc::managed::ErrorCode::Unavailable,
"provider cleanup was interrupted",
))
})
})
.unwrap_or_else(|_| Err(timeout()));
(connection.name.clone(), result)
})
.buffer_unordered(core.limits.max_parallel_operations)
.collect::<Vec<_>>()
.await;
let provider_errors = providers
.into_iter()
.filter_map(|(name, result)| result.err().map(|error| (name, error)))
.collect();
let registrations = records
.iter()
.map(|record| lock(&record.status).clone())
.collect();
lock(&core.gate).phase = OwnerPhase::Closed;
core.close.send_replace(Some(CloseReport {
registrations,
provider_errors,
deadline_elapsed: Instant::now() >= deadline,
}));
changed(&core);
}