kvbm_engine/leader/session/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! # Session Module
5//!
6//! This module provides session management for distributed block transfers.
7//!
8//! ## Core Building Blocks
9//!
10//! Composable building blocks for session management:
11//!
12//! - `BlockHolder<T>`: RAII container for holding blocks during sessions
13//! - `SessionEndpoint`: Point-to-point session primitive with state machine
14//! - `SessionHandle`: Unified handle for controlling remote sessions
15//! - `SessionMessage`: Unified message protocol with bidirectional control
16//! - `SessionPhase`, `ControlRole`, `AttachmentState`: State machine types
17//!
18//! ## Session Implementations
19//!
20//! - `ServerSession`: Server-side session (merges former EndpointSession + ControllableSession)
21//! - `InitiatorSession`: Multi-peer search orchestrator (OnboardMessage)
22//! - `ResponderSession`: Responds to search requests (OnboardMessage)
23
24// Core session building blocks
25mod blocks;
26mod endpoint;
27mod handle;
28mod server_session;
29mod staging;
30mod state;
31
32// Session implementations
33mod initiator;
34mod messages;
35mod responder;
36pub mod transport;
37
38// =============================================================================
39// Core Building Blocks
40// =============================================================================
41
42/// RAII container for holding blocks during sessions.
43pub use blocks::BlockHolder;
44
45/// Point-to-point session endpoint with state machine.
46pub use endpoint::{SessionEndpoint, SessionMessageTx, session_message_channel};
47
48/// Server-side session (unified replacement for EndpointSession + ControllableSession).
49pub use server_session::{
50    ServerSession, ServerSessionCommand, ServerSessionHandle, ServerSessionOptions,
51    create_server_session,
52};
53
54// Backwards-compatible aliases for the former EndpointSession types.
55pub use server_session::ServerSessionCommand as EndpointSessionCommand;
56pub use server_session::ServerSessionHandle as EndpointSessionHandle;
57
58/// Unified handle for controlling remote sessions.
59pub use handle::{SessionHandle, SessionHandleStateTx, session_handle_state_channel};
60
61/// State machine types for the unified session model.
62pub use state::{AttachmentState, ControlRole, SessionPhase};
63
64/// Unified session message protocol.
65pub use messages::{BlockInfo, SessionMessage, SessionStateSnapshot};
66
67// =============================================================================
68// Session Implementations
69// =============================================================================
70
71/// Session implementations for initiator and responder patterns.
72pub use initiator::InitiatorSession;
73pub use responder::ResponderSession;
74
75/// Backwards-compatible re-exports (ControllableSessionResult is still used externally).
76pub use server_session::ServerSessionOptions as ControllableSessionOptions;
77
78/// Result of creating a controllable/server session.
79#[derive(Debug, Clone)]
80pub struct ControllableSessionResult {
81    /// The unique session ID.
82    pub session_id: super::SessionId,
83    /// Number of G2 blocks found.
84    pub local_g2_count: usize,
85    /// Number of G3 blocks found.
86    pub local_g3_count: usize,
87}
88
89/// Message types for session communication.
90pub use messages::{BlockMatch, OnboardMessage};
91
92/// Transport types.
93pub use transport::{LocalTransport, MessageTransport, VeloTransport};
94
95use anyhow::Result;
96use dashmap::DashMap;
97use tokio::sync::mpsc;
98
99pub type SessionId = uuid::Uuid;
100pub type OnboardSessionTx = mpsc::Sender<OnboardMessage>;
101
102/// Route an [`OnboardMessage`] to its per-session task channel.
103///
104/// Looks up the session ID in the `DashMap` registry and forwards the message
105/// through the session's mpsc sender. Each session processes messages serially
106/// via its channel, so ordering is preserved per-session.
107pub async fn dispatch_onboard_message(
108    sessions: &DashMap<SessionId, OnboardSessionTx>,
109    message: OnboardMessage,
110) -> Result<()> {
111    let session_id = message.session_id();
112
113    let sender = sessions.get(&session_id).map(|entry| entry.value().clone());
114    if let Some(sender) = sender {
115        sender
116            .send(message)
117            .await
118            .map_err(|e| anyhow::anyhow!("failed to send to session {session_id}: {e}"))?;
119        return Ok(());
120    }
121
122    anyhow::bail!("no session task registered for session {session_id}");
123}
124
125/// Route a unified [`SessionMessage`] to its session task.
126///
127/// All message variants are routed through a single `DashMap<SessionId, SessionMessageTx>`
128/// registry.
129pub async fn dispatch_session_message(
130    sessions: &DashMap<SessionId, SessionMessageTx>,
131    message: SessionMessage,
132) -> Result<()> {
133    let session_id = message.session_id();
134
135    let sender = sessions.get(&session_id).map(|entry| entry.value().clone());
136    if let Some(sender) = sender {
137        sender
138            .send(message)
139            .await
140            .map_err(|e| anyhow::anyhow!("failed to send to session {session_id}: {e}"))?;
141        return Ok(());
142    }
143
144    anyhow::bail!("no session registered for session {session_id}");
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[tokio::test]
152    async fn test_dispatch_onboard_message() {
153        let sessions: DashMap<SessionId, OnboardSessionTx> = DashMap::new();
154        let session_id = SessionId::new_v4();
155        let (tx, mut rx) = mpsc::channel(16);
156        sessions.insert(session_id, tx);
157
158        let msg = OnboardMessage::CloseSession {
159            requester: crate::InstanceId::new_v4(),
160            session_id,
161        };
162
163        dispatch_onboard_message(&sessions, msg).await.unwrap();
164
165        let received = rx.recv().await.unwrap();
166        assert_eq!(received.session_id(), session_id);
167    }
168
169    #[tokio::test]
170    async fn test_dispatch_session_message() {
171        let sessions: DashMap<SessionId, SessionMessageTx> = DashMap::new();
172        let session_id = SessionId::new_v4();
173        let (tx, mut rx) = mpsc::channel(16);
174        sessions.insert(session_id, tx);
175
176        let msg = SessionMessage::Close { session_id };
177
178        dispatch_session_message(&sessions, msg).await.unwrap();
179
180        let received = rx.recv().await.unwrap();
181        assert_eq!(received.session_id(), session_id);
182    }
183
184    #[tokio::test]
185    async fn test_dispatch_missing_onboard_session() {
186        let sessions: DashMap<SessionId, OnboardSessionTx> = DashMap::new();
187        let session_id = SessionId::new_v4();
188
189        let msg = OnboardMessage::CloseSession {
190            requester: crate::InstanceId::new_v4(),
191            session_id,
192        };
193
194        let result = dispatch_onboard_message(&sessions, msg).await;
195        assert!(result.is_err());
196    }
197
198    #[tokio::test]
199    async fn test_dispatch_missing_session_message() {
200        let sessions: DashMap<SessionId, SessionMessageTx> = DashMap::new();
201        let session_id = SessionId::new_v4();
202
203        let msg = SessionMessage::Close { session_id };
204
205        let result = dispatch_session_message(&sessions, msg).await;
206        assert!(result.is_err());
207    }
208}
kvbm_engine/leader/session/mod.rs

kvbm_engine/leader/session/
mod.rs
