dynamo_runtime/discovery/
registration.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{sync::Arc, time::Duration};
5
6use anyhow::{Context, Result};
7use dashmap::{DashMap, mapref::entry::Entry as DashEntry};
8use tokio::sync::Mutex;
9use tokio_util::sync::CancellationToken;
10
11use super::{
12    Discovery, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, DiscoverySpec,
13    EndpointInstanceId,
14};
15
16struct RegistrationEntry {
17    instance: DiscoveryInstance,
18    leases: usize,
19    owned: bool,
20}
21
22type RegistrationSlot = Arc<Mutex<Option<RegistrationEntry>>>;
23
24/// Coordinates endpoint registration ownership across all users of one distributed runtime.
25pub(crate) struct EndpointRegistrationManager {
26    discovery: Arc<dyn Discovery>,
27    slots: DashMap<EndpointInstanceId, RegistrationSlot>,
28    runtime: tokio::runtime::Handle,
29    cancellation: CancellationToken,
30}
31
32impl EndpointRegistrationManager {
33    pub(crate) fn new(
34        discovery: Arc<dyn Discovery>,
35        runtime: tokio::runtime::Handle,
36        cancellation: CancellationToken,
37    ) -> Arc<Self> {
38        Arc::new(Self {
39            discovery,
40            slots: DashMap::new(),
41            runtime,
42            cancellation,
43        })
44    }
45
46    pub(crate) async fn register(
47        self: &Arc<Self>,
48        spec: DiscoverySpec,
49    ) -> Result<EndpointRegistrationLease> {
50        let (result_tx, result_rx) = tokio::sync::oneshot::channel();
51        let manager = self.clone();
52        self.runtime.spawn(async move {
53            let result = manager.acquire(spec).await;
54            let _ = result_tx.send(result);
55        });
56        result_rx
57            .await
58            .context("endpoint discovery registration task ended without a result")?
59    }
60
61    async fn acquire(self: &Arc<Self>, spec: DiscoverySpec) -> Result<EndpointRegistrationLease> {
62        let candidate = spec.clone().into_instance(self.discovery.instance_id());
63        let DiscoveryInstanceId::Endpoint(id) = candidate.id() else {
64            anyhow::bail!("endpoint registration leases require an endpoint specification");
65        };
66        let slot = self
67            .slots
68            .entry(id.clone())
69            .or_insert_with(|| Arc::new(Mutex::new(None)))
70            .clone();
71        let mut entry = slot.lock().await;
72        if let Some(entry) = entry.as_mut() {
73            anyhow::ensure!(
74                entry.instance == candidate,
75                "endpoint registration lease conflicts with the existing specification"
76            );
77            entry.leases += 1;
78            return Ok(EndpointRegistrationLease::new(self.clone(), id));
79        }
80
81        let registration = async {
82            let existing = self
83                .discovery
84                .list(DiscoveryQuery::Endpoint {
85                    namespace: id.namespace.clone(),
86                    component: id.component.clone(),
87                    endpoint: id.endpoint.clone(),
88                })
89                .await?
90                .into_iter()
91                .find(|instance| instance.id() == DiscoveryInstanceId::Endpoint(id.clone()));
92            match existing {
93                Some(existing) => {
94                    anyhow::ensure!(
95                        existing == candidate,
96                        "endpoint registration lease conflicts with a pre-existing specification"
97                    );
98                    Ok((existing, false))
99                }
100                None => Ok((self.discovery.register(spec).await?, true)),
101            }
102        }
103        .await;
104        let (instance, owned) = match registration {
105            Ok(registration) => registration,
106            Err(error) => {
107                drop(entry);
108                self.remove_unused_slot(&id, &slot);
109                return Err(error);
110            }
111        };
112        *entry = Some(RegistrationEntry {
113            instance,
114            leases: 1,
115            owned,
116        });
117        Ok(EndpointRegistrationLease::new(self.clone(), id))
118    }
119
120    async fn release(self: Arc<Self>, id: EndpointInstanceId) {
121        let Some(slot) = self.slots.get(&id).map(|slot| slot.clone()) else {
122            return;
123        };
124        let mut entry = slot.lock().await;
125        let Some(registration) = entry.as_mut() else {
126            return;
127        };
128        registration.leases = registration.leases.saturating_sub(1);
129        if registration.leases != 0 {
130            return;
131        }
132        if registration.owned {
133            let mut retry_delay = Duration::from_millis(50);
134            loop {
135                match self
136                    .discovery
137                    .unregister(registration.instance.clone())
138                    .await
139                {
140                    Ok(()) => break,
141                    Err(error) => {
142                        tracing::warn!(
143                            %error,
144                            instance = ?registration.instance.id(),
145                            "Failed to release endpoint discovery registration; retrying"
146                        );
147                    }
148                }
149                tokio::select! {
150                    _ = self.cancellation.cancelled() => return,
151                    _ = tokio::time::sleep(retry_delay) => {}
152                }
153                retry_delay = (retry_delay * 2).min(Duration::from_secs(5));
154            }
155        }
156        *entry = None;
157        drop(entry);
158        self.remove_unused_slot(&id, &slot);
159    }
160
161    fn remove_unused_slot(&self, id: &EndpointInstanceId, slot: &RegistrationSlot) {
162        if let DashEntry::Occupied(entry) = self.slots.entry(id.clone())
163            && Arc::ptr_eq(entry.get(), slot)
164            && Arc::strong_count(slot) == 2
165        {
166            entry.remove();
167        }
168    }
169}
170
171/// Keeps a discovery endpoint registered until the last lease is dropped.
172pub struct EndpointRegistrationLease {
173    manager: Arc<EndpointRegistrationManager>,
174    id: Option<EndpointInstanceId>,
175}
176
177impl EndpointRegistrationLease {
178    fn new(manager: Arc<EndpointRegistrationManager>, id: EndpointInstanceId) -> Self {
179        Self {
180            manager,
181            id: Some(id),
182        }
183    }
184}
185
186impl std::fmt::Debug for EndpointRegistrationLease {
187    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
188        formatter
189            .debug_struct("EndpointRegistrationLease")
190            .field("id", &self.id)
191            .finish()
192    }
193}
194
195impl Drop for EndpointRegistrationLease {
196    fn drop(&mut self) {
197        let Some(id) = self.id.take() else {
198            return;
199        };
200        let manager = self.manager.clone();
201        self.manager.runtime.spawn(async move {
202            manager.release(id).await;
203        });
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210    use crate::{
211        component::TransportType,
212        discovery::{DiscoveryStream, MockDiscovery, SharedMockRegistry},
213    };
214    use async_trait::async_trait;
215
216    struct BlockingRegistrationDiscovery {
217        inner: MockDiscovery,
218        registered: Arc<tokio::sync::Notify>,
219        release: Arc<tokio::sync::Notify>,
220    }
221
222    #[async_trait]
223    impl Discovery for BlockingRegistrationDiscovery {
224        fn instance_id(&self) -> u64 {
225            self.inner.instance_id()
226        }
227
228        async fn register_internal(&self, spec: DiscoverySpec) -> Result<DiscoveryInstance> {
229            let instance = self.inner.register_internal(spec).await?;
230            self.registered.notify_one();
231            self.release.notified().await;
232            Ok(instance)
233        }
234
235        async fn unregister(&self, instance: DiscoveryInstance) -> Result<()> {
236            self.inner.unregister(instance).await
237        }
238
239        async fn list(&self, query: DiscoveryQuery) -> Result<Vec<DiscoveryInstance>> {
240            self.inner.list(query).await
241        }
242
243        async fn list_and_watch(
244            &self,
245            query: DiscoveryQuery,
246            cancel_token: Option<CancellationToken>,
247        ) -> Result<DiscoveryStream> {
248            self.inner.list_and_watch(query, cancel_token).await
249        }
250    }
251
252    fn endpoint_spec() -> DiscoverySpec {
253        DiscoverySpec::Endpoint {
254            namespace: "ns".to_string(),
255            component: "frontend".to_string(),
256            endpoint: "router".to_string(),
257            transport: TransportType::Tcp("127.0.0.1:1/7/router".to_string()),
258            device_type: None,
259            request_plane_codec: None,
260        }
261    }
262
263    async fn wait_for_endpoint_count(discovery: &dyn Discovery, expected: usize) {
264        tokio::time::timeout(Duration::from_secs(1), async {
265            loop {
266                let instances = discovery.list(DiscoveryQuery::AllEndpoints).await.unwrap();
267                if instances.len() == expected {
268                    break;
269                }
270                tokio::task::yield_now().await;
271            }
272        })
273        .await
274        .expect("endpoint registration count did not converge");
275    }
276
277    #[tokio::test]
278    async fn cancelled_registration_releases_endpoint_after_register_completes() {
279        let registered = Arc::new(tokio::sync::Notify::new());
280        let release = Arc::new(tokio::sync::Notify::new());
281        let discovery: Arc<dyn Discovery> = Arc::new(BlockingRegistrationDiscovery {
282            inner: MockDiscovery::new(Some(7), SharedMockRegistry::new()),
283            registered: registered.clone(),
284            release: release.clone(),
285        });
286        let manager = EndpointRegistrationManager::new(
287            discovery.clone(),
288            tokio::runtime::Handle::current(),
289            CancellationToken::new(),
290        );
291        let caller = tokio::spawn({
292            let manager = manager.clone();
293            async move { manager.register(endpoint_spec()).await }
294        });
295
296        registered.notified().await;
297        caller.abort();
298        let _ = caller.await;
299        wait_for_endpoint_count(discovery.as_ref(), 1).await;
300        release.notify_one();
301        wait_for_endpoint_count(discovery.as_ref(), 0).await;
302    }
303
304    #[tokio::test]
305    async fn endpoint_remains_registered_until_last_lease_drops() {
306        let discovery: Arc<dyn Discovery> =
307            Arc::new(MockDiscovery::new(Some(7), SharedMockRegistry::new()));
308        let manager = EndpointRegistrationManager::new(
309            discovery.clone(),
310            tokio::runtime::Handle::current(),
311            CancellationToken::new(),
312        );
313        let first = manager.register(endpoint_spec()).await.unwrap();
314        let second = manager.register(endpoint_spec()).await.unwrap();
315        wait_for_endpoint_count(discovery.as_ref(), 1).await;
316
317        drop(first);
318        tokio::task::yield_now().await;
319        wait_for_endpoint_count(discovery.as_ref(), 1).await;
320        drop(second);
321        wait_for_endpoint_count(discovery.as_ref(), 0).await;
322    }
323}
dynamo_runtime/discovery/registration.rs

dynamo_runtime/discovery/
registration.rs
