use std::io::{BufRead, BufReader, Read, Write};
use std::net::TcpListener;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use ytsaurus_client::{Client, RetryPolicy};
use ytsaurus_yson::{YsonFormat, YsonNode, YsonValue, from_slice};
const TX: &str = "3-5bc70-10001-387a";
#[test]
fn a_started_handles_drop_still_aborts() {
let cluster = StubCluster::answering(Answers::default());
{
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let tx = client.start_transaction().expect("starts");
assert_eq!(tx.id(), TX);
}
let aborts = cluster.heads_for("abort_transaction");
assert_eq!(
aborts.len(),
1,
"a dropped started handle must abort exactly once: {:?}",
cluster.request_lines()
);
assert_eq!(
str_param(&aborts[0], "transaction_id").as_deref(),
Some(TX),
"the abort named the wrong transaction:\n{}",
aborts[0]
);
}
#[test]
fn a_detached_transaction_is_neither_aborted_nor_pinged_again() {
let cluster = StubCluster::answering(Answers::default());
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let tx = client
.start_transaction_with(Duration::from_secs(3))
.expect("starts");
let id = tx.detach();
assert_eq!(id, TX, "detach must hand back the transaction's own id");
let settled = cluster.request_count();
std::thread::sleep(Duration::from_millis(1600));
assert!(
cluster.heads_for("abort_transaction").is_empty(),
"detach sent an abort: {:?}",
cluster.request_lines()
);
assert_eq!(
cluster.request_count(),
settled,
"something was sent after detach returned: {:?}",
cluster.request_lines()
);
}
#[test]
fn detach_with_a_ping_in_flight_waits_for_it() {
const ANSWERED_AFTER: Duration = Duration::from_millis(500);
let cluster = StubCluster::answering(Answers {
ping_delay: ANSWERED_AFTER,
..Answers::default()
});
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let tx = client
.start_transaction_with(Duration::from_secs(3))
.expect("starts");
let deadline = Instant::now() + Duration::from_secs(5);
while cluster.heads_for("ping_transaction").is_empty() {
assert!(
Instant::now() < deadline,
"no ping arrived within 5 s of starting a 3 s transaction"
);
std::thread::sleep(Duration::from_millis(10));
}
let waited = Instant::now();
let id = tx.detach(); let waited = waited.elapsed();
assert_eq!(id, TX);
assert!(
waited >= ANSWERED_AFTER / 2,
"detach returned in {waited:?} with a ping being answered {ANSWERED_AFTER:?} late: \
it did not wait for the ping in flight"
);
let settled = cluster.request_count();
std::thread::sleep(Duration::from_millis(1600));
assert!(
cluster.heads_for("abort_transaction").is_empty(),
"detach under a ping in flight aborted: {:?}",
cluster.request_lines()
);
assert_eq!(
cluster.request_count(),
settled,
"a request started after detach returned: {:?}",
cluster.request_lines()
);
}
#[test]
fn attach_reads_the_timeout_pings_and_its_drop_does_not_abort() {
let cluster = StubCluster::answering(Answers {
timeout_ms: 3000,
..Answers::default()
});
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let tx = client.attach_transaction(TX).expect("attaches");
assert_eq!(tx.id(), TX);
let gets = cluster.heads_for("/api/v4/get");
assert_eq!(gets.len(), 1, "{:?}", cluster.request_lines());
assert_eq!(
str_param(&gets[0], "path").as_deref(),
Some(format!("#{TX}/@timeout").as_str()),
"the timeout was read from somewhere else:\n{}",
gets[0]
);
let deadline = Instant::now() + Duration::from_secs(5);
let ping = loop {
if let Some(head) = cluster.heads_for("ping_transaction").into_iter().next() {
break head;
}
assert!(
Instant::now() < deadline,
"an attached handle sent no ping within 5 s"
);
std::thread::sleep(Duration::from_millis(10));
};
assert_eq!(str_param(&ping, "transaction_id").as_deref(), Some(TX));
drop(tx);
std::thread::sleep(Duration::from_millis(1000));
let settled = cluster.request_count();
std::thread::sleep(Duration::from_millis(1600));
assert!(
cluster.heads_for("abort_transaction").is_empty(),
"an attached handle's drop aborted the owner's transaction: {:?}",
cluster.request_lines()
);
assert_eq!(
cluster.request_count(),
settled,
"the pings did not stop when the attached handle dropped: {:?}",
cluster.request_lines()
);
}
#[test]
fn attaching_pings_before_it_returns() {
let cluster = StubCluster::answering(Answers {
timeout_ms: 30_000,
..Answers::default()
});
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let tx = client.attach_transaction(TX).expect("attaches");
let pings = cluster.heads_for("ping_transaction");
assert_eq!(
pings.len(),
1,
"attach must restart the transaction's clock before handing back a handle: {:?}",
cluster.request_lines()
);
assert_eq!(str_param(&pings[0], "transaction_id").as_deref(), Some(TX));
drop(tx);
}
#[test]
fn attaching_to_a_transaction_that_dies_between_the_two_reads_fails_here() {
let cluster = StubCluster::answering(Answers {
ping_gone: true,
..Answers::default()
});
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let error = client
.attach_transaction(TX)
.expect_err("the transaction is gone");
let message = error.to_string();
for expected in ["attach", TX, "No such transaction"] {
assert!(
message.contains(expected),
"the error does not say {expected:?}: {message}"
);
}
let before = cluster.request_count();
std::thread::sleep(Duration::from_millis(300));
assert_eq!(
cluster.request_count(),
before,
"a failed attach left a ping thread behind: {:?}",
cluster.request_lines()
);
}
#[test]
fn a_handle_says_so_when_the_cluster_says_the_transaction_is_gone() {
let cluster = StubCluster::answering(Answers {
ping_gone: true,
..Answers::default()
});
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let tx = client
.start_transaction_with(Duration::from_secs(3))
.expect("starts");
assert!(!tx.is_lost(), "nothing has been answered yet");
let deadline = Instant::now() + Duration::from_secs(5);
while !tx.is_lost() {
assert!(
Instant::now() < deadline,
"the handle never noticed: {:?}",
cluster.request_lines()
);
std::thread::sleep(Duration::from_millis(20));
}
let settled = cluster.request_count();
std::thread::sleep(Duration::from_millis(1600));
assert_eq!(
cluster.request_count(),
settled,
"the thread gave up and went on pinging: {:?}",
cluster.request_lines()
);
}
#[test]
fn an_attached_handles_explicit_abort_still_aborts() {
let cluster = StubCluster::answering(Answers {
timeout_ms: 30_000,
..Answers::default()
});
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let tx = client.attach_transaction(TX).expect("attaches");
tx.abort().expect("aborts");
let aborts = cluster.heads_for("abort_transaction");
assert_eq!(
aborts.len(),
1,
"an attached handle's explicit abort was swallowed: {:?}",
cluster.request_lines()
);
assert_eq!(str_param(&aborts[0], "transaction_id").as_deref(), Some(TX));
}
#[test]
fn detaching_an_attached_handle_sends_nothing() {
let cluster = StubCluster::answering(Answers {
timeout_ms: 3000,
..Answers::default()
});
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let tx = client.attach_transaction(TX).expect("attaches");
assert_eq!(tx.detach(), TX, "detach must hand back the same id");
let settled = cluster.request_count();
std::thread::sleep(Duration::from_millis(1600));
assert!(
cluster.heads_for("abort_transaction").is_empty(),
"detaching an attached handle aborted: {:?}",
cluster.request_lines()
);
assert!(
cluster.heads_for("commit_transaction").is_empty(),
"detaching an attached handle committed: {:?}",
cluster.request_lines()
);
assert_eq!(
cluster.request_count(),
settled,
"something was sent after detach returned: {:?}",
cluster.request_lines()
);
}
#[test]
fn an_attached_handle_commits_like_an_owner() {
let cluster = StubCluster::answering(Answers {
timeout_ms: 30_000,
..Answers::default()
});
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let tx = client.attach_transaction(TX).expect("attaches");
tx.commit().expect("commits");
let commits = cluster.heads_for("commit_transaction");
assert_eq!(commits.len(), 1, "{:?}", cluster.request_lines());
assert_eq!(
str_param(&commits[0], "transaction_id").as_deref(),
Some(TX)
);
assert!(
param_of(&commits[0], "mutation_id").is_some(),
"a commit is not idempotent and must carry a mutation id:\n{}",
commits[0]
);
assert!(
cluster.heads_for("abort_transaction").is_empty(),
"an abort followed a successful commit: {:?}",
cluster.request_lines()
);
}
#[test]
fn attaching_to_a_transaction_that_is_gone_is_a_clear_error() {
let cluster = StubCluster::answering(Answers {
missing: true,
..Answers::default()
});
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
let error = client
.attach_transaction("0-0-0-1")
.expect_err("there is nothing to attach to");
let message = error.to_string();
for expected in ["attach", "0-0-0-1", "No such object"] {
assert!(
message.contains(expected),
"the error does not say {expected:?}: {message}"
);
}
std::thread::sleep(Duration::from_millis(200));
assert!(
cluster.heads_for("ping_transaction").is_empty(),
"a failed attach left a ping thread behind: {:?}",
cluster.request_lines()
);
}
#[test]
fn finishing_someone_elses_transaction_takes_only_the_id() {
let cluster = StubCluster::answering(Answers::default());
let client = Client::new(&cluster.url()).with_retries(RetryPolicy::none());
client.ping_transaction(TX).expect("pings");
client.commit_transaction(TX).expect("commits");
client.abort_transaction(TX).expect("aborts");
for command in [
"ping_transaction",
"commit_transaction",
"abort_transaction",
] {
let heads = cluster.heads_for(command);
assert_eq!(heads.len(), 1, "{command}: {:?}", cluster.request_lines());
assert!(
heads[0].starts_with(&format!("POST /api/v4/{command} ")),
"{command} used the wrong verb or path:\n{}",
heads[0]
);
assert_eq!(
str_param(&heads[0], "transaction_id").as_deref(),
Some(TX),
"{command} named the wrong transaction:\n{}",
heads[0]
);
}
let commit = &cluster.heads_for("commit_transaction")[0];
assert!(
param_of(commit, "mutation_id").is_some(),
"a by-id commit must ride under a mutation id:\n{commit}"
);
for freely in ["ping_transaction", "abort_transaction"] {
let head = &cluster.heads_for(freely)[0];
assert!(
param_of(head, "mutation_id").is_none(),
"{freely} is idempotent on its own and must carry no mutation id:\n{head}"
);
assert!(
param_of(head, "retry").is_none(),
"{freely} sent a retry flag with no mutation id to go with it:\n{head}"
);
}
}
struct Answers {
timeout_ms: i64,
ping_delay: Duration,
missing: bool,
ping_gone: bool,
}
impl Default for Answers {
fn default() -> Self {
Self {
timeout_ms: 30_000,
ping_delay: Duration::ZERO,
missing: false,
ping_gone: false,
}
}
}
impl Answers {
fn answer(&self, head: &str) -> Vec<u8> {
let path = head
.split_whitespace()
.nth(1)
.unwrap_or_default()
.to_owned();
match path.as_str() {
"/api/v4/start_transaction" => ok(format!(r#"{{"transaction_id"="{TX}"}}"#).as_bytes()),
"/api/v4/get" if self.missing => resolve_error(),
"/api/v4/get" => ok(format!(r#"{{"value"={}}}"#, self.timeout_ms).as_bytes()),
"/api/v4/ping_transaction" => {
std::thread::sleep(self.ping_delay);
if self.ping_gone {
return no_such_transaction();
}
ok(b"{}")
}
_ => ok(b"{}"),
}
}
}
struct StubCluster {
address: std::net::SocketAddr,
seen: Arc<Mutex<Vec<String>>>,
}
impl StubCluster {
fn answering(answers: Answers) -> Self {
let listener = TcpListener::bind("127.0.0.1:0").expect("binds");
let address = listener.local_addr().expect("has an address");
let seen = Arc::new(Mutex::new(Vec::new()));
let served = Arc::clone(&seen);
let answers = Arc::new(answers);
std::thread::spawn(move || {
for stream in listener.incoming() {
let Ok(stream) = stream else { return };
let answers = Arc::clone(&answers);
let seen = Arc::clone(&served);
std::thread::spawn(move || serve(stream, &answers, &seen));
}
});
Self { address, seen }
}
fn url(&self) -> String {
format!("http://{}", self.address)
}
fn heads_for(&self, what: &str) -> Vec<String> {
self.seen
.lock()
.expect("nothing panicked holding it")
.iter()
.filter(|head| head.lines().next().is_some_and(|line| line.contains(what)))
.cloned()
.collect()
}
fn request_count(&self) -> usize {
self.seen.lock().expect("nothing panicked holding it").len()
}
fn request_lines(&self) -> Vec<String> {
self.seen
.lock()
.expect("nothing panicked holding it")
.iter()
.map(|head| head.lines().next().unwrap_or_default().to_owned())
.collect()
}
}
fn serve(mut stream: std::net::TcpStream, answers: &Answers, seen: &Arc<Mutex<Vec<String>>>) {
let mut reader = BufReader::new(stream.try_clone().expect("clones"));
loop {
let mut head = String::new();
loop {
let mut line = String::new();
match reader.read_line(&mut line) {
Ok(0) => return,
Ok(_) if line == "\r\n" => break,
Ok(_) => head.push_str(&line),
Err(_) => return,
}
}
if head.is_empty() {
return;
}
if let Some(length) = content_length(&head) {
let mut body = vec![0_u8; length];
if reader.read_exact(&mut body).is_err() {
return;
}
} else if head.to_lowercase().contains("transfer-encoding: chunked") {
drain_chunked(&mut reader);
}
seen.lock()
.expect("nothing panicked holding it")
.push(head.clone());
let answer = answers.answer(&head);
if stream.write_all(&answer).is_err() {
return;
}
stream.flush().ok();
}
}
fn ok(body: &[u8]) -> Vec<u8> {
let mut reply = format!(
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nContent-Type: application/x-yt-yson-text\r\n\r\n",
body.len()
)
.into_bytes();
reply.extend_from_slice(body);
reply
}
fn resolve_error() -> Vec<u8> {
let document = r#"{"code":500,"message":"Error resolving path #0-0-0-1/@timeout","inner_errors":[{"code":500,"message":"No such object 0-0-0-1","attributes":{"missing_object_id":"0-0-0-1"}}]}"#;
format!("HTTP/1.1 200 OK\r\nX-YT-Error: {document}\r\nContent-Length: 0\r\n\r\n").into_bytes()
}
fn no_such_transaction() -> Vec<u8> {
let document = format!(
r#"{{"code":11000,"message":"No such transaction {TX}","attributes":{{"transaction_id":"{TX}"}}}}"#
);
format!("HTTP/1.1 200 OK\r\nX-YT-Error: {document}\r\nContent-Length: 0\r\n\r\n").into_bytes()
}
fn content_length(head: &str) -> Option<usize> {
head.lines()
.find(|line| line.to_lowercase().starts_with("content-length:"))
.and_then(|line| line.split_once(':'))
.and_then(|(_, value)| value.trim().parse().ok())
}
fn drain_chunked(reader: &mut BufReader<std::net::TcpStream>) {
loop {
let mut header = String::new();
if reader.read_line(&mut header).is_err() {
return;
}
let size = usize::from_str_radix(header.trim(), 16).unwrap_or(0);
let mut chunk = vec![0_u8; size + 2];
if reader.read_exact(&mut chunk).is_err() || size == 0 {
return;
}
}
}
fn params_of(head: &str) -> YsonValue {
let line = head
.lines()
.find(|line| {
line.split_once(':')
.is_some_and(|(name, _)| name.eq_ignore_ascii_case("x-yt-parameters"))
})
.unwrap_or_else(|| panic!("no X-YT-Parameters header in:\n{head}"));
let value = line
.split_once(':')
.expect("the header has a value")
.1
.trim();
from_slice(value.as_bytes(), YsonFormat::Text)
.unwrap_or_else(|e| panic!("parameters are not text YSON ({e}): {value}"))
}
fn param_of(head: &str, key: &str) -> Option<YsonValue> {
match params_of(head).node {
YsonNode::Map(mut m) => m.remove(key.as_bytes()),
_ => None,
}
}
fn str_param(head: &str, key: &str) -> Option<String> {
match param_of(head, key)?.node {
YsonNode::String(bytes) => Some(String::from_utf8_lossy(&bytes).into_owned()),
_ => None,
}
}