kevy_cli/
bulk.rs

1//! **v2.10** — prefix bulk ops + diagnostics (RFC D4/D5):
2//! `copy-prefix` / `delete-prefix` (token-bucket rate limit,
3//! `--dry-run`), `digest`, `diff`, `inspect`.
4
5use std::io::{self, Write};
6use std::time::{Duration, Instant};
7
8use kevy_resp::Reply;
9use kevy_resp_client::RespClient;
10
11/// Token bucket: `rate` ops/second, starting EMPTY (strict pacing —
12/// a full-bucket start lets a small job burn its whole burst
13/// unthrottled, defeating the point of `--rate` for short sweeps).
14/// `rate == 0` = unlimited.
15pub struct RateLimiter {
16    rate: u64,
17    tokens: f64,
18    last: Instant,
19}
20
21impl RateLimiter {
22    /// New limiter at `rate` ops/s (0 = off).
23    pub fn new(rate: u64) -> Self {
24        Self { rate, tokens: 0.0, last: Instant::now() }
25    }
26
27    /// Block until one op is admitted.
28    pub fn take(&mut self) {
29        if self.rate == 0 {
30            return;
31        }
32        loop {
33            let now = Instant::now();
34            self.tokens = (self.tokens + now.duration_since(self.last).as_secs_f64() * self.rate as f64)
35                .min(self.rate as f64);
36            self.last = now;
37            if self.tokens >= 1.0 {
38                self.tokens -= 1.0;
39                return;
40            }
41            std::thread::sleep(Duration::from_millis(2));
42        }
43    }
44}
45
46fn scan_page(client: &mut RespClient, cursor: &[u8], pattern: &[u8]) -> io::Result<(Vec<u8>, Vec<Vec<u8>>)> {
47    let reply = client.request_borrowed(&[b"SCAN", cursor, b"MATCH", pattern, b"COUNT", b"512"])?;
48    let Reply::Array(items) = reply else {
49        return Err(io::Error::new(io::ErrorKind::InvalidData, "SCAN reply shape"));
50    };
51    let (Some(Reply::Bulk(next)), Some(Reply::Array(keys))) = (items.first(), items.get(1)) else {
52        return Err(io::Error::new(io::ErrorKind::InvalidData, "SCAN reply shape"));
53    };
54    let keys = keys
55        .iter()
56        .filter_map(|k| if let Reply::Bulk(b) = k { Some(b.clone()) } else { None })
57        .collect();
58    Ok((next.clone(), keys))
59}
60
61/// `delete-prefix`: SCAN + UNLINK, rate-limited. Returns deleted count.
62pub fn run_delete_prefix(
63    client: &mut RespClient,
64    prefix: &[u8],
65    rate: u64,
66    dry_run: bool,
67) -> io::Result<u64> {
68    let mut pattern = prefix.to_vec();
69    pattern.push(b'*');
70    let mut cursor: Vec<u8> = b"0".to_vec();
71    let mut limiter = RateLimiter::new(rate);
72    let mut n = 0u64;
73    loop {
74        let (next, keys) = scan_page(client, &cursor, &pattern)?;
75        for key in &keys {
76            if dry_run {
77                n += 1;
78                continue;
79            }
80            limiter.take();
81            if let Reply::Int(d) = client.request_borrowed(&[b"UNLINK", key])? {
82                n += d as u64;
83            }
84        }
85        cursor = next;
86        if cursor == b"0" {
87            return Ok(n);
88        }
89    }
90}
91
92/// `copy-prefix`: SCAN src prefix, re-key under dst prefix via
93/// read+rebuild frames (the server has no COPY verb; TTL carried as
94/// absolute PEXPIREAT). Rate-limited per source key.
95pub fn run_copy_prefix(
96    client: &mut RespClient,
97    src_prefix: &[u8],
98    dst_prefix: &[u8],
99    rate: u64,
100) -> io::Result<u64> {
101    let mut pattern = src_prefix.to_vec();
102    pattern.push(b'*');
103    let mut cursor: Vec<u8> = b"0".to_vec();
104    let mut limiter = RateLimiter::new(rate);
105    let mut n = 0u64;
106    loop {
107        let (next, keys) = scan_page(client, &cursor, &pattern)?;
108        for key in &keys {
109            limiter.take();
110            let mut dst = dst_prefix.to_vec();
111            dst.extend_from_slice(&key[src_prefix.len()..]);
112            let Some(frames) = crate::migrate::rebuild_frames(client, key, &dst)? else {
113                continue;
114            };
115            let n_cmds = count_commands(&frames);
116            for r in client.pipeline_raw(&frames, n_cmds)? {
117                if let Reply::Error(e) = r {
118                    return Err(io::Error::new(
119                        io::ErrorKind::InvalidData,
120                        String::from_utf8_lossy(&e).into_owned(),
121                    ));
122                }
123            }
124            n += 1;
125        }
126        cursor = next;
127        if cursor == b"0" {
128            return Ok(n);
129        }
130    }
131}
132
133/// Number of RESP commands in a frame buffer (top-level '*' headers).
134fn count_commands(mut b: &[u8]) -> usize {
135    let mut n = 0;
136    while let Some(len) = crate::migrate::command_len_pub(b) {
137        n += 1;
138        b = &b[len..];
139    }
140    n
141}
142
143/// `digest <prefix>` → (count, hex).
144pub fn run_digest(client: &mut RespClient, prefix: &[u8]) -> io::Result<(i64, String)> {
145    let r = client.request_borrowed(&[b"PREFIX.DIGEST", prefix])?;
146    let Reply::Array(items) = r else {
147        return Err(io::Error::new(io::ErrorKind::InvalidData, "PREFIX.DIGEST reply"));
148    };
149    let (Some(Reply::Int(n)), Some(Reply::Bulk(hex))) = (items.first(), items.get(1)) else {
150        return Err(io::Error::new(io::ErrorKind::InvalidData, "PREFIX.DIGEST reply"));
151    };
152    Ok((*n, String::from_utf8_lossy(hex).into_owned()))
153}
154
155/// `diff`: compare prefixes across two servers. Returns mismatching
156/// prefixes.
157pub fn run_diff(
158    a: &mut RespClient,
159    b: &mut RespClient,
160    prefixes: &[Vec<u8>],
161    out: &mut impl Write,
162) -> io::Result<Vec<Vec<u8>>> {
163    let mut bad = Vec::new();
164    for p in prefixes {
165        let (na, da) = run_digest(a, p)?;
166        let (nb, db) = run_digest(b, p)?;
167        let ok = na == nb && da == db;
168        writeln!(
169            out,
170            "{}  A: {na} keys {da}  B: {nb} keys {db}  {}",
171            String::from_utf8_lossy(p),
172            if ok { "OK" } else { "MISMATCH" }
173        )?;
174        if !ok {
175            bad.push(p.clone());
176        }
177    }
178    Ok(bad)
179}
180
181/// `inspect <prefix>`: sample keys, type distribution, sizes.
182pub fn run_inspect(client: &mut RespClient, prefix: &[u8], out: &mut impl Write) -> io::Result<()> {
183    let mut pattern = prefix.to_vec();
184    pattern.push(b'*');
185    let mut cursor: Vec<u8> = b"0".to_vec();
186    let mut total = 0u64;
187    let mut by_type: Vec<(String, u64)> = Vec::new();
188    let mut samples: Vec<String> = Vec::new();
189    loop {
190        let (next, keys) = scan_page(client, &cursor, &pattern)?;
191        for key in &keys {
192            total += 1;
193            if samples.len() < 8 {
194                samples.push(String::from_utf8_lossy(key).into_owned());
195            }
196            if let Reply::Simple(t) = client.request_borrowed(&[b"TYPE", key])? {
197                let t = String::from_utf8_lossy(&t).into_owned();
198                match by_type.iter_mut().find(|(n, _)| *n == t) {
199                    Some((_, c)) => *c += 1,
200                    None => by_type.push((t, 1)),
201                }
202            }
203        }
204        cursor = next;
205        if cursor == b"0" {
206            break;
207        }
208    }
209    writeln!(out, "prefix {}: {total} keys", String::from_utf8_lossy(prefix))?;
210    by_type.sort_by_key(|(_, c)| std::cmp::Reverse(*c));
211    for (t, c) in &by_type {
212        writeln!(out, "  {t}: {c}")?;
213    }
214    if !samples.is_empty() {
215        writeln!(out, "  samples: {}", samples.join(", "))?;
216    }
217    Ok(())
218}
kevy_cli/bulk.rs

kevy_cli/
bulk.rs
