Skip to main content

velo_ext/
discovery.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Peer and service discovery extension traits.
5//!
6//! Out-of-tree backends (etcd, Consul, custom k8s, …) implement these traits
7//! against `velo-ext` and integrate with the Velo runtime without pulling
8//! the runtime crate as a build dep.
9
10use std::pin::Pin;
11
12use anyhow::Result;
13use futures::Stream;
14use futures::future::BoxFuture;
15
16use crate::id::{InstanceId, PeerInfo, WorkerId};
17
18// ---------------------------------------------------------------------------
19// Peer discovery
20// ---------------------------------------------------------------------------
21
22/// Abstraction over peer discovery mechanisms.
23///
24/// Higher-level crates implement this trait to integrate with concrete discovery
25/// backends (e.g., etcd, consul) without pulling those dependencies into the
26/// messenger layer.
27pub trait PeerDiscovery: Send + Sync {
28    /// Discover a peer by its worker ID.
29    fn discover_by_worker_id(&self, worker_id: WorkerId) -> BoxFuture<'_, Result<PeerInfo>>;
30
31    /// Discover a peer by its instance ID.
32    fn discover_by_instance_id(&self, instance_id: InstanceId) -> BoxFuture<'_, Result<PeerInfo>>;
33}
34
35/// RAII guard that unregisters a peer when dropped or explicitly unregistered.
36///
37/// Backends implement this trait to provide async cleanup via `unregister()`.
38/// Callers that can await should prefer `unregister()` over relying on `Drop`.
39pub trait PeerRegistrationGuard: Send {
40    /// Explicitly unregister and clean up resources.
41    fn unregister(&mut self) -> BoxFuture<'_, Result<()>>;
42}
43
44// ---------------------------------------------------------------------------
45// Service discovery
46// ---------------------------------------------------------------------------
47
48/// Event emitted by a service instance watch stream.
49#[derive(Clone, Debug, PartialEq, Eq)]
50pub enum ServiceEvent {
51    /// Initial snapshot of all known instances for the service.
52    Initial(Vec<InstanceId>),
53    /// A new instance registered for the service.
54    Added(InstanceId),
55    /// An instance was removed from the service.
56    Removed(InstanceId),
57    /// The watch stream lost its connection to the backend.
58    ///
59    /// Callers should treat the current instance set as stale and re-establish
60    /// the watch via [`ServiceDiscovery::watch_instances`]. This is the last
61    /// event the stream will emit before ending.
62    Disconnected,
63}
64
65/// Abstraction over service discovery mechanisms.
66///
67/// Maps named services to sets of [`InstanceId`]s. For example, finding all
68/// instances that expose a "rhino-router" service. Clients should treat results
69/// as best-effort — returned instance IDs may refer to instances that have
70/// already departed. The caller is responsible for handling failures when
71/// communicating with discovered instances.
72pub trait ServiceDiscovery: Send + Sync {
73    /// List all service names that have at least one registered instance.
74    fn list_services(&self) -> BoxFuture<'_, Result<Vec<String>>>;
75
76    /// Get all instances currently registered for a service.
77    fn get_instances(&self, service_name: &str) -> BoxFuture<'_, Result<Vec<InstanceId>>>;
78
79    /// Watch for changes to instances registered for a service.
80    ///
81    /// The stream emits an [`ServiceEvent::Initial`] event with the current set
82    /// of instances, followed by [`ServiceEvent::Added`] / [`ServiceEvent::Removed`]
83    /// events as instances come and go.
84    fn watch_instances(
85        &self,
86        service_name: &str,
87    ) -> BoxFuture<'_, Result<Pin<Box<dyn Stream<Item = ServiceEvent> + Send>>>>;
88}
89
90/// RAII guard that unregisters a service instance when dropped or explicitly unregistered.
91///
92/// Backends implement this trait to provide async cleanup via `unregister()`.
93/// Callers that can await should prefer `unregister()` over relying on `Drop`.
94pub trait ServiceRegistrationGuard: Send {
95    /// Explicitly unregister and clean up resources.
96    fn unregister(&mut self) -> BoxFuture<'_, Result<()>>;
97}