use std::io::{Read, Write};
use std::net::{TcpStream, ToSocketAddrs};
use std::time::Duration;
use super::Server;
use super::args::{self, Args};
use crate::reply::Out;
use yo_common::num::push_i64;
use yo_common::{Code, Error, Result};
use yo_kv::Ask;
const KEYS_NEEDS_EMPTY: &str =
"When using MIGRATE KEYS option, the key argument must be set to the empty string";
const TARGET_ERROR: &[u8] = b"ERR Target instance replied with error: ";
const CONNECT_FAILED: &[u8] = b"IOERR error or timeout connecting to the client";
const WRITE_FAILED: &[u8] = b"IOERR error or timeout writing to target instance";
const READ_FAILED: &[u8] = b"IOERR error or timeout reading to target instance";
const DEFAULT_TIMEOUT_MS: i64 = 1000;
const CACHE_MAX: usize = 64;
const CACHE_TTL_MS: u64 = 10_000;
const LINE_MAX: usize = 1024;
#[derive(Default)]
pub(super) struct Peers {
open: Vec<Peer>,
}
struct Peer {
at: String,
sock: TcpStream,
last_db: i64,
used_ms: u64,
}
impl Peers {
fn expire(&mut self, now: u64) {
self.open
.retain(|p| now.saturating_sub(p.used_ms) < CACHE_TTL_MS);
}
fn find(&self, at: &str) -> Option<usize> {
self.open.iter().position(|p| p.at == at)
}
fn close(&mut self, at: &str) {
self.open.retain(|p| p.at != at);
}
}
struct Plan<'a> {
host: &'a [u8],
port: &'a [u8],
db: i64,
timeout: Duration,
copy: bool,
replace: bool,
user: Option<&'a [u8]>,
pass: Option<&'a [u8]>,
first_key: usize,
num_keys: usize,
}
struct Going<'a> {
key: &'a [u8],
ttl: i64,
payload: Vec<u8>,
}
pub(super) fn execute(server: &mut Server, at: usize, args: Args<'_>, out: &mut Out) -> Result<()> {
yo_alloc::allow(|| run(server, at, args, out))
}
fn run(server: &mut Server, at: usize, args: Args<'_>, out: &mut Out) -> Result<()> {
let plan = parse(args)?;
let db = server.db(at);
let now = db.clock().now_ms();
let mut going: Vec<Going<'_>> = Vec::new();
for i in 0..plan.num_keys {
let key = args.get(plan.first_key + i);
let ttl = match db.deadline_of(key) {
Ask::Missing => continue,
Ask::NoDeadline => 0,
Ask::At(when) if when <= now => continue,
Ask::At(when) => (when - now).max(1) as i64,
};
let Some(payload) = db.dump(key) else {
continue;
};
going.push(Going { key, ttl, payload });
}
if going.is_empty() {
out.simple(b"NOKEY");
return Ok(());
}
let mut name = String::new();
name.push_str(&String::from_utf8_lossy(plan.host));
name.push(':');
name.push_str(&String::from_utf8_lossy(plan.port));
server.peers.expire(now);
let mut retry = true;
loop {
match attempt(server, &name, &plan, &going, now) {
Ok(replies) => {
finish(server, at, &name, &plan, &going, &replies, out);
return Ok(());
}
Err(Broke { failed, timed_out }) => {
server.peers.close(&name);
if retry && !timed_out {
retry = false;
continue;
}
out.error(failed);
return Ok(());
}
}
}
}
struct Broke {
failed: &'static [u8],
timed_out: bool,
}
struct Replies {
auth: Option<Vec<u8>>,
select: Option<Vec<u8>>,
keys: Vec<Vec<u8>>,
}
fn attempt(
server: &mut Server,
name: &str,
plan: &Plan<'_>,
going: &[Going<'_>],
now: u64,
) -> std::result::Result<Replies, Broke> {
let held = server.peers.find(name);
let last_db = held.map_or(-1, |i| server.peers.open[i].last_db);
let select = last_db != plan.db;
if held.is_none() {
let sock = connect(plan)?;
if server.peers.open.len() >= CACHE_MAX {
let oldest = server
.peers
.open
.iter()
.enumerate()
.min_by_key(|(_, p)| p.used_ms)
.map_or(0, |(i, _)| i);
server.peers.open.swap_remove(oldest);
}
server.peers.open.push(Peer {
at: name.to_string(),
sock,
last_db: -1,
used_ms: now,
});
}
let i = server
.peers
.find(name)
.expect("the socket was just put here");
let peer = &mut server.peers.open[i];
peer.used_ms = now;
let mut cmd =
Vec::with_capacity(64 + going.iter().map(|g| g.payload.len() + 64).sum::<usize>());
if let Some(pass) = plan.pass {
array(&mut cmd, if plan.user.is_some() { 3 } else { 2 });
bulk(&mut cmd, b"AUTH");
if let Some(user) = plan.user {
bulk(&mut cmd, user);
}
bulk(&mut cmd, pass);
}
if select {
array(&mut cmd, 2);
bulk(&mut cmd, b"SELECT");
int(&mut cmd, plan.db);
}
for g in going {
array(&mut cmd, if plan.replace { 5 } else { 4 });
bulk(&mut cmd, b"RESTORE");
bulk(&mut cmd, g.key);
int(&mut cmd, g.ttl);
bulk(&mut cmd, &g.payload);
if plan.replace {
bulk(&mut cmd, b"REPLACE");
}
}
peer.sock.write_all(&cmd).map_err(broke(WRITE_FAILED))?;
peer.sock.flush().map_err(broke(WRITE_FAILED))?;
let auth = match plan.pass {
Some(_) => Some(line(&mut peer.sock)?),
None => None,
};
let select = if select {
Some(line(&mut peer.sock)?)
} else {
None
};
let mut keys = Vec::with_capacity(going.len());
for _ in going {
keys.push(line(&mut peer.sock)?);
}
Ok(Replies { auth, select, keys })
}
fn finish(
server: &mut Server,
at: usize,
name: &str,
plan: &Plan<'_>,
going: &[Going<'_>],
replies: &Replies,
out: &mut Out,
) {
fn bad(r: &Option<Vec<u8>>) -> Option<&[u8]> {
r.as_deref().filter(|line| line.first() == Some(&b'-'))
}
let before = bad(&replies.auth).or_else(|| bad(&replies.select));
let mut told: Option<&[u8]> = None;
for (g, reply) in going.iter().zip(&replies.keys) {
let failed = before.or_else(|| {
if reply.first() == Some(&b'-') {
Some(reply.as_slice())
} else {
None
}
});
match failed {
Some(line) => told = told.or(Some(line)),
None if !plan.copy => {
server.db(at).del(g.key);
}
None => {}
}
}
if let Some(i) = server.peers.find(name) {
server.peers.open[i].last_db = if told.is_some() { -1 } else { plan.db };
}
match told {
Some(line) => out.error_line(TARGET_ERROR, &line[1..]),
None => out.ok(),
}
}
fn connect(plan: &Plan<'_>) -> std::result::Result<TcpStream, Broke> {
let host = std::str::from_utf8(plan.host).map_err(|_| Broke {
failed: WRITE_FAILED,
timed_out: false,
})?;
let mut addrs = (host, port_of(plan.port))
.to_socket_addrs()
.map_err(|_| Broke {
failed: WRITE_FAILED,
timed_out: false,
})?;
let addr = addrs.next().ok_or(Broke {
failed: WRITE_FAILED,
timed_out: false,
})?;
let sock = TcpStream::connect_timeout(&addr, plan.timeout).map_err(|e| {
let timed_out = is_timeout(&e);
Broke {
failed: if timed_out {
CONNECT_FAILED
} else {
WRITE_FAILED
},
timed_out,
}
})?;
sock.set_read_timeout(Some(plan.timeout))
.and_then(|()| sock.set_write_timeout(Some(plan.timeout)))
.and_then(|()| sock.set_nodelay(true))
.map_err(|_| Broke {
failed: WRITE_FAILED,
timed_out: false,
})?;
Ok(sock)
}
fn broke(failed: &'static [u8]) -> impl Fn(std::io::Error) -> Broke {
move |e| Broke {
failed,
timed_out: is_timeout(&e),
}
}
fn is_timeout(e: &std::io::Error) -> bool {
matches!(
e.kind(),
std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
)
}
fn line(sock: &mut TcpStream) -> std::result::Result<Vec<u8>, Broke> {
let mut out = Vec::with_capacity(32);
let mut b = [0u8; 1];
loop {
match sock.read(&mut b) {
Ok(0) => {
return Err(Broke {
failed: READ_FAILED,
timed_out: false,
});
}
Ok(_) => {}
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(e) => return Err(broke(READ_FAILED)(e)),
}
if b[0] == b'\n' {
if out.last() == Some(&b'\r') {
out.pop();
}
return Ok(out);
}
if out.len() >= LINE_MAX {
return Err(Broke {
failed: READ_FAILED,
timed_out: false,
});
}
out.push(b[0]);
}
}
fn parse(args: Args<'_>) -> Result<Plan<'_>> {
let mut plan = Plan {
host: args.get(1),
port: args.get(2),
db: 0,
timeout: Duration::from_millis(DEFAULT_TIMEOUT_MS as u64),
copy: false,
replace: false,
user: None,
pass: None,
first_key: 3,
num_keys: 1,
};
let mut i = 6;
while i < args.len() {
let arg = args.get(i);
let more = args.len() - i - 1;
if args::is(arg, b"copy") {
plan.copy = true;
} else if args::is(arg, b"replace") {
plan.replace = true;
} else if args::is(arg, b"auth") {
if more < 1 {
return Err(args::syntax());
}
plan.pass = Some(args.get(i + 1));
i += 1;
} else if args::is(arg, b"auth2") {
if more < 2 {
return Err(args::syntax());
}
plan.user = Some(args.get(i + 1));
plan.pass = Some(args.get(i + 2));
i += 2;
} else if args::is(arg, b"keys") {
if !args.get(3).is_empty() {
return Err(Error::new(Code::Invalid, KEYS_NEEDS_EMPTY));
}
plan.first_key = i + 1;
plan.num_keys = args.len() - i - 1;
break;
} else {
return Err(args::syntax());
}
i += 1;
}
plan.db = args.int(4)?;
let ms = args.int(5)?;
plan.timeout = Duration::from_millis(if ms <= 0 {
DEFAULT_TIMEOUT_MS as u64
} else {
ms as u64
});
Ok(plan)
}
fn port_of(s: &[u8]) -> u16 {
let mut i = 0;
while i < s.len() && s[i].is_ascii_whitespace() {
i += 1;
}
let negative = i < s.len() && s[i] == b'-';
if i < s.len() && (s[i] == b'-' || s[i] == b'+') {
i += 1;
}
let mut n: i64 = 0;
while i < s.len() && s[i].is_ascii_digit() {
n = n.saturating_mul(10).saturating_add(i64::from(s[i] - b'0'));
i += 1;
}
if negative {
n = -n;
}
n as u16
}
fn array(buf: &mut Vec<u8>, n: i64) {
buf.push(b'*');
push_i64(buf, n);
buf.extend_from_slice(b"\r\n");
}
fn bulk(buf: &mut Vec<u8>, s: &[u8]) {
buf.push(b'$');
push_i64(buf, s.len() as i64);
buf.extend_from_slice(b"\r\n");
buf.extend_from_slice(s);
buf.extend_from_slice(b"\r\n");
}
fn int(buf: &mut Vec<u8>, n: i64) {
let mut digits = Vec::with_capacity(20);
push_i64(&mut digits, n);
bulk(buf, &digits);
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dispatch::{Args, Session, execute};
use crate::proto::{Limits, Proto};
use crate::request::Argv;
use std::net::TcpListener;
use std::sync::{Arc, Mutex};
struct Fake {
port: String,
seen: Arc<Mutex<Vec<u8>>>,
}
impl Fake {
fn seen(&self) -> Vec<Vec<Vec<u8>>> {
commands(&self.seen.lock().expect("the peer thread has not panicked"))
}
fn words(&self) -> Vec<Vec<String>> {
self.seen()
.iter()
.map(|cmd| {
cmd.iter()
.map(|a| String::from_utf8_lossy(a).into_owned())
.collect()
})
.collect()
}
}
fn fake(rounds: Vec<Vec<&'static str>>) -> Fake {
let listener = TcpListener::bind("127.0.0.1:0").expect("a free port");
let port = listener
.local_addr()
.expect("the port it bound")
.port()
.to_string();
let seen = Arc::new(Mutex::new(Vec::new()));
let mine = Arc::clone(&seen);
std::thread::spawn(move || {
let Ok((mut sock, _)) = listener.accept() else {
return;
};
let mut buf = Vec::new();
let mut chunk = [0u8; 8192];
let mut done = 0usize;
for round in rounds {
let want = done + round.len();
while commands(&buf).len() < want {
match sock.read(&mut chunk) {
Ok(0) | Err(_) => return,
Ok(n) => buf.extend_from_slice(&chunk[..n]),
}
}
*mine.lock().expect("nobody poisoned it") = buf.clone();
for r in round {
if sock.write_all(r.as_bytes()).is_err() {
return;
}
}
done = want;
}
*mine.lock().expect("nobody poisoned it") = buf;
while sock.read(&mut chunk).is_ok_and(|n| n > 0) {}
});
Fake { port, seen }
}
fn commands(buf: &[u8]) -> Vec<Vec<Vec<u8>>> {
let mut out = Vec::new();
let mut i = 0;
while i < buf.len() {
let Some((n, after)) = header(buf, i, b'*') else {
break;
};
let mut cmd = Vec::new();
let mut at = after;
let mut whole = true;
for _ in 0..n {
let Some((len, body)) = header(buf, at, b'$') else {
whole = false;
break;
};
if body + len + 2 > buf.len() {
whole = false;
break;
}
cmd.push(buf[body..body + len].to_vec());
at = body + len + 2;
}
if !whole {
break;
}
out.push(cmd);
i = at;
}
out
}
fn header(buf: &[u8], at: usize, kind: u8) -> Option<(usize, usize)> {
if buf.get(at) != Some(&kind) {
return None;
}
let end = at + buf[at..].windows(2).position(|w| w == b"\r\n")?;
let n = std::str::from_utf8(&buf[at + 1..end]).ok()?.parse().ok()?;
Some((n, end + 2))
}
struct At {
server: Server,
session: Session,
argv: Argv,
out: Out,
}
impl At {
fn new() -> At {
At {
server: Server::new(),
session: Session::new(3),
argv: Argv::new(),
out: Out::new(Proto::Resp2),
}
}
fn run(&mut self, parts: &[&[u8]]) -> String {
let mut wire = format!("*{}\r\n", parts.len()).into_bytes();
for p in parts {
wire.extend_from_slice(format!("${}\r\n", p.len()).as_bytes());
wire.extend_from_slice(p);
wire.extend_from_slice(b"\r\n");
}
self.argv
.decode(&wire, &Limits::default())
.expect("the test wrote a legal command");
self.out.clear();
execute(
&mut self.server,
&mut self.session,
Args::new(&self.argv, &wire),
&mut self.out,
);
String::from_utf8_lossy(self.out.as_slice()).into_owned()
}
fn advance(&mut self, ms: u64) {
self.server.db(0).clock_mut().advance(ms);
}
}
#[test]
fn a_key_goes_over_and_leaves_nothing_behind() {
let peer = fake(vec![vec!["+OK\r\n", "+OK\r\n"]]);
let mut at = At::new();
at.run(&[b"set", b"k", b"v"]);
let reply = at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"k",
b"0",
b"1000",
]);
assert_eq!(reply, "+OK\r\n");
assert_eq!(at.run(&[b"exists", b"k"]), ":0\r\n");
let sent = peer.words();
assert_eq!(sent[0], ["SELECT", "0"]);
assert_eq!(sent[1][0], "RESTORE");
assert_eq!(sent[1][1], "k");
assert_eq!(sent[1][2], "0", "no deadline goes over as a zero ttl");
assert_eq!(sent[1].len(), 4, "and no REPLACE was asked for");
let mut other = At::new();
let payload = peer.seen()[1][3].clone();
assert_eq!(
other.run(&[b"restore", b"k", b"0", &payload]),
"+OK\r\n",
"the bytes that went over are a payload RESTORE takes"
);
assert_eq!(other.run(&[b"get", b"k"]), "$1\r\nv\r\n");
}
#[test]
fn copy_leaves_the_key_where_it_is_and_replace_is_passed_on() {
let peer = fake(vec![vec!["+OK\r\n", "+OK\r\n"]]);
let mut at = At::new();
at.run(&[b"set", b"k", b"v"]);
let reply = at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"k",
b"0",
b"1000",
b"COPY",
b"REPLACE",
]);
assert_eq!(reply, "+OK\r\n");
assert_eq!(at.run(&[b"exists", b"k"]), ":1\r\n");
let sent = peer.words();
assert_eq!(sent[1].len(), 5);
assert_eq!(sent[1][4], "REPLACE");
}
#[test]
fn a_key_that_is_not_here_is_not_an_error() {
let peer = fake(vec![]);
let mut at = At::new();
assert_eq!(
at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"gone",
b"0",
b"1000"
]),
"+NOKEY\r\n"
);
assert!(
peer.seen().is_empty(),
"and nothing was said to the other server"
);
}
#[test]
fn a_key_that_has_expired_is_the_same_nokey() {
let peer = fake(vec![]);
let mut at = At::new();
at.run(&[b"set", b"k", b"v", b"PX", b"50"]);
at.advance(60);
assert_eq!(
at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"k",
b"0",
b"1000"
]),
"+NOKEY\r\n"
);
}
#[test]
fn a_deadline_goes_over_as_what_is_left_of_it() {
let peer = fake(vec![vec!["+OK\r\n", "+OK\r\n"]]);
let mut at = At::new();
at.run(&[b"set", b"k", b"v", b"PX", b"5000"]);
at.advance(1000);
at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"k",
b"0",
b"1000",
]);
assert_eq!(peer.words()[1][2], "4000");
}
#[test]
fn the_keys_form_takes_the_ones_that_are_here() {
let peer = fake(vec![vec!["+OK\r\n", "+OK\r\n", "+OK\r\n"]]);
let mut at = At::new();
at.run(&[b"mset", b"a", b"1", b"b", b"2"]);
let reply = at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"",
b"0",
b"1000",
b"KEYS",
b"a",
b"missing",
b"b",
]);
assert_eq!(reply, "+OK\r\n");
let sent = peer.words();
assert_eq!(sent.len(), 3, "one SELECT and two RESTOREs");
assert_eq!(sent[1][1], "a");
assert_eq!(sent[2][1], "b");
assert_eq!(at.run(&[b"exists", b"a", b"b"]), ":0\r\n");
}
#[test]
fn the_keys_form_wants_the_key_argument_empty() {
let mut at = At::new();
assert_eq!(
at.run(&[
b"migrate",
b"127.0.0.1",
b"1",
b"k",
b"0",
b"1000",
b"KEYS",
b"k"
]),
format!("-ERR {KEYS_NEEDS_EMPTY}\r\n")
);
}
#[test]
fn the_empty_string_is_a_key_like_any_other() {
let peer = fake(vec![vec!["+OK\r\n", "+OK\r\n"]]);
let mut at = At::new();
at.run(&[b"set", b"", b"v"]);
assert_eq!(
at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"",
b"0",
b"1000"
]),
"+OK\r\n"
);
assert_eq!(peer.words()[1][1], "");
}
#[test]
fn an_option_written_after_keys_is_a_key() {
let peer = fake(vec![vec!["+OK\r\n", "+OK\r\n"]]);
let mut at = At::new();
at.run(&[b"set", b"COPY", b"v"]);
at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"",
b"0",
b"1000",
b"KEYS",
b"COPY",
]);
assert_eq!(peer.words()[1][1], "COPY");
assert_eq!(
at.run(&[b"exists", b"COPY"]),
":0\r\n",
"so it was migrated and not read as an option"
);
}
#[test]
fn the_options_are_read_before_the_numbers() {
let mut at = At::new();
assert_eq!(
at.run(&[
b"migrate",
b"127.0.0.1",
b"1",
b"k",
b"notanum",
b"alsonot",
b"BOGUS"
]),
"-ERR syntax error\r\n"
);
assert_eq!(
at.run(&[b"migrate", b"127.0.0.1", b"1", b"k", b"notanum", b"1000"]),
"-ERR value is not an integer or out of range\r\n"
);
}
#[test]
fn auth_and_auth2_go_out_in_front() {
let peer = fake(vec![vec!["+OK\r\n", "+OK\r\n", "+OK\r\n"]]);
let mut at = At::new();
at.run(&[b"set", b"k", b"v"]);
at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"k",
b"0",
b"1000",
b"AUTH2",
b"someone",
b"hunter2",
]);
let sent = peer.words();
assert_eq!(sent[0], ["AUTH", "someone", "hunter2"]);
assert_eq!(sent[1], ["SELECT", "0"]);
assert_eq!(sent[2][0], "RESTORE");
}
#[test]
fn auth_and_auth2_want_their_arguments() {
let mut at = At::new();
for bad in [
&[b"migrate".as_slice(), b"h", b"1", b"k", b"0", b"1", b"AUTH"][..],
&[b"migrate", b"h", b"1", b"k", b"0", b"1", b"AUTH2"],
&[b"migrate", b"h", b"1", b"k", b"0", b"1", b"AUTH2", b"user"],
] {
assert_eq!(at.run(bad), "-ERR syntax error\r\n");
}
}
#[test]
fn a_key_the_peer_refuses_stays_here() {
let peer = fake(vec![vec![
"+OK\r\n",
"-BUSYKEY Target key name already exists.\r\n",
]]);
let mut at = At::new();
at.run(&[b"set", b"k", b"v"]);
let reply = at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"k",
b"0",
b"1000",
]);
assert_eq!(
reply,
"-ERR Target instance replied with error: BUSYKEY Target key name already exists.\r\n"
);
assert_eq!(at.run(&[b"exists", b"k"]), ":1\r\n");
}
#[test]
fn one_key_refused_does_not_hold_the_others_back() {
let peer = fake(vec![vec!["+OK\r\n", "-BUSYKEY no\r\n", "+OK\r\n"]]);
let mut at = At::new();
at.run(&[b"mset", b"a", b"1", b"b", b"2"]);
let reply = at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"",
b"0",
b"1000",
b"KEYS",
b"a",
b"b",
]);
assert_eq!(
reply,
"-ERR Target instance replied with error: BUSYKEY no\r\n"
);
assert_eq!(
at.run(&[b"exists", b"a"]),
":1\r\n",
"the refused one stays"
);
assert_eq!(at.run(&[b"exists", b"b"]), ":0\r\n", "the taken one goes");
}
#[test]
fn a_select_that_fails_fails_every_key() {
let peer = fake(vec![vec![
"-ERR DB index is out of range\r\n",
"+OK\r\n",
"+OK\r\n",
]]);
let mut at = At::new();
at.run(&[b"mset", b"a", b"1", b"b", b"2"]);
let reply = at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"",
b"9",
b"1000",
b"KEYS",
b"a",
b"b",
]);
assert_eq!(
reply,
"-ERR Target instance replied with error: ERR DB index is out of range\r\n"
);
assert_eq!(at.run(&[b"exists", b"a", b"b"]), ":2\r\n");
}
#[test]
fn the_socket_is_kept_and_so_is_what_it_is_selected_on() {
let peer = fake(vec![
vec!["+OK\r\n", "+OK\r\n"],
vec!["+OK\r\n"],
vec!["-BUSYKEY no\r\n"],
vec!["+OK\r\n", "+OK\r\n"],
]);
let mut at = At::new();
at.run(&[b"mset", b"a", b"1", b"b", b"2", b"c", b"3", b"d", b"4"]);
let port = peer.port.clone();
let go = |at: &mut At, key: &[u8]| {
at.run(&[
b"migrate",
b"127.0.0.1",
port.as_bytes(),
key,
b"0",
b"1000",
])
};
assert_eq!(go(&mut at, b"a"), "+OK\r\n");
assert_eq!(go(&mut at, b"b"), "+OK\r\n");
assert_eq!(
go(&mut at, b"c"),
"-ERR Target instance replied with error: BUSYKEY no\r\n"
);
assert_eq!(go(&mut at, b"d"), "+OK\r\n");
let sent = peer.words();
let selects: Vec<&Vec<String>> = sent.iter().filter(|c| c[0] == "SELECT").collect();
assert_eq!(selects.len(), 2, "one at the start and one after the error");
assert_eq!(sent[0][0], "SELECT");
assert_eq!(sent[1][0], "RESTORE");
assert_eq!(sent[2][0], "RESTORE", "the second one did not select again");
assert_eq!(sent[3][0], "RESTORE");
assert_eq!(sent[4][0], "SELECT", "and the fourth one had to");
}
#[test]
fn a_peer_that_is_not_listening_is_an_io_error() {
let port = TcpListener::bind("127.0.0.1:0")
.expect("a free port")
.local_addr()
.expect("the port it bound")
.port()
.to_string();
let mut at = At::new();
at.run(&[b"set", b"k", b"v"]);
let reply = at.run(&[
b"migrate",
b"127.0.0.1",
port.as_bytes(),
b"k",
b"0",
b"200",
]);
let either = [CONNECT_FAILED, WRITE_FAILED]
.map(|line| format!("-{}\r\n", String::from_utf8_lossy(line)));
assert!(either.contains(&reply), "got {reply:?}");
assert_eq!(at.run(&[b"exists", b"k"]), ":1\r\n", "and the key stays");
}
#[test]
fn a_peer_that_never_answers_is_a_read_error() {
let peer = fake(vec![]);
let mut at = At::new();
at.run(&[b"set", b"k", b"v"]);
assert_eq!(
at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"k",
b"0",
b"200"
]),
format!("-{}\r\n", String::from_utf8_lossy(READ_FAILED))
);
assert_eq!(at.run(&[b"exists", b"k"]), ":1\r\n");
}
#[test]
fn a_timeout_that_is_not_positive_is_one_second() {
let peer = fake(vec![vec!["+OK\r\n", "+OK\r\n"]]);
let mut at = At::new();
at.run(&[b"set", b"k", b"v"]);
assert_eq!(
at.run(&[
b"migrate",
b"127.0.0.1",
peer.port.as_bytes(),
b"k",
b"0",
b"-5"
]),
"+OK\r\n"
);
}
#[test]
fn a_port_is_read_the_way_atoi_reads_one() {
assert_eq!(port_of(b"6379"), 6379);
assert_eq!(port_of(b" 6379"), 6379);
assert_eq!(port_of(b"63rubbish"), 63);
assert_eq!(port_of(b"rubbish"), 0);
assert_eq!(port_of(b""), 0);
assert_eq!(
port_of(b"-1"),
65535,
"the low sixteen bits, as htons sees it"
);
assert_eq!(port_of(b"65536"), 0);
}
}