spark_storage/snapshot_swap/
wire.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3// ───────────────────────────── control protocol ─────────────────────────────
4//
5// The TCP control channel between a paging client and the peer. It rides the
6// SAME stream the peer used for the RDMA handshake (which today just idles).
7// v2-only: EVERY client's first bytes are
8// `[u64 PAGING_MAGIC_V2][u8 kind][u64 arena_bytes][u64 blob_bytes]`.
9// `blob_bytes == 0` selects the RAW one-sided mode (per-connection arena,
10// client-owned allocator, no residency — the legacy data plane,
11// now selected explicitly); anything else is a paging arena. The retired v1
12// magic and the bare-`total_bytes` legacy escape are affirmatively rejected.
13//
14// After the shared rail handshake, the loop is: client sends [op][key], peer
15// replies [status] (+ [offset] for ALLOC/GET-hit). Data still moves one-sided
16// over RDMA into/out of `slot_offset(slot)`; only tiny control messages cross
17// TCP. Peer and client halves deliberately share this ONE module so the wire
18// format (byte-frozen, golden-pinned — it is what the fleet peer binary
19// speaks) can never drift.
20
21use std::io::{Read, Write};
22
23use anyhow::{Context, Result, bail};
24
25use super::{Residency, SlotArena, SwapStore};
26
27/// The RETIRED v1 first-u64 ("PAGE" + 1). Recognized ONLY to reject it with a
28/// dedicated diagnostic — the v1 dialect was deleted (no kind byte) and the
29/// bare-`total_bytes` legacy escape, so a stale binary fails legibly at
30/// handshake instead of being reinterpreted.
31const PAGING_MAGIC_V1_RETIRED: u64 = 0x5041_4745_0000_0001;
32
33/// The ONLY accepted first u64 on the RW paging port ("PAGE" + 2, v2-only
34///): after the magic comes a `[u8 kind]` byte so ONE peer serves
35/// a registry of per-(kind, shape) arenas, then `[u64 arena_bytes]`
36/// `[u64 blob_bytes]` — `blob_bytes == 0` selects the RAW one-sided mode.
37pub const PAGING_MAGIC_V2: u64 = 0x5041_4745_0000_0002;
38
39/// The tier a paging arena serves. Only the RW paging kinds (SSM, KV-as-paging)
40/// ride the `CacheServerParams` single-base+rkey reply; the read-only tiers
41/// (experts/weights/lora) speak a different manifest+VerbsServerParams dialect
42/// and are NOT accepted on this handshake (rejected in `parse_paging_header`).
43#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
44pub struct PagingKind(pub u8);
45impl PagingKind {
46    pub const SSM: PagingKind = PagingKind(0);
47    pub const KV: PagingKind = PagingKind(1);
48    /// Whether this kind is servable on the RW paging (CacheServerParams) path.
49    pub fn is_paging_rw(self) -> bool {
50        self.0 <= 1
51    }
52}
53
54/// Parse the paging handshake header after the caller has read the first u64.
55/// v2-only: the first u64 MUST be [`PAGING_MAGIC_V2`], then
56/// `[u8 kind][u64 arena_bytes][u64 blob_bytes]`. `blob_bytes == 0` = the RAW
57/// one-sided mode (per-connection arena, client-owned allocator — the caller
58/// routes it OFF the paging registry). Rejects unsupported kinds (≥2), the
59/// retired v1 magic (dedicated diagnostic so a stale binary fails legibly),
60/// and any other first u64 (e.g. a bare legacy `total_bytes`).
61pub fn parse_paging_header<R: Read>(first: u64, r: &mut R) -> Result<(PagingKind, u64, u64)> {
62    if first == PAGING_MAGIC_V1_RETIRED {
63        bail!(
64            "paging: v1 client no longer supported (magic {first:#x}); rebuild the client — \
65             every connect now sends [u64 PAGING_MAGIC_V2][u8 kind][u64 arena_bytes][u64 blob_bytes]"
66        );
67    }
68    if first != PAGING_MAGIC_V2 {
69        bail!(
70            "paging: first u64 {first:#x} is not PAGING_MAGIC_V2 (the bare legacy total_bytes \
71             handshake was retired; RAW one-sided clients send the v2 header with \
72             blob_bytes == 0)"
73        );
74    }
75    let mut kb = [0u8; 1];
76    r.read_exact(&mut kb).context("read paging kind")?;
77    let kind = PagingKind(kb[0]);
78    if !kind.is_paging_rw() {
79        bail!(
80            "paging: unsupported kind {} (only SSM/KV ride this handshake)",
81            kb[0]
82        );
83    }
84    let mut b8 = [0u8; 8];
85    r.read_exact(&mut b8).context("read paging arena_bytes")?;
86    let arena_bytes = u64::from_le_bytes(b8);
87    r.read_exact(&mut b8).context("read paging blob_bytes")?;
88    let blob_bytes = u64::from_le_bytes(b8);
89    Ok((kind, arena_bytes, blob_bytes))
90}
91
92/// Encode the CLIENT half of the v2 paging handshake header — what EVERY
93/// paging client sends first: KV paging, SSM paging (`connect_paging`), and
94/// both RAW one-sided modes (via `blob_bytes == 0`):
95/// `[u64 PAGING_MAGIC_V2 LE][u8 kind][u64 arena_bytes LE][u64 blob_bytes LE]`
96/// — 25 bytes, followed by the unchanged `[u8 n_rails]` RailSet exchange.
97/// Lives in this ONE shared module (beside `parse_paging_header`, its peer
98/// half) so writer and reader can never drift; byte-frozen and golden-pinned
99/// in `wire_tests.rs`.
100pub fn encode_paging_v2_header(kind: PagingKind, arena_bytes: u64, blob_bytes: u64) -> [u8; 25] {
101    let mut w = [0u8; 25];
102    w[0..8].copy_from_slice(&PAGING_MAGIC_V2.to_le_bytes());
103    w[8] = kind.0;
104    w[9..17].copy_from_slice(&arena_bytes.to_le_bytes());
105    w[17..25].copy_from_slice(&blob_bytes.to_le_bytes());
106    w
107}
108
109/// Round-robin stripe a `blob_bytes` transfer into `chunk_bytes` chunks across
110/// `n_rails`, returning per-rail lists of `(offset, len)`. The offset is the
111/// chunk's position in BOTH the (single, contiguous) staging buffer and the peer
112/// arena slot — so whichever rail fetches chunk j, it lands at its true offset
113/// and one memcpy reassembles the blob (the verified inc-6 reassembly fix). The
114/// tail chunk carries the short remainder, never `chunk_bytes`.
115pub fn stripe_plan(
116    blob_bytes: usize,
117    chunk_bytes: usize,
118    n_rails: usize,
119) -> Vec<Vec<(usize, usize)>> {
120    let n = n_rails.max(1);
121    let cb = chunk_bytes.max(1);
122    let mut rails: Vec<Vec<(usize, usize)>> = vec![Vec::new(); n];
123    let mut off = 0usize;
124    let mut j = 0usize;
125    while off < blob_bytes {
126        let len = cb.min(blob_bytes - off);
127        rails[j % n].push((off, len));
128        off += len;
129        j += 1;
130    }
131    rails
132}
133
134/// Chunk size for the striped snapshot pipeline (ATLAS_SSM_CHUNK_BYTES, default
135/// 1 MiB) and pipeline depth (ATLAS_SSM_PIPELINE_DEPTH, default 16, clamped
136/// 1..=128, mirroring the KV backend).
137pub fn staging_chunk_bytes() -> usize {
138    std::env::var("ATLAS_SSM_CHUNK_BYTES")
139        .ok()
140        .and_then(|s| s.parse::<usize>().ok())
141        .filter(|&v| v >= 4096)
142        .unwrap_or(1024 * 1024)
143}
144pub fn staging_depth() -> usize {
145    std::env::var("ATLAS_SSM_PIPELINE_DEPTH")
146        .ok()
147        .and_then(|s| s.parse::<usize>().ok())
148        .unwrap_or(16)
149        .clamp(1, 128)
150}
151
152pub const OP_BYE: u8 = 0;
153pub const OP_ALLOC: u8 = 1;
154pub const OP_COMMIT: u8 = 2;
155pub const OP_GET: u8 = 3;
156pub const OP_REMOVE: u8 = 4;
157
158pub const ST_OK: u8 = 0;
159pub const ST_MISS: u8 = 1;
160pub const ST_ERR: u8 = 2;
161
162/// Result of one control request, ready to serialize.
163#[derive(Debug, PartialEq, Eq)]
164pub enum PagingReply {
165    /// ST_OK with a following u64 arena offset (ALLOC and GET-hit).
166    Located(u64),
167    /// ST_OK with no payload (COMMIT, REMOVE).
168    Ok,
169    /// ST_MISS — unknown key (GET).
170    Miss,
171    /// ST_ERR — operation failed (e.g. arena exhausted by reservations).
172    Err,
173    /// Client asked to close.
174    Bye,
175}
176
177/// Execute one control op against the residency and return the reply. Pure over
178/// the (already unit-tested) `Residency`, so the protocol is testable
179/// without a socket or RDMA.
180pub fn dispatch<A: SlotArena, S: SwapStore>(
181    res: &mut Residency<A, S>,
182    op: u8,
183    key: u64,
184) -> PagingReply {
185    match op {
186        OP_BYE => PagingReply::Bye,
187        OP_ALLOC => match res.alloc(key) {
188            Ok(slot) => PagingReply::Located(res.slot_offset(slot)),
189            Err(e) => {
190                tracing::warn!("paging ALLOC {key:#x} failed: {e:#}");
191                PagingReply::Err
192            }
193        },
194        OP_COMMIT => match res.commit(key) {
195            Ok(()) => PagingReply::Ok,
196            Err(e) => {
197                tracing::warn!("paging COMMIT {key:#x} failed: {e:#}");
198                PagingReply::Err
199            }
200        },
201        OP_GET => match res.locate(key) {
202            Ok(Some(slot)) => PagingReply::Located(res.slot_offset(slot)),
203            Ok(None) => PagingReply::Miss,
204            Err(e) => {
205                tracing::warn!("paging GET {key:#x} failed: {e:#}");
206                PagingReply::Err
207            }
208        },
209        OP_REMOVE => {
210            res.remove(key);
211            PagingReply::Ok
212        }
213        other => {
214            tracing::warn!("paging: unknown op {other}");
215            PagingReply::Err
216        }
217    }
218}
219
220fn write_reply<W: Write>(w: &mut W, reply: &PagingReply) -> Result<()> {
221    match reply {
222        PagingReply::Located(off) => {
223            w.write_all(&[ST_OK])?;
224            w.write_all(&off.to_le_bytes())?;
225        }
226        PagingReply::Ok => w.write_all(&[ST_OK])?,
227        PagingReply::Miss => w.write_all(&[ST_MISS])?,
228        PagingReply::Err => w.write_all(&[ST_ERR])?,
229        PagingReply::Bye => {}
230    }
231    w.flush()?;
232    Ok(())
233}
234
235/// The peer-side control loop: read `[op][u64 key]` requests, dispatch against
236/// `res`, write replies, until BYE or hangup. Generic over the stream so it runs
237/// against a real `TcpStream` in the peer and a fake duplex in tests.
238/// One control op with connection-scoped read-pin lifecycle. Releases this
239/// connection's previous GET read-pin — its RDMA READ has necessarily drained,
240/// because the client is synchronous and only sends its NEXT op after the read
241/// completes — then dispatches, and pins a fresh GET hit so a concurrent ALLOC
242/// on another connection cannot evict the slot mid-read. `pinned` threads the
243/// connection's currently-pinned key across calls. Needs NO new opcode and no
244/// client change (auto-release on next op / disconnect) → wire-compatible with
245/// the frozen control protocol.
246fn handle_paging_op<A: SlotArena, S: SwapStore>(
247    res: &mut Residency<A, S>,
248    op: u8,
249    key: u64,
250    pinned: &mut Option<u64>,
251) -> PagingReply {
252    if let Some(prev) = pinned.take() {
253        res.unpin_read(prev);
254    }
255    let reply = dispatch(res, op, key);
256    if op == OP_GET && matches!(reply, PagingReply::Located(_)) {
257        res.pin_read(key);
258        *pinned = Some(key);
259    }
260    reply
261}
262
263pub fn run_paging_loop<T: Read + Write, A: SlotArena, S: SwapStore>(
264    stream: &mut T,
265    res: &mut Residency<A, S>,
266) -> Result<()> {
267    let mut pinned: Option<u64> = None;
268    loop {
269        let mut op = [0u8; 1];
270        if stream.read_exact(&mut op).is_err() {
271            break; // client hung up
272        }
273        if op[0] == OP_BYE {
274            break;
275        }
276        let mut kb = [0u8; 8];
277        stream.read_exact(&mut kb).context("read paging key")?;
278        let key = u64::from_le_bytes(kb);
279        let reply = handle_paging_op(res, op[0], key, &mut pinned);
280        write_reply(stream, &reply)?;
281    }
282    // Release this connection's outstanding read-pin on hangup / BYE.
283    if let Some(pk) = pinned {
284        res.unpin_read(pk);
285    }
286    Ok(())
287}
288
289// ─────────────────────── client-side protocol helpers ───────────────────────
290//
291// The CLIENT half of the control channel, sharing the wire format above so peer
292// and client can never drift. Each sends `[op][u64 key]` and reads the reply.
293// The RDMA data-plane WRITE/READ (client-side) happens between `client_alloc`
294// and `client_commit` (PUT) or after `client_get` (GET) — see RdmaSnapshotArena.
295
296fn send_req<T: Write>(s: &mut T, op: u8, key: u64) -> Result<()> {
297    let mut buf = [0u8; 9];
298    buf[0] = op;
299    buf[1..].copy_from_slice(&key.to_le_bytes());
300    s.write_all(&buf)?;
301    s.flush()?;
302    Ok(())
303}
304
305fn read_status<T: Read>(s: &mut T) -> Result<u8> {
306    let mut st = [0u8; 1];
307    s.read_exact(&mut st).context("read paging status")?;
308    Ok(st[0])
309}
310
311fn read_offset<T: Read>(s: &mut T) -> Result<u64> {
312    let mut b = [0u8; 8];
313    s.read_exact(&mut b).context("read paging offset")?;
314    Ok(u64::from_le_bytes(b))
315}
316
317/// PUT step 1: reserve a slot for `key`; returns the arena offset to RDMA-WRITE.
318pub fn client_alloc<T: Read + Write>(s: &mut T, key: u64) -> Result<u64> {
319    send_req(s, OP_ALLOC, key)?;
320    match read_status(s)? {
321        ST_OK => read_offset(s),
322        st => bail!("paging ALLOC {key:#x} refused (status {st})"),
323    }
324}
325
326/// PUT step 2: the RDMA-WRITE has drained; mark `key` resident.
327pub fn client_commit<T: Read + Write>(s: &mut T, key: u64) -> Result<()> {
328    send_req(s, OP_COMMIT, key)?;
329    match read_status(s)? {
330        ST_OK => Ok(()),
331        st => bail!("paging COMMIT {key:#x} failed (status {st})"),
332    }
333}
334
335/// GET: `Some(offset)` to RDMA-READ, or `None` if the peer has no such key.
336pub fn client_get<T: Read + Write>(s: &mut T, key: u64) -> Result<Option<u64>> {
337    send_req(s, OP_GET, key)?;
338    match read_status(s)? {
339        ST_OK => Ok(Some(read_offset(s)?)),
340        ST_MISS => Ok(None),
341        st => bail!("paging GET {key:#x} error (status {st})"),
342    }
343}
344
345/// Drop `key` from the peer cache.
346pub fn client_remove<T: Read + Write>(s: &mut T, key: u64) -> Result<()> {
347    send_req(s, OP_REMOVE, key)?;
348    match read_status(s)? {
349        ST_OK => Ok(()),
350        st => bail!("paging REMOVE {key:#x} failed (status {st})"),
351    }
352}
353
354/// Politely tell the peer to close the paging loop.
355pub fn client_bye<T: Write>(s: &mut T) -> Result<()> {
356    send_req(s, OP_BYE, 0)
357}
358
359/// Shared variant of [`run_paging_loop`]: many connection threads drive ONE
360/// process-global residency, locking it per request. This is what makes the
361/// peer a SHARED warm cache — a snapshot PUT by one client is GET-able by
362/// another (same namespace). The lock is held only for the (fast) map op + any
363/// spill/fault byte move, never across a TCP read.
364pub fn run_paging_loop_shared<T: Read + Write, A: SlotArena, S: SwapStore>(
365    stream: &mut T,
366    res: &std::sync::Mutex<Residency<A, S>>,
367) -> Result<()> {
368    // Per-connection read-pin (see `handle_paging_op`): a GET hit is pinned OUT
369    // of the LRU under the same lock as the dispatch, so a concurrent ALLOC on
370    // another connection can't evict the slot while THIS client RDMA-reads it.
371    // Released on the connection's next op (its read has drained) or disconnect.
372    let mut pinned: Option<u64> = None;
373    loop {
374        let mut op = [0u8; 1];
375        if stream.read_exact(&mut op).is_err() {
376            break;
377        }
378        if op[0] == OP_BYE {
379            break;
380        }
381        let mut kb = [0u8; 8];
382        stream.read_exact(&mut kb).context("read paging key")?;
383        let key = u64::from_le_bytes(kb);
384        let reply = {
385            let mut g = res.lock().expect("shared residency mutex poisoned");
386            handle_paging_op(&mut g, op[0], key, &mut pinned)
387        };
388        write_reply(stream, &reply)?;
389    }
390    if let Some(pk) = pinned {
391        res.lock()
392            .expect("shared residency mutex poisoned")
393            .unpin_read(pk);
394    }
395    Ok(())
396}
397
398#[cfg(test)]
399#[path = "wire_tests.rs"]
400mod tests;
spark_storage/snapshot_swap/wire.rs

spark_storage/snapshot_swap/
wire.rs
