kevy_cli/
migrate.rs

1//! Migration toolchain: `export` / `import`.
2//!
3//! Wire format = a RESP command stream of rebuild frames (SET / HSET /
4//! RPUSH / SADD / ZADD / PEXPIREAT) — bidirectionally compatible with
5//! `redis-cli --pipe`. Export walks SCAN cursors (per-key
6//! point-in-time; SCAN-class consistency). Import pipelines 512
7//! commands per batch with a fsynced progress file for `--resume`.
8//! Every key's frames start with DEL, so replay REBUILDS the key from
9//! scratch — genuinely idempotent for every type (RPUSH would
10//! otherwise append on re-import; the round-trip test caught exactly
11//! that). No cross-batch atomicity to lose.
12
13use std::fs::{File, OpenOptions};
14use std::io::{self, BufWriter, Read, Seek, SeekFrom, Write};
15use std::path::Path;
16
17use kevy_resp::{Reply, encode_command_borrowed};
18use kevy_resp_client::RespClient;
19
20const PIPELINE: usize = 512;
21
22/// Run `export` — walk the keyspace (optionally under `prefix`) and
23/// write rebuild frames to `out_path`. Returns exported key count.
24/// What an export did — including what it did NOT do. The skipped map
25/// is the half that used to be invisible: a type with no rebuild verb
26/// produced no frames, no error and no mention, so a migration could
27/// report success while leaving a whole type behind.
28pub struct Export {
29    /// Keys whose frames are in the file.
30    pub keys: u64,
31    /// Type name -> keys left out because nothing here rebuilds them.
32    pub skipped: std::collections::BTreeMap<Vec<u8>, u64>,
33}
34
35/// Walk the keyspace (optionally under `prefix`) and write rebuild
36/// frames to `out_path`, reporting both what went in and what could
37/// not.
38pub fn run_export(
39    client: &mut RespClient,
40    prefix: Option<&[u8]>,
41    out_path: &Path,
42) -> io::Result<Export> {
43    let mut out = BufWriter::new(File::create(out_path)?);
44    let mut cursor: Vec<u8> = b"0".to_vec();
45    let mut pattern = prefix.unwrap_or_default().to_vec();
46    pattern.push(b'*');
47    let mut n = 0u64;
48    // Types with no rebuild verb, counted by name. Silence here is how
49    // a migration loses a whole type and reports success.
50    let mut skipped: std::collections::BTreeMap<Vec<u8>, u64> = Default::default();
51    loop {
52        let reply =
53            client.request_borrowed(&[b"SCAN", &cursor, b"MATCH", &pattern, b"COUNT", b"512"])?;
54        let Reply::Array(items) = reply else {
55            return Err(io::Error::new(io::ErrorKind::InvalidData, "SCAN reply shape"));
56        };
57        let (Some(Reply::Bulk(next)), Some(Reply::Array(keys))) = (items.first(), items.get(1))
58        else {
59            return Err(io::Error::new(io::ErrorKind::InvalidData, "SCAN reply shape"));
60        };
61        let next = next.clone();
62        for k in keys {
63            let Reply::Bulk(key) = k else { continue };
64            let key = key.clone();
65            match export_key(client, &key, &mut out)? {
66                Some(None) => n += 1,
67                Some(Some(ty)) => *skipped.entry(ty).or_insert(0u64) += 1,
68                None => {}
69            }
70        }
71        cursor = next;
72        if cursor == b"0" {
73            break;
74        }
75    }
76    out.flush()?;
77    Ok(Export { keys: n, skipped })
78}
79
80/// Emit one key's rebuild frames. `Ok(None)` = nothing written; the
81/// type name comes back when the reason was "no rebuild verb", so the
82/// caller can report what it is leaving behind rather than count it as
83/// a key that happened to vanish.
84fn export_key(
85    client: &mut RespClient,
86    key: &[u8],
87    out: &mut impl Write,
88) -> io::Result<Option<Option<Vec<u8>>>> {
89    match rebuild_frames(client, key, key)? {
90        Rebuilt::Frames(frame) => {
91            out.write_all(&frame)?;
92            Ok(Some(None))
93        }
94        Rebuilt::Vanished => Ok(None),
95        Rebuilt::UnsupportedType(ty) => Ok(Some(Some(ty))),
96    }
97}
98
99/// Why a key produced no frames. The two reasons are not the same and
100/// were indistinguishable: a vanished key is a race the walk expects,
101/// an unsupported type is data the caller is about to leave behind.
102pub(crate) enum Rebuilt {
103    /// The rebuild frames, ready to write.
104    Frames(Vec<u8>),
105    /// The key was gone between SCAN and read — expected, uncounted.
106    Vanished,
107    /// The type has no rebuild verb here. Carries the type name so the
108    /// caller can tell someone rather than skip in silence.
109    UnsupportedType(Vec<u8>),
110}
111
112/// Read `key` and produce DEL+rebuild frames addressed to `dst`
113/// (`dst == key` for export; a re-prefixed name for copy-prefix —
114/// the server has no COPY verb, so copying IS read+rebuild).
115pub(crate) fn rebuild_frames(
116    client: &mut RespClient,
117    key: &[u8],
118    dst: &[u8],
119) -> io::Result<Rebuilt> {
120    let ty = match client.request_borrowed(&[b"TYPE", key])? {
121        Reply::Simple(t) => t,
122        _ => return Ok(Rebuilt::Vanished),
123    };
124    if ty == b"none" {
125        return Ok(Rebuilt::Vanished);
126    }
127    let mut frame = Vec::new();
128    // DEL first: replay rebuilds from scratch (idempotence for
129    // append-shaped verbs like RPUSH).
130    encode_command_borrowed(&mut frame, &[b"DEL", dst]);
131    match encode_body(client, key, dst, &ty, &mut frame)? {
132        Some(()) => {}
133        None => return Ok(Rebuilt::Vanished),
134    }
135    if frame.len() == encoded_del_len(dst) {
136        return Ok(Rebuilt::UnsupportedType(ty));
137    }
138    append_ttl_frame(client, key, dst, &mut frame)?;
139    Ok(Rebuilt::Frames(frame))
140}
141
142/// The `DEL <dst>` prologue's encoded length — how `rebuild_frames`
143/// tells "the body wrote nothing" from "the body wrote frames".
144fn encoded_del_len(dst: &[u8]) -> usize {
145    let mut probe = Vec::new();
146    encode_command_borrowed(&mut probe, &[b"DEL", dst]);
147    probe.len()
148}
149
150/// Append the type's rebuild verbs to `frame`. `None` = the key
151/// vanished mid-read; leaving `frame` untouched = no verb for this
152/// type, which the caller turns into `UnsupportedType`.
153fn encode_body(
154    client: &mut RespClient,
155    key: &[u8],
156    dst: &[u8],
157    ty: &[u8],
158    frame: &mut Vec<u8>,
159) -> io::Result<Option<()>> {
160    match ty {
161        b"string" => {
162            let Reply::Bulk(v) = client.request_borrowed(&[b"GET", key])? else {
163                return Ok(None);
164            };
165            encode_command_borrowed(frame, &[b"SET", dst, &v]);
166        }
167        b"hash" => {
168            let Some(items) = fetch_bulks(client, &[b"HGETALL", key])? else {
169                return Ok(None);
170            };
171            encode_multi(frame, b"HSET", dst, &items);
172        }
173        b"list" => {
174            let Some(vals) = fetch_bulks(client, &[b"LRANGE", key, b"0", b"-1"])? else {
175                return Ok(None);
176            };
177            encode_multi(frame, b"RPUSH", dst, &vals);
178        }
179        b"set" => {
180            let Some(ms) = fetch_bulks(client, &[b"SMEMBERS", key])? else {
181                return Ok(None);
182            };
183            encode_multi(frame, b"SADD", dst, &ms);
184        }
185        b"zset" => {
186            let zrange: &[&[u8]] = &[b"ZRANGE", key, b"0", b"-1", b"WITHSCORES"];
187            let Some(flat) = fetch_bulks(client, zrange)? else {
188                return Ok(None);
189            };
190            encode_zadd(frame, dst, &flat);
191        }
192        // Streams and anything added later: no rebuild verb here. The
193        // caller reports it by name — a migration that drops a type
194        // must say which one.
195        _ => return Ok(Some(())),
196    }
197    Ok(Some(()))
198}
199
200/// Issue `cmd` and unwrap its Array reply into bulk payloads.
201/// `None` when the reply isn't an array or the array is empty (the key
202/// vanished / changed type between TYPE and read).
203fn fetch_bulks(client: &mut RespClient, cmd: &[&[u8]]) -> io::Result<Option<Vec<Vec<u8>>>> {
204    let Reply::Array(items) = client.request_borrowed(cmd)? else {
205        return Ok(None);
206    };
207    if items.is_empty() {
208        return Ok(None);
209    }
210    Ok(Some(
211        items
212            .into_iter()
213            .filter_map(|r| if let Reply::Bulk(b) = r { Some(b) } else { None })
214            .collect(),
215    ))
216}
217
218/// Encode `<verb> <dst> <vals…>` onto `frame` (HSET / RPUSH / SADD).
219fn encode_multi(frame: &mut Vec<u8>, verb: &[u8], dst: &[u8], vals: &[Vec<u8>]) {
220    let mut argv: Vec<&[u8]> = vec![verb, dst];
221    argv.extend(vals.iter().map(Vec::as_slice));
222    encode_command_borrowed(frame, &argv);
223}
224
225/// Encode `ZADD <dst> score member …` onto `frame`.
226/// ZADD wants score member; ZRANGE gives member score.
227fn encode_zadd(frame: &mut Vec<u8>, dst: &[u8], flat: &[Vec<u8>]) {
228    let mut argv: Vec<&[u8]> = vec![b"ZADD", dst];
229    for pair in flat.chunks(2) {
230        if pair.len() == 2 {
231            argv.push(&pair[1]);
232            argv.push(&pair[0]);
233        }
234    }
235    encode_command_borrowed(frame, &argv);
236}
237
238/// TTL rides as an absolute PEXPIREAT follow-up.
239fn append_ttl_frame(
240    client: &mut RespClient,
241    key: &[u8],
242    dst: &[u8],
243    frame: &mut Vec<u8>,
244) -> io::Result<()> {
245    if let Reply::Int(ms) = client.request_borrowed(&[b"PTTL", key])?
246        && ms > 0
247    {
248        let now = std::time::SystemTime::now()
249            .duration_since(std::time::UNIX_EPOCH)
250            .map_err(io::Error::other)?
251            .as_millis() as i64;
252        encode_command_borrowed(frame, &[b"PEXPIREAT", dst, (now + ms).to_string().as_bytes()]);
253    }
254    Ok(())
255}
256
257/// Import stats.
258pub struct ImportReport {
259    /// Commands sent successfully.
260    pub sent: u64,
261    /// -ERR replies (counted, not fatal unless `strict`).
262    pub errors: u64,
263    /// Byte offset reached in the source file.
264    pub offset: u64,
265}
266
267/// Run `import` — stream `src` (a RESP command file) into the server,
268/// `PIPELINE` commands per batch. The progress file `<src>.progress`
269/// records the safely-applied byte offset after every batch (fsynced);
270/// `resume` starts there. Idempotent replay.
271pub fn run_import(
272    client: &mut RespClient,
273    src: &Path,
274    resume: bool,
275    strict: bool,
276) -> io::Result<ImportReport> {
277    let (mut progress, start) = open_progress(src, resume)?;
278    let mut f = File::open(src)?;
279    f.seek(SeekFrom::Start(start))?;
280    let mut pending: Vec<u8> = Vec::with_capacity(1 << 20);
281    let mut report = ImportReport { sent: 0, errors: 0, offset: start };
282    let mut chunk = vec![0u8; 1 << 20];
283    let mut batch_bytes = 0usize;
284    let mut batch_cmds = 0usize;
285    loop {
286        let n = f.read(&mut chunk)?;
287        if n == 0 {
288            break;
289        }
290        pending.extend_from_slice(&chunk[..n]);
291        // carve complete commands off `pending`
292        while let Some(used) = command_len(&pending[batch_bytes..]) {
293            batch_bytes += used;
294            batch_cmds += 1;
295            if batch_cmds == PIPELINE {
296                flush_batch(client, &pending[..batch_bytes], batch_cmds, strict, &mut report)?;
297                pending.drain(..batch_bytes);
298                write_progress(&mut progress, report.offset)?;
299                batch_bytes = 0;
300                batch_cmds = 0;
301            }
302        }
303    }
304    if batch_cmds > 0 {
305        flush_batch(client, &pending[..batch_bytes], batch_cmds, strict, &mut report)?;
306        write_progress(&mut progress, report.offset)?;
307    }
308    Ok(report)
309}
310
311fn flush_batch(
312    client: &mut RespClient,
313    raw: &[u8],
314    n: usize,
315    strict: bool,
316    report: &mut ImportReport,
317) -> io::Result<()> {
318    let replies = client.pipeline_raw(raw, n)?;
319    for r in replies {
320        if let Reply::Error(e) = r {
321            report.errors += 1;
322            if strict {
323                return Err(io::Error::new(
324                    io::ErrorKind::InvalidData,
325                    format!("server error (strict): {}", String::from_utf8_lossy(&e)),
326                ));
327            }
328        } else {
329            report.sent += 1;
330        }
331    }
332    report.offset += raw.len() as u64;
333    Ok(())
334}
335
336/// The progress file, opened, and the offset to start from.
337///
338/// The name is `<src>.progress`, APPENDED — the doc always said so, but
339/// `with_extension` turned dump.kevy into dump.progress, and everything
340/// written against the documented name (the migration drill's cleanup
341/// included) touched a file this code never read.
342///
343/// A fresh (non-resume) import resets the file immediately, before any
344/// batch. A completed earlier import leaves offset = EOF on disk; kill a
345/// fresh import before its first batch lands and that stale EOF
346/// survives — a later `--resume` then seeks to the end, imports nothing,
347/// and reports success. The migration drill hit exactly that window.
348fn open_progress(src: &Path, resume: bool) -> io::Result<(File, u64)> {
349    let path = {
350        let mut os = src.as_os_str().to_owned();
351        os.push(".progress");
352        std::path::PathBuf::from(os)
353    };
354    let mut start = 0u64;
355    if resume && let Ok(text) = std::fs::read_to_string(&path) {
356        start = text.trim().parse().unwrap_or(0);
357    }
358    let mut f = OpenOptions::new().create(true).truncate(false).write(true).open(&path)?;
359    if !resume {
360        write_progress(&mut f, 0)?;
361    }
362    Ok((f, start))
363}
364
365fn write_progress(f: &mut File, offset: u64) -> io::Result<()> {
366    f.set_len(0)?;
367    f.seek(SeekFrom::Start(0))?;
368    f.write_all(offset.to_string().as_bytes())?;
369    f.sync_data()
370}
371
372/// Crate-visible alias for [`command_len`] (bulk copy counts frames).
373pub(crate) fn command_len_pub(b: &[u8]) -> Option<usize> {
374    command_len(b)
375}
376
377/// Length of one complete RESP command at the head of `b`, or `None`.
378fn command_len(b: &[u8]) -> Option<usize> {
379    let mut pos = 0usize;
380    let line = take_line(b, &mut pos)?;
381    if line.first() != Some(&b'*') {
382        return None;
383    }
384    let n: usize = std::str::from_utf8(&line[1..]).ok()?.trim().parse().ok()?;
385    for _ in 0..n {
386        let hdr = take_line(b, &mut pos)?;
387        if hdr.first() != Some(&b'$') {
388            return None;
389        }
390        let len: usize = std::str::from_utf8(&hdr[1..]).ok()?.trim().parse().ok()?;
391        if b.len() < pos + len + 2 {
392            return None;
393        }
394        pos += len + 2;
395    }
396    Some(pos)
397}
398
399fn take_line<'b>(b: &'b [u8], pos: &mut usize) -> Option<&'b [u8]> {
400    let rest = &b[*pos..];
401    let idx = rest.windows(2).position(|w| w == b"\r\n")?;
402    let line = &rest[..idx];
403    *pos += idx + 2;
404    Some(line)
405}
kevy_cli/migrate.rs

kevy_cli/
migrate.rs
