kvbm_engine/leader/velo/
service.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use ::velo::{Handler, Messenger};
5use anyhow::Result;
6use bytes::Bytes;
7use dashmap::DashMap;
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::Arc;
11
12use crate::leader::session::{
13    OnboardMessage, OnboardSessionTx, SessionId, SessionMessage, SessionMessageTx,
14    dispatch_onboard_message, dispatch_session_message,
15};
16use kvbm_physical::manager::SerializedLayout;
17
18/// Type alias for async export metadata callback.
19/// Returns a boxed future that resolves to `Vec<SerializedLayout>`.
20pub type ExportMetadataCallback = Arc<
21    dyn Fn() -> Pin<Box<dyn Future<Output = Result<Vec<SerializedLayout>>> + Send>> + Send + Sync,
22>;
23
24/// Velo leader service for handling distributed onboarding messages.
25///
26/// This service registers handlers for:
27/// 1. OnboardMessage: Standard find_matches flow (initiator → responder)
28/// 2. SessionMessage: Unified session protocol
29/// 3. Export metadata RPC: Returns worker layout metadata for RDMA
30pub struct VeloLeaderService {
31    messenger: Arc<Messenger>,
32    sessions: Arc<DashMap<SessionId, OnboardSessionTx>>,
33    /// Callback to spawn new responder sessions.
34    /// Takes the CreateSession message and creates a new responder task.
35    spawn_responder: Option<Arc<dyn Fn(OnboardMessage) -> Result<()> + Send + Sync>>,
36
37    // Unified session protocol
38    /// Map of unified session receivers.
39    session_sessions: Option<Arc<DashMap<SessionId, SessionMessageTx>>>,
40
41    // RDMA metadata export
42    /// Callback to export worker metadata for RDMA transfers.
43    export_metadata: Option<ExportMetadataCallback>,
44}
45
46impl VeloLeaderService {
47    pub fn new(
48        messenger: Arc<Messenger>,
49        sessions: Arc<DashMap<SessionId, OnboardSessionTx>>,
50    ) -> Self {
51        Self {
52            messenger,
53            sessions,
54            spawn_responder: None,
55            session_sessions: None,
56            export_metadata: None,
57        }
58    }
59
60    /// Set the callback for spawning responder sessions.
61    pub fn with_spawn_responder<F>(mut self, f: F) -> Self
62    where
63        F: Fn(OnboardMessage) -> Result<()> + Send + Sync + 'static,
64    {
65        self.spawn_responder = Some(Arc::new(f));
66        self
67    }
68
69    /// Set the unified session sessions map.
70    pub fn with_session_sessions(
71        mut self,
72        sessions: Arc<DashMap<SessionId, SessionMessageTx>>,
73    ) -> Self {
74        self.session_sessions = Some(sessions);
75        self
76    }
77
78    /// Set the callback for exporting worker metadata (RDMA).
79    ///
80    /// This callback is invoked when a remote leader requests metadata
81    /// to enable RDMA transfers. The callback should return `Vec<SerializedLayout>`
82    /// containing metadata from all workers.
83    pub fn with_export_metadata(mut self, callback: ExportMetadataCallback) -> Self {
84        self.export_metadata = Some(callback);
85        self
86    }
87
88    /// Register all Velo handlers for leader-to-leader communication.
89    pub fn register_handlers(self) -> Result<()> {
90        self.register_onboard_handler()?;
91
92        // Register session handler if unified protocol is configured
93        if self.session_sessions.is_some() {
94            self.register_session_handler()?;
95        }
96
97        // Register export_metadata handler if callback is configured
98        if self.export_metadata.is_some() {
99            self.register_export_metadata_handler()?;
100        }
101
102        Ok(())
103    }
104
105    /// Register the "kvbm.leader.onboard" handler.
106    ///
107    /// This handler is intentionally simple and fast:
108    /// - Deserializes the message
109    /// - If CreateSession and session doesn't exist, spawns responder
110    /// - Dispatches to session channel
111    /// - Returns immediately (< 1ms)
112    fn register_onboard_handler(&self) -> Result<()> {
113        let sessions = self.sessions.clone();
114        let spawn_responder = self.spawn_responder.clone();
115
116        let handler = Handler::am_handler_async("kvbm.leader.onboard", move |ctx| {
117            let sessions = sessions.clone();
118            let spawn_responder = spawn_responder.clone();
119
120            async move {
121                // Fast path: just deserialize and dispatch
122                let message: OnboardMessage = serde_json::from_slice(&ctx.payload)
123                    .map_err(|e| anyhow::anyhow!("failed to deserialize OnboardMessage: {e}"))?;
124
125                let session_id = message.session_id();
126
127                tracing::debug!(
128                    variant = message.variant_name(),
129                    %session_id,
130                    "Received onboard message"
131                );
132
133                // If this is a CreateSession and no session exists, spawn responder
134                if matches!(message, OnboardMessage::CreateSession { .. })
135                    && !sessions.contains_key(&session_id)
136                {
137                    tracing::debug!(%session_id, "Spawning new ResponderSession");
138                    if let Some(ref spawner) = spawn_responder {
139                        spawner(message.clone()).ok(); // Best-effort spawn
140                    }
141                }
142
143                // Dispatch to session channel (will create if needed by spawner above)
144                tracing::debug!(%session_id, "Dispatching message to session");
145                dispatch_onboard_message(&sessions, message).await?;
146
147                Ok(())
148            }
149        })
150        .build();
151
152        self.messenger.register_handler(handler)?;
153
154        Ok(())
155    }
156
157    /// Register the "kvbm.leader.session" handler.
158    ///
159    /// This handler supports the unified session protocol.
160    /// Routes SessionMessages to the appropriate session endpoint.
161    fn register_session_handler(&self) -> Result<()> {
162        let session_sessions = self
163            .session_sessions
164            .clone()
165            .expect("session_sessions required for handler registration");
166
167        let handler = Handler::am_handler_async("kvbm.leader.session", move |ctx| {
168            let session_sessions = session_sessions.clone();
169
170            async move {
171                let message: SessionMessage = serde_json::from_slice(&ctx.payload)
172                    .map_err(|e| anyhow::anyhow!("failed to deserialize SessionMessage: {e}"))?;
173
174                let session_id = message.session_id();
175
176                tracing::debug!(
177                    variant = message.variant_name(),
178                    %session_id,
179                    "Received session message"
180                );
181
182                // Dispatch to session endpoint
183                dispatch_session_message(&session_sessions, message).await?;
184
185                Ok(())
186            }
187        })
188        .build();
189
190        self.messenger.register_handler(handler)?;
191
192        Ok(())
193    }
194
195    /// Register the "kvbm.leader.export_metadata" handler.
196    ///
197    /// This handler returns `Vec<SerializedLayout>` containing metadata from all workers.
198    /// Used by remote leaders to enable RDMA transfers.
199    fn register_export_metadata_handler(&self) -> Result<()> {
200        let export_metadata = self
201            .export_metadata
202            .clone()
203            .expect("export_metadata callback required for handler registration");
204
205        let handler = Handler::unary_handler_async("kvbm.leader.export_metadata", move |_ctx| {
206            let export_metadata = export_metadata.clone();
207
208            async move {
209                tracing::debug!("Received export_metadata request");
210
211                // Call the async callback to get metadata from all workers
212                let metadata_vec = export_metadata().await?;
213
214                // Serialize the Vec<SerializedLayout> for transport
215                let serialized = serde_json::to_vec(&metadata_vec)?;
216
217                tracing::debug!(
218                    count = metadata_vec.len(),
219                    "Returning worker metadata entries"
220                );
221
222                Ok(Some(Bytes::from(serialized)))
223            }
224        })
225        .build();
226
227        self.messenger.register_handler(handler)?;
228
229        Ok(())
230    }
231}
kvbm_engine/leader/velo/service.rs

kvbm_engine/leader/velo/
service.rs
