pub mod assembly;
pub mod config;
mod observation;
mod owner;
pub mod providers;
mod registration;
mod status;
mod transport;
pub use status::{ConnectionStatus, RuntimeObserver, RuntimeStatus};
pub use transport::TransportLease;
pub use nadisc::managed;
pub use observation::{DiscoveryHandle, Observation, Snapshot};
pub use owner::{
CloseReport, CloseTicket, ConnectionSpec, DiscoveryOwner, OwnerBuilder, OwnerPhase,
QuiesceTicket,
};
pub use registration::{
Eligibility, EligibilityController, RegistrationSpec, RegistrationState, RegistrationStatus,
};
use nadisc::managed::{DiscoveryError, ErrorCode, ManagedProvider, Result};
use std::{
collections::BTreeMap,
sync::{
atomic::{AtomicU64, Ordering},
Arc, Mutex, MutexGuard,
},
time::Duration,
};
use tokio::{
sync::{watch, Semaphore},
task::JoinHandle,
time::Instant,
};
use tokio_util::sync::CancellationToken;
pub(crate) struct Core {
pub id: String,
pub gate: Mutex<Gate>,
pub connections: BTreeMap<String, Arc<Connection>>,
pub default_connection: Option<String>,
pub limits: config::Limits,
pub quota: Arc<Mutex<Quota>>,
pub operations: Semaphore,
pub watches: Semaphore,
pub tasks: Mutex<Vec<JoinHandle<()>>>,
pub registration_tasks: Mutex<Vec<JoinHandle<()>>>,
pub withdrawal_task: Mutex<Option<JoinHandle<()>>>,
pub close_task: Mutex<Option<JoinHandle<()>>>,
pub bridges: Arc<Semaphore>,
pub transports: Arc<Semaphore>,
pub stop: CancellationToken,
pub registrations_stop: CancellationToken,
pub changes: watch::Sender<u64>,
pub close: watch::Sender<Option<CloseReport>>,
pub withdrawn: watch::Sender<Option<Vec<RegistrationStatus>>>,
}
pub(crate) struct Gate {
pub phase: OwnerPhase,
pub registrations: BTreeMap<String, Arc<registration::Record>>,
pub frozen: bool,
pub close_started: bool,
pub failed: bool,
pub shutdown_deadline: Option<Instant>,
}
pub(crate) struct Connection {
pub name: String,
pub provider: Arc<dyn ManagedProvider>,
pub policy: config::ConnectionPolicy,
pub slots: Mutex<BTreeMap<String, Arc<observation::Slot>>>,
pub prepared: watch::Sender<bool>,
pub error: Mutex<Option<DiscoveryError>>,
pub epoch: AtomicU64,
pub revocation: AtomicU64,
pub queries: AtomicU64,
pub query_errors: AtomicU64,
pub query_micros: AtomicU64,
}
#[derive(Default)]
pub(crate) struct Quota {
instances: usize,
bytes: usize,
}
pub(crate) fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
pub(crate) fn millis(value: u64) -> Duration {
Duration::from_millis(value)
}
pub(crate) fn changed(core: &Core) {
core.changes
.send_modify(|sequence| *sequence = sequence.wrapping_add(1));
}
pub(crate) async fn operation<T>(
core: &Core,
deadline: Instant,
operation: impl std::future::Future<Output = Result<T>>,
) -> Result<T> {
if Instant::now() >= deadline {
return Err(timeout());
}
tokio::select! {
biased;
_ = core.stop.cancelled() => Err(closed()),
result = tokio::time::timeout_at(deadline, async {
let _permit = core.operations.acquire().await.map_err(|_| closed())?;
tokio::pin!(operation);
futures_util::future::poll_fn(|context| {
if Instant::now() >= deadline { std::task::Poll::Ready(Err(timeout())) }
else { operation.as_mut().poll(context) }
}).await
}) => result.unwrap_or_else(|_| Err(timeout())),
}
}
pub(crate) fn spawn_owned(
core: &Arc<Core>,
future: impl std::future::Future<Output = ()> + Send + 'static,
) -> tokio::task::AbortHandle {
spawn_task(core, false, future)
}
pub(crate) fn spawn_registration(
core: &Arc<Core>,
future: impl std::future::Future<Output = ()> + Send + 'static,
) {
spawn_task(core, true, future);
}
fn spawn_task(
core: &Arc<Core>,
registration: bool,
future: impl std::future::Future<Output = ()> + Send + 'static,
) -> tokio::task::AbortHandle {
use futures_util::FutureExt;
let shared = core.clone();
let task = tokio::spawn(async move {
if std::panic::AssertUnwindSafe(future)
.catch_unwind()
.await
.is_err()
{
let mut gate = lock(&shared.gate);
gate.failed = true;
gate.phase = OwnerPhase::Quiescing;
shared.stop.cancel();
shared.registrations_stop.cancel();
drop(gate);
changed(&shared);
}
});
let mut tasks = lock(if registration {
&core.registration_tasks
} else {
&core.tasks
});
tasks.retain(|task| !task.is_finished());
let abort = task.abort_handle();
tasks.push(task);
abort
}
pub(crate) fn closed() -> DiscoveryError {
DiscoveryError::new(ErrorCode::OwnerClosed, "discovery owner is closed")
}
pub(crate) fn timeout() -> DiscoveryError {
DiscoveryError::new(ErrorCode::DeadlineExceeded, "discovery deadline elapsed")
}
pub(crate) fn retry_delay(
core: &Core,
connection: &Connection,
maximum_ms: u64,
attempt: u64,
) -> Duration {
use std::hash::{Hash, Hasher};
let mut hash = std::collections::hash_map::DefaultHasher::new();
(&core.id, &connection.name, attempt).hash(&mut hash);
let floor = maximum_ms.div_ceil(2);
millis(floor + hash.finish() % (maximum_ms - floor + 1))
}
pub(crate) fn revoke_connection(core: &Core, gate: &Gate, connection: &Connection) {
connection.revocation.fetch_add(1, Ordering::AcqRel);
for slot in lock(&connection.slots).values() {
let mut state = lock(&slot.state);
state.valid_until = Some(Instant::now());
state.error = Some(DiscoveryError::new(
ErrorCode::Unauthenticated,
"connection identity is unavailable",
));
slot.changes
.send_modify(|value| *value = value.wrapping_add(1));
}
for record in gate
.registrations
.values()
.filter(|record| std::ptr::eq(record.connection.as_ref(), connection))
{
let mut status = lock(&record.status);
status.valid_until = None;
if matches!(
status.state,
RegistrationState::Active
| RegistrationState::PendingHealth
| RegistrationState::Suspended
) {
status.state = RegistrationState::Suspended;
status.error = Some(DiscoveryError::new(
ErrorCode::Unauthenticated,
"registration identity is unavailable",
));
}
}
changed(core);
}