agave_scheduling_utils/handshake/
shared.rs

1use {
2    agave_scheduler_bindings::{
3        PackToWorkerMessage, ProgressMessage, TpuToPackMessage, WorkerToPackMessage,
4    },
5    rts_alloc::Allocator,
6    thiserror::Error,
7};
8
9pub(crate) type RtsAllocError = rts_alloc::error::Error;
10pub(crate) type ShaqError = shaq::error::Error;
11
12pub const MAX_WORKERS: usize = 64;
13
14/// Protocol version.
15pub(crate) const VERSION: u64 = 4;
16pub(crate) const LOGON_SUCCESS: u8 = 0x01;
17pub(crate) const LOGON_FAILURE: u8 = 0x02;
18pub(crate) const MAX_ALLOCATOR_HANDLES: usize = 128;
19pub(crate) const GLOBAL_ALLOCATORS: usize = 1;
20
21/// The logon message sent by the client to the server.
22#[derive(Debug, Default, Clone, Copy)]
23#[repr(C)]
24pub struct ClientLogon {
25    /// The number of Agave worker threads that will be spawned to handle packing requests.
26    pub worker_count: usize,
27    /// The minimum allocator file size in bytes, this is shared by all allocator handles.
28    pub allocator_size: usize,
29    /// The number of [`rts_alloc::Allocator`] handles the external process is requesting.
30    pub allocator_handles: usize,
31    /// The minimum capacity of the `tpu_to_pack` queue in messages.
32    pub tpu_to_pack_capacity: usize,
33    /// The minimum capacity of the `progress_tracker` queue in messages.
34    pub progress_tracker_capacity: usize,
35    /// The minimum capacity of the `pack_to_worker` queue in messages.
36    pub pack_to_worker_capacity: usize,
37    /// The minimum capacity of the `worker_to_pack` queue in messages.
38    pub worker_to_pack_capacity: usize,
39    /// Flags that control the behavior of the new scheduling session.
40    pub flags: u16,
41    // NB: If adding more fields please ensure:
42    // - The fields are zeroable.
43    // - If possible the fields are backwards compatible:
44    //   - Added to the end of the struct.
45    //   - 0 bytes is valid default (older clients will not have the field and thus send zeroes).
46    // - If not backwards compatible, increment the version counter.
47}
48
49impl ClientLogon {
50    pub fn try_from_bytes(buffer: &[u8]) -> Option<Self> {
51        if buffer.len() != core::mem::size_of::<Self>() {
52            return None;
53        }
54
55        // SAFETY:
56        // - buffer is correctly sized, initialized and readable.
57        // - `Self` is valid for any byte pattern
58        Some(unsafe { core::ptr::read_unaligned(buffer.as_ptr().cast()) })
59    }
60}
61
62pub mod logon_flags {}
63
64/// The complete initialized scheduling session.
65pub struct ClientSession {
66    pub allocators: Vec<Allocator>,
67    pub tpu_to_pack: shaq::spsc::Consumer<TpuToPackMessage>,
68    pub progress_tracker: shaq::spsc::Consumer<ProgressMessage>,
69    pub workers: Vec<ClientWorkerSession>,
70}
71
72/// A per worker scheduling session.
73pub struct ClientWorkerSession {
74    pub pack_to_worker: shaq::spsc::Producer<PackToWorkerMessage>,
75    pub worker_to_pack: shaq::spsc::Consumer<WorkerToPackMessage>,
76}
77
78/// Potential errors that can occur during the client's side of the handshake.
79#[derive(Debug, Error)]
80pub enum ClientHandshakeError {
81    #[error("Io; err={0}")]
82    Io(#[from] std::io::Error),
83    #[error("Timed out")]
84    TimedOut,
85    #[error("Protocol violation")]
86    ProtocolViolation,
87    #[error("Rejected; reason={0}")]
88    Rejected(String),
89    #[error("Rts alloc; err={0}")]
90    RtsAlloc(#[from] RtsAllocError),
91    #[error("Shaq; err={0}")]
92    Shaq(#[from] ShaqError),
93}
94
95/// An initialized scheduling session.
96pub struct AgaveSession {
97    pub flags: u16,
98    pub tpu_to_pack: AgaveTpuToPackSession,
99    pub progress_tracker: shaq::spsc::Producer<ProgressMessage>,
100    pub workers: Vec<AgaveWorkerSession>,
101}
102
103/// Shared memory objects for the tpu to pack worker.
104pub struct AgaveTpuToPackSession {
105    pub allocator: Allocator,
106    pub producer: shaq::spsc::Producer<TpuToPackMessage>,
107}
108
109/// Shared memory objects for a single banking worker.
110pub struct AgaveWorkerSession {
111    pub allocator: Allocator,
112    pub pack_to_worker: shaq::spsc::Consumer<PackToWorkerMessage>,
113    pub worker_to_pack: shaq::spsc::Producer<WorkerToPackMessage>,
114}
115
116/// Potential errors that can occur during the Agave side of the handshake.
117///
118/// # Note
119///
120/// These errors are stringified (up to 256 bytes then truncated) and sent to the client.
121#[derive(Debug, Error)]
122pub enum AgaveHandshakeError {
123    #[error("Io; err={0}")]
124    Io(#[from] std::io::Error),
125    #[error("Timeout")]
126    Timeout,
127    #[error("Close during handshake")]
128    EofDuringHandshake,
129    #[error("Version; server={server}; client={client}")]
130    Version { server: u64, client: u64 },
131    #[error("Worker count; count={0}")]
132    WorkerCount(usize),
133    #[error("Allocator handles; count={0}")]
134    AllocatorHandles(usize),
135    #[error("Rts alloc; err={0:?}")]
136    RtsAlloc(#[from] RtsAllocError),
137    #[error("Shaq; err={0:?}")]
138    Shaq(#[from] ShaqError),
139}
agave_scheduling_utils/handshake/shared.rs

agave_scheduling_utils/handshake/
shared.rs
