metrale_storage/snapshot_swap/
wire.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The paging control protocol on the TCP stream between a paging client
4//! and the peer: the handshake header, the request loop, and the client calls.
5//!
6//! Owner: storage, SSM snapshot tier.
7//! Invariants:
8//! - A connection holds at most one read pin, taken on a GET hit. It is released at
9//!   the connection's next request, or when the loop ends on BYE or hangup; a failed
10//!   key read or reply write returns without releasing it.
11//!
12//! Every client first sends `[u64 PAGING_MAGIC_V2][u8 kind][u64 arena_bytes]
13//! [u64 blob_bytes]` on the stream the RDMA handshake then uses. `blob_bytes == 0`
14//! selects raw mode (a private arena, placement by the client, no residency); any
15//! other value a paging arena. After the rail handshake the client sends
16//! `[op][u64 key]` and the peer replies `[status]`, plus a `[u64 offset]` for ALLOC
17//! and a GET hit. Blobs move over one-sided RDMA to and from that offset; only
18//! control messages cross TCP. The peer and client halves live here together, and
19//! `wire_tests.rs` pins the byte layout.
20
21use std::io::{Read, Write};
22
23use anyhow::{Context, Result, bail};
24
25use super::{Residency, SlotArena, SwapStore};
26
27/// 2026-09-25: ASCII "PAGE" + 1, recognised only to be refused with its own error.
28const PAGING_MAGIC_V1_RETIRED: u64 = 0x5041_4745_0000_0001;
29
30/// 2026-09-25: The only accepted first u64 on the paging port, ASCII "PAGE" + 2. The
31/// `[u8 kind]` that follows lets one peer keep an arena per `(kind, blob_bytes)`.
32pub const PAGING_MAGIC_V2: u64 = 0x5041_4745_0000_0002;
33
34/// 2026-09-25: The tier a paging arena serves. Only `SSM` and `KV` are accepted
35/// (`parse_paging_header`); the read-only expert and weight peers use their own
36/// handshake.
37#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
38pub struct PagingKind(pub u8);
39impl PagingKind {
40    pub const SSM: PagingKind = PagingKind(0);
41    pub const KV: PagingKind = PagingKind(1);
42    /// 2026-09-25: `true` for `SSM` and `KV`.
43    pub fn is_paging_rw(self) -> bool {
44        self.0 <= 1
45    }
46}
47
48/// 2026-09-25: Parse the rest of the handshake header, given the first u64 the caller
49/// has read: `[u8 kind][u64 arena_bytes][u64 blob_bytes]`. A first u64 other than
50/// [`PAGING_MAGIC_V2`] (with its own message for the "PAGE" + 1 value), or a kind
51/// other than `SSM`/`KV`, is an error.
52pub fn parse_paging_header<R: Read>(first: u64, r: &mut R) -> Result<(PagingKind, u64, u64)> {
53    if first == PAGING_MAGIC_V1_RETIRED {
54        bail!(
55            "paging: v1 client no longer supported (magic {first:#x}); rebuild the client — \
56             every connect now sends [u64 PAGING_MAGIC_V2][u8 kind][u64 arena_bytes][u64 blob_bytes]"
57        );
58    }
59    if first != PAGING_MAGIC_V2 {
60        bail!(
61            "paging: first u64 {first:#x} is not PAGING_MAGIC_V2 (the bare legacy total_bytes \
62             handshake was retired; RAW one-sided clients send the v2 header with \
63             blob_bytes == 0)"
64        );
65    }
66    let mut kb = [0u8; 1];
67    r.read_exact(&mut kb).context("read paging kind")?;
68    let kind = PagingKind(kb[0]);
69    if !kind.is_paging_rw() {
70        bail!(
71            "paging: unsupported kind {} (only SSM/KV ride this handshake)",
72            kb[0]
73        );
74    }
75    let mut b8 = [0u8; 8];
76    r.read_exact(&mut b8).context("read paging arena_bytes")?;
77    let arena_bytes = u64::from_le_bytes(b8);
78    r.read_exact(&mut b8).context("read paging blob_bytes")?;
79    let blob_bytes = u64::from_le_bytes(b8);
80    Ok((kind, arena_bytes, blob_bytes))
81}
82
83/// 2026-09-25: The 25-byte header every client sends first, in paging and raw mode:
84/// `[u64 PAGING_MAGIC_V2 LE][u8 kind][u64 arena_bytes LE][u64 blob_bytes LE]`. The
85/// `RailSet` exchange follows it.
86pub fn encode_paging_v2_header(kind: PagingKind, arena_bytes: u64, blob_bytes: u64) -> [u8; 25] {
87    let mut w = [0u8; 25];
88    w[0..8].copy_from_slice(&PAGING_MAGIC_V2.to_le_bytes());
89    w[8] = kind.0;
90    w[9..17].copy_from_slice(&arena_bytes.to_le_bytes());
91    w[17..25].copy_from_slice(&blob_bytes.to_le_bytes());
92    w
93}
94
95/// 2026-09-25: Split a `blob_bytes` transfer into `chunk_bytes` chunks dealt
96/// round-robin to `n_rails`, as per-rail lists of `(offset, len)`. The offset is the
97/// chunk's position in both the staging buffer and the peer slot, so one copy
98/// assembles the blob whichever rail moved each chunk. Only the last chunk may be
99/// short. Zero `chunk_bytes` or `n_rails` counts as 1.
100pub fn stripe_plan(
101    blob_bytes: usize,
102    chunk_bytes: usize,
103    n_rails: usize,
104) -> Vec<Vec<(usize, usize)>> {
105    let n = n_rails.max(1);
106    let cb = chunk_bytes.max(1);
107    let mut rails: Vec<Vec<(usize, usize)>> = vec![Vec::new(); n];
108    let mut off = 0usize;
109    let mut j = 0usize;
110    while off < blob_bytes {
111        let len = cb.min(blob_bytes - off);
112        rails[j % n].push((off, len));
113        off += len;
114        j += 1;
115    }
116    rails
117}
118
119/// 2026-09-25: Chunk size of the striped snapshot transfer: `METRALE_SSM_CHUNK_BYTES`,
120/// or 1 MiB when it is unset, unparsable or below 4096.
121pub fn staging_chunk_bytes() -> usize {
122    std::env::var("METRALE_SSM_CHUNK_BYTES")
123        .ok()
124        .and_then(|s| s.parse::<usize>().ok())
125        .filter(|&v| v >= 4096)
126        .unwrap_or(1024 * 1024)
127}
128/// 2026-09-25: Chunks in flight per rail: `METRALE_SSM_PIPELINE_DEPTH`, or 16 when it
129/// is unset or unparsable, clamped to 1..=128.
130pub fn staging_depth() -> usize {
131    std::env::var("METRALE_SSM_PIPELINE_DEPTH")
132        .ok()
133        .and_then(|s| s.parse::<usize>().ok())
134        .unwrap_or(16)
135        .clamp(1, 128)
136}
137
138pub const OP_BYE: u8 = 0;
139pub const OP_ALLOC: u8 = 1;
140pub const OP_COMMIT: u8 = 2;
141pub const OP_GET: u8 = 3;
142pub const OP_REMOVE: u8 = 4;
143
144pub const ST_OK: u8 = 0;
145pub const ST_MISS: u8 = 1;
146pub const ST_ERR: u8 = 2;
147
148/// 2026-09-25: Result of one control request, ready to serialise.
149#[derive(Debug, PartialEq, Eq)]
150pub enum PagingReply {
151    /// 2026-09-25: `ST_OK` and a u64 arena offset (ALLOC, and a GET hit).
152    Located(u64),
153    /// 2026-09-25: `ST_OK` alone (COMMIT, REMOVE).
154    Ok,
155    /// 2026-09-25: `ST_MISS`: GET of an unknown key.
156    Miss,
157    /// 2026-09-25: `ST_ERR`: the operation failed, or the op is unknown.
158    Err,
159    /// 2026-09-25: The client asked to close; nothing is written.
160    Bye,
161}
162
163/// 2026-09-25: Execute one control op against `res` and return the reply; no I/O, so
164/// the protocol is testable without a socket or RDMA.
165pub fn dispatch<A: SlotArena, S: SwapStore>(
166    res: &mut Residency<A, S>,
167    op: u8,
168    key: u64,
169) -> PagingReply {
170    match op {
171        OP_BYE => PagingReply::Bye,
172        OP_ALLOC => match res.alloc(key) {
173            Ok(slot) => PagingReply::Located(res.slot_offset(slot)),
174            Err(e) => {
175                tracing::warn!("paging ALLOC {key:#x} failed: {e:#}");
176                PagingReply::Err
177            }
178        },
179        OP_COMMIT => match res.commit(key) {
180            Ok(()) => PagingReply::Ok,
181            Err(e) => {
182                tracing::warn!("paging COMMIT {key:#x} failed: {e:#}");
183                PagingReply::Err
184            }
185        },
186        OP_GET => match res.locate(key) {
187            Ok(Some(slot)) => PagingReply::Located(res.slot_offset(slot)),
188            Ok(None) => PagingReply::Miss,
189            Err(e) => {
190                tracing::warn!("paging GET {key:#x} failed: {e:#}");
191                PagingReply::Err
192            }
193        },
194        OP_REMOVE => {
195            res.remove(key);
196            PagingReply::Ok
197        }
198        other => {
199            tracing::warn!("paging: unknown op {other}");
200            PagingReply::Err
201        }
202    }
203}
204
205fn write_reply<W: Write>(w: &mut W, reply: &PagingReply) -> Result<()> {
206    match reply {
207        PagingReply::Located(off) => {
208            w.write_all(&[ST_OK])?;
209            w.write_all(&off.to_le_bytes())?;
210        }
211        PagingReply::Ok => w.write_all(&[ST_OK])?,
212        PagingReply::Miss => w.write_all(&[ST_MISS])?,
213        PagingReply::Err => w.write_all(&[ST_ERR])?,
214        PagingReply::Bye => {}
215    }
216    w.flush()?;
217    Ok(())
218}
219
220/// 2026-09-25: One control op with the connection's read pin: release the pin of the
221/// previous GET, dispatch, and pin a GET hit so an ALLOC on another connection cannot
222/// evict the slot while this client reads it. Releasing on the next op is sound only
223/// for clients that finish each READ before sending another op, as
224/// `RdmaSnapshotArena` and `KvPagingBackend` do.
225fn handle_paging_op<A: SlotArena, S: SwapStore>(
226    res: &mut Residency<A, S>,
227    op: u8,
228    key: u64,
229    pinned: &mut Option<u64>,
230) -> PagingReply {
231    if let Some(prev) = pinned.take() {
232        res.unpin_read(prev);
233    }
234    let reply = dispatch(res, op, key);
235    if op == OP_GET && matches!(reply, PagingReply::Located(_)) {
236        res.pin_read(key);
237        *pinned = Some(key);
238    }
239    reply
240}
241
242/// 2026-09-25: The peer's control loop over one residency: read `[op][u64 key]`,
243/// dispatch, reply, until BYE or a failed read of the op byte.
244pub fn run_paging_loop<T: Read + Write, A: SlotArena, S: SwapStore>(
245    stream: &mut T,
246    res: &mut Residency<A, S>,
247) -> Result<()> {
248    let mut pinned: Option<u64> = None;
249    loop {
250        let mut op = [0u8; 1];
251        if stream.read_exact(&mut op).is_err() {
252            break;
253        }
254        if op[0] == OP_BYE {
255            break;
256        }
257        let mut kb = [0u8; 8];
258        stream.read_exact(&mut kb).context("read paging key")?;
259        let key = u64::from_le_bytes(kb);
260        let reply = handle_paging_op(res, op[0], key, &mut pinned);
261        write_reply(stream, &reply)?;
262    }
263    if let Some(pk) = pinned {
264        res.unpin_read(pk);
265    }
266    Ok(())
267}
268
269// 2026-09-25: Client side. Each call sends `[op][u64 key]` and reads the reply. A PUT
270// is `client_alloc`, the RDMA WRITE, then `client_commit`; a GET is `client_get`,
271// then the RDMA READ.
272
273fn send_req<T: Write>(s: &mut T, op: u8, key: u64) -> Result<()> {
274    let mut buf = [0u8; 9];
275    buf[0] = op;
276    buf[1..].copy_from_slice(&key.to_le_bytes());
277    s.write_all(&buf)?;
278    s.flush()?;
279    Ok(())
280}
281
282fn read_status<T: Read>(s: &mut T) -> Result<u8> {
283    let mut st = [0u8; 1];
284    s.read_exact(&mut st).context("read paging status")?;
285    Ok(st[0])
286}
287
288fn read_offset<T: Read>(s: &mut T) -> Result<u64> {
289    let mut b = [0u8; 8];
290    s.read_exact(&mut b).context("read paging offset")?;
291    Ok(u64::from_le_bytes(b))
292}
293
294/// 2026-09-25: Reserve a slot for `key` and return its arena offset; a non-OK status
295/// is an error.
296pub fn client_alloc<T: Read + Write>(s: &mut T, key: u64) -> Result<u64> {
297    send_req(s, OP_ALLOC, key)?;
298    match read_status(s)? {
299        ST_OK => read_offset(s),
300        st => bail!("paging ALLOC {key:#x} refused (status {st})"),
301    }
302}
303
304/// 2026-09-25: After the WRITE has completed, mark `key` resident.
305pub fn client_commit<T: Read + Write>(s: &mut T, key: u64) -> Result<()> {
306    send_req(s, OP_COMMIT, key)?;
307    match read_status(s)? {
308        ST_OK => Ok(()),
309        st => bail!("paging COMMIT {key:#x} failed (status {st})"),
310    }
311}
312
313/// 2026-09-25: `Some(offset)` to READ from, or `None` when the peer has no such key.
314pub fn client_get<T: Read + Write>(s: &mut T, key: u64) -> Result<Option<u64>> {
315    send_req(s, OP_GET, key)?;
316    match read_status(s)? {
317        ST_OK => Ok(Some(read_offset(s)?)),
318        ST_MISS => Ok(None),
319        st => bail!("paging GET {key:#x} error (status {st})"),
320    }
321}
322
323/// 2026-09-25: Ask the peer to drop `key`.
324pub fn client_remove<T: Read + Write>(s: &mut T, key: u64) -> Result<()> {
325    send_req(s, OP_REMOVE, key)?;
326    match read_status(s)? {
327        ST_OK => Ok(()),
328        st => bail!("paging REMOVE {key:#x} failed (status {st})"),
329    }
330}
331
332/// 2026-09-25: Tell the peer to end the paging loop; no reply is read.
333pub fn client_bye<T: Write>(s: &mut T) -> Result<()> {
334    send_req(s, OP_BYE, 0)
335}
336
337/// 2026-09-25: [`run_paging_loop`] over a residency shared by many connections, so a
338/// blob one client PUTs another can GET under the same key. The lock is held for
339/// each dispatch, which may move bytes to or from swap, and never across a TCP read
340/// or write.
341pub fn run_paging_loop_shared<T: Read + Write, A: SlotArena, S: SwapStore>(
342    stream: &mut T,
343    res: &std::sync::Mutex<Residency<A, S>>,
344) -> Result<()> {
345    // 2026-09-25: This connection's read pin (`handle_paging_op`), taken under the
346    // same lock as the dispatch.
347    let mut pinned: Option<u64> = None;
348    loop {
349        let mut op = [0u8; 1];
350        if stream.read_exact(&mut op).is_err() {
351            break;
352        }
353        if op[0] == OP_BYE {
354            break;
355        }
356        let mut kb = [0u8; 8];
357        stream.read_exact(&mut kb).context("read paging key")?;
358        let key = u64::from_le_bytes(kb);
359        let reply = {
360            let mut g = res.lock().expect("shared residency mutex poisoned");
361            handle_paging_op(&mut g, op[0], key, &mut pinned)
362        };
363        write_reply(stream, &reply)?;
364    }
365    if let Some(pk) = pinned {
366        res.lock()
367            .expect("shared residency mutex poisoned")
368            .unpin_read(pk);
369    }
370    Ok(())
371}
372
373#[cfg(test)]
374#[path = "wire_tests.rs"]
375mod tests;
metrale_storage/snapshot_swap/wire.rs

metrale_storage/snapshot_swap/
wire.rs
