1pub mod assembly;
7pub mod config;
8mod observation;
9mod owner;
10pub mod providers;
11mod registration;
12mod status;
13mod transport;
14pub use status::{ConnectionStatus, RuntimeObserver, RuntimeStatus};
15pub use transport::TransportLease;
16
17pub use nadisc::managed;
18pub use observation::{DiscoveryHandle, Observation, Snapshot};
19pub use owner::{
20 CloseReport, CloseTicket, ConnectionSpec, DiscoveryOwner, OwnerBuilder, OwnerPhase,
21 QuiesceTicket,
22};
23pub use registration::{
24 Eligibility, EligibilityController, RegistrationSpec, RegistrationState, RegistrationStatus,
25};
26
27use nadisc::managed::{DiscoveryError, ErrorCode, ManagedProvider, Result};
28use std::{
29 collections::BTreeMap,
30 sync::{
31 atomic::{AtomicU64, Ordering},
32 Arc, Mutex, MutexGuard,
33 },
34 time::Duration,
35};
36use tokio::{
37 sync::{watch, Semaphore},
38 task::JoinHandle,
39 time::Instant,
40};
41use tokio_util::sync::CancellationToken;
42
43pub(crate) struct Core {
44 pub id: String,
45 pub gate: Mutex<Gate>,
46 pub connections: BTreeMap<String, Arc<Connection>>,
47 pub default_connection: Option<String>,
48 pub limits: config::Limits,
49 pub quota: Arc<Mutex<Quota>>,
50 pub operations: Semaphore,
51 pub watches: Semaphore,
52 pub tasks: Mutex<Vec<JoinHandle<()>>>,
53 pub registration_tasks: Mutex<Vec<JoinHandle<()>>>,
54 pub withdrawal_task: Mutex<Option<JoinHandle<()>>>,
55 pub close_task: Mutex<Option<JoinHandle<()>>>,
56 pub bridges: Arc<Semaphore>,
57 pub transports: Arc<Semaphore>,
58 pub stop: CancellationToken,
59 pub registrations_stop: CancellationToken,
60 pub changes: watch::Sender<u64>,
61 pub close: watch::Sender<Option<CloseReport>>,
62 pub withdrawn: watch::Sender<Option<Vec<RegistrationStatus>>>,
63}
64
65pub(crate) struct Gate {
66 pub phase: OwnerPhase,
67 pub registrations: BTreeMap<String, Arc<registration::Record>>,
68 pub frozen: bool,
69 pub close_started: bool,
70 pub failed: bool,
71 pub shutdown_deadline: Option<Instant>,
72}
73
74pub(crate) struct Connection {
75 pub name: String,
76 pub provider: Arc<dyn ManagedProvider>,
77 pub policy: config::ConnectionPolicy,
78 pub slots: Mutex<BTreeMap<String, Arc<observation::Slot>>>,
79 pub prepared: watch::Sender<bool>,
80 pub error: Mutex<Option<DiscoveryError>>,
81 pub epoch: AtomicU64,
82 pub revocation: AtomicU64,
83 pub queries: AtomicU64,
84 pub query_errors: AtomicU64,
85 pub query_micros: AtomicU64,
86}
87
88#[derive(Default)]
89pub(crate) struct Quota {
90 instances: usize,
91 bytes: usize,
92}
93
94pub(crate) fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
98 mutex
99 .lock()
100 .unwrap_or_else(std::sync::PoisonError::into_inner)
101}
102
103pub(crate) fn millis(value: u64) -> Duration {
107 Duration::from_millis(value)
108}
109
110pub(crate) fn changed(core: &Core) {
114 core.changes
115 .send_modify(|sequence| *sequence = sequence.wrapping_add(1));
116}
117
118pub(crate) async fn operation<T>(
122 core: &Core,
123 deadline: Instant,
124 operation: impl std::future::Future<Output = Result<T>>,
125) -> Result<T> {
126 if Instant::now() >= deadline {
127 return Err(timeout());
128 }
129 tokio::select! {
130 biased;
131 _ = core.stop.cancelled() => Err(closed()),
132 result = tokio::time::timeout_at(deadline, async {
133 let _permit = core.operations.acquire().await.map_err(|_| closed())?;
134 tokio::pin!(operation);
135 futures_util::future::poll_fn(|context| {
137 if Instant::now() >= deadline { std::task::Poll::Ready(Err(timeout())) }
138 else { operation.as_mut().poll(context) }
139 }).await
140 }) => result.unwrap_or_else(|_| Err(timeout())),
141 }
142}
143
144pub(crate) fn spawn_owned(
148 core: &Arc<Core>,
149 future: impl std::future::Future<Output = ()> + Send + 'static,
150) -> tokio::task::AbortHandle {
151 spawn_task(core, false, future)
152}
153
154pub(crate) fn spawn_registration(
158 core: &Arc<Core>,
159 future: impl std::future::Future<Output = ()> + Send + 'static,
160) {
161 spawn_task(core, true, future);
162}
163
164fn spawn_task(
168 core: &Arc<Core>,
169 registration: bool,
170 future: impl std::future::Future<Output = ()> + Send + 'static,
171) -> tokio::task::AbortHandle {
172 use futures_util::FutureExt;
173 let shared = core.clone();
174 let task = tokio::spawn(async move {
175 if std::panic::AssertUnwindSafe(future)
176 .catch_unwind()
177 .await
178 .is_err()
179 {
180 let mut gate = lock(&shared.gate);
182 gate.failed = true;
183 gate.phase = OwnerPhase::Quiescing;
184 shared.stop.cancel();
185 shared.registrations_stop.cancel();
186 drop(gate);
187 changed(&shared);
188 }
189 });
190 let mut tasks = lock(if registration {
191 &core.registration_tasks
192 } else {
193 &core.tasks
194 });
195 tasks.retain(|task| !task.is_finished());
196 let abort = task.abort_handle();
197 tasks.push(task);
198 abort
199}
200
201pub(crate) fn closed() -> DiscoveryError {
205 DiscoveryError::new(ErrorCode::OwnerClosed, "discovery owner is closed")
206}
207
208pub(crate) fn timeout() -> DiscoveryError {
212 DiscoveryError::new(ErrorCode::DeadlineExceeded, "discovery deadline elapsed")
213}
214
215pub(crate) fn retry_delay(
219 core: &Core,
220 connection: &Connection,
221 maximum_ms: u64,
222 attempt: u64,
223) -> Duration {
224 use std::hash::{Hash, Hasher};
225 let mut hash = std::collections::hash_map::DefaultHasher::new();
226 (&core.id, &connection.name, attempt).hash(&mut hash);
227 let floor = maximum_ms.div_ceil(2);
228 millis(floor + hash.finish() % (maximum_ms - floor + 1))
229}
230
231pub(crate) fn revoke_connection(core: &Core, gate: &Gate, connection: &Connection) {
235 connection.revocation.fetch_add(1, Ordering::AcqRel);
236 for slot in lock(&connection.slots).values() {
237 let mut state = lock(&slot.state);
238 state.valid_until = Some(Instant::now());
240 state.error = Some(DiscoveryError::new(
241 ErrorCode::Unauthenticated,
242 "connection identity is unavailable",
243 ));
244 slot.changes
245 .send_modify(|value| *value = value.wrapping_add(1));
246 }
247 for record in gate
249 .registrations
250 .values()
251 .filter(|record| std::ptr::eq(record.connection.as_ref(), connection))
252 {
253 let mut status = lock(&record.status);
254 status.valid_until = None;
255 if matches!(
256 status.state,
257 RegistrationState::Active
258 | RegistrationState::PendingHealth
259 | RegistrationState::Suspended
260 ) {
261 status.state = RegistrationState::Suspended;
262 status.error = Some(DiscoveryError::new(
263 ErrorCode::Unauthenticated,
264 "registration identity is unavailable",
265 ));
266 }
267 }
268 changed(core);
269}