#![cfg(unix)]
use rch_common::e2e::{
DaemonConfigFixture, HarnessResult, TestHarness, TestHarnessBuilder, WorkerFixture,
WorkersFixture,
};
use serial_test::serial;
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Write};
use std::os::unix::net::UnixStream;
use std::path::PathBuf;
use std::time::Duration;
fn create_multi_worker_harness(test_name: &str) -> HarnessResult<TestHarness> {
let target_dir = if let Ok(target_dir) = std::env::var("CARGO_TARGET_DIR") {
PathBuf::from(target_dir).join("debug")
} else {
let manifest_dir = std::env::var("CARGO_MANIFEST_DIR")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from("."));
let project_root = manifest_dir.parent().unwrap_or(&manifest_dir);
project_root.join("target/debug")
};
TestHarnessBuilder::new(test_name)
.cleanup_on_success(true)
.cleanup_on_failure(false)
.default_timeout(Duration::from_secs(60))
.rchd_binary(target_dir.join("rchd"))
.rch_binary(target_dir.join("rch"))
.rch_wkr_binary(target_dir.join("rch-wkr"))
.env("RCH_MOCK_SSH", "1")
.env("RCH_MOCK_SSH_STDOUT", "health_check\n")
.build()
}
fn setup_daemon_with_workers(
harness: &TestHarness,
workers: &WorkersFixture,
) -> HarnessResult<PathBuf> {
let socket_path = harness.test_dir().join("rch.sock");
let daemon_config = DaemonConfigFixture::minimal(&socket_path);
harness.create_daemon_config(&daemon_config.to_toml())?;
harness.create_workers_config(&workers.to_toml())?;
Ok(socket_path)
}
fn start_daemon_with_socket(
harness: &TestHarness,
socket_path: &std::path::Path,
extra_args: &[&str],
) -> HarnessResult<u32> {
let socket_str = socket_path.to_string_lossy();
let mut args = vec![
"--foreground",
"--metrics-port",
"0",
"--socket",
&*socket_str,
];
args.extend(extra_args);
harness.start_daemon(&args)
}
fn send_socket_request(socket_path: &std::path::Path, request: &str) -> std::io::Result<String> {
let mut last_error = None;
for attempt in 0..5 {
let result: std::io::Result<String> = (|| {
let mut stream = UnixStream::connect(socket_path)?;
stream.set_read_timeout(Some(Duration::from_secs(5)))?;
stream.set_write_timeout(Some(Duration::from_secs(5)))?;
writeln!(stream, "{}", request)?;
stream.flush()?;
let mut reader = BufReader::new(stream);
let mut response = String::new();
loop {
let mut line = String::new();
match reader.read_line(&mut line) {
Ok(0) => break,
Ok(_) => response.push_str(&line),
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
Err(e) => return Err(e),
}
}
Ok(response)
})();
match result {
Ok(response) if !response.is_empty() => return Ok(response),
Ok(_) => {}
Err(err) => last_error = Some(err),
}
if attempt < 4 {
std::thread::sleep(Duration::from_millis(50));
}
}
Err(last_error.unwrap_or_else(|| {
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "empty daemon response")
}))
}
fn extract_json_body(response: &str) -> Option<&str> {
if let Some(pos) = response.find("\r\n\r\n") {
Some(&response[pos + 4..])
} else if let Some(pos) = response.find("\n\n") {
Some(&response[pos + 2..])
} else {
None
}
}
fn create_worker(id: &str, host: &str, slots: u32, priority: u32) -> WorkerFixture {
WorkerFixture {
id: id.to_string(),
host: host.to_string(),
user: current_user(),
identity_file: "~/.ssh/id_rsa".to_string(),
total_slots: slots,
priority,
}
}
fn current_user() -> String {
std::env::var("USER")
.or_else(|_| std::env::var("LOGNAME"))
.unwrap_or_else(|_| "user".to_string())
}
fn parse_selection_response(body: &str) -> Option<serde_json::Value> {
serde_json::from_str(body).ok()
}
fn selected_worker(response: &str) -> rch_common::SelectedWorker {
assert!(
response.starts_with("HTTP/1.0 200 OK\r\n"),
"Unexpected selection HTTP response: {response}"
);
let body = extract_json_body(response).expect("selection response has an HTTP body");
let parsed = serde_json::from_str::<rch_common::SelectionResponse>(body);
assert!(
parsed.is_ok(),
"Invalid selection response: {body}; error={:?}",
parsed.as_ref().err()
);
let selection = parsed.expect("validated selection response");
assert!(
selection.worker.is_some(),
"Expected an admitted worker for every sample: {body}"
);
selection.worker.expect("validated worker selection")
}
#[test]
#[serial]
fn test_load_balance_distribution() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("load_balance_distribution").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_load_balance_distribution");
let worker1 = create_worker("worker-1", "localhost", 8, 100);
let worker2 = create_worker("worker-2", "localhost", 4, 100);
let worker3 = create_worker("worker-3", "localhost", 8, 100);
let workers = WorkersFixture::empty()
.add_worker(worker1)
.add_worker(worker2)
.add_worker(worker3);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
harness.logger.info(
"[e2e::multi_worker] Fleet: worker-1 (8 slots), worker-2 (4 slots), worker-3 (8 slots)",
);
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let mut distribution: HashMap<String, u32> = HashMap::new();
let num_selections = 20;
for i in 0..num_selections {
let response = send_socket_request(
&socket_path,
&format!("GET /select-worker?project=load-test-{}&cores=1", i),
)
.unwrap();
let worker = selected_worker(&response);
*distribution.entry(worker.id.to_string()).or_insert(0) += 1;
}
harness.logger.info(format!(
"[e2e::multi_worker] Distribution: {:?}",
distribution
));
let total: u32 = distribution.values().sum();
harness.logger.info(format!(
"[e2e::multi_worker] Total selections: {} (expected: {})",
total, num_selections
));
assert_eq!(total, num_selections, "every request must select a worker");
let workers_with_selections = distribution.len();
assert!(
workers_with_selections >= 2,
"Expected at least 2 workers to receive selections, got {}",
workers_with_selections
);
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_load_balance_distribution");
harness.mark_passed();
}
#[test]
#[serial]
fn test_worker_prioritization() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("worker_prioritization").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_worker_prioritization");
let high_priority = create_worker("high-priority", "localhost", 64, 200);
let low_priority = create_worker("low-priority", "localhost", 64, 50);
let workers = WorkersFixture::empty()
.add_worker(high_priority)
.add_worker(low_priority);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
harness.logger.info(
"[e2e::multi_worker] WORKERS: [high-priority (slots=64, priority=200), low-priority (slots=64, priority=50)]",
);
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let mut high_priority_count = 0;
let num_selections = 50;
for i in 0..num_selections {
let response = send_socket_request(
&socket_path,
&format!("GET /select-worker?project=priority-test-{}&cores=1", i),
)
.unwrap();
let worker = selected_worker(&response);
if worker.id.as_str() == "high-priority" {
high_priority_count += 1;
}
}
harness.logger.info(format!(
"[e2e::multi_worker] High-priority selected: {}/{}",
high_priority_count, num_selections
));
let min_expected = num_selections * 2 / 5; assert!(
high_priority_count >= min_expected,
"Expected high-priority worker to be selected at least 40% of the time ({}/{}), got {}/{}",
min_expected,
num_selections,
high_priority_count,
num_selections
);
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_worker_prioritization");
harness.mark_passed();
}
#[test]
#[serial]
fn test_cached_project_locality() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("cached_project_locality").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_cached_project_locality");
let worker1 = create_worker("worker-1", "localhost", 4, 100);
let worker2 = create_worker("worker-2", "localhost", 4, 100);
let worker3 = create_worker("worker-3", "localhost", 4, 100);
let workers = WorkersFixture::empty()
.add_worker(worker1)
.add_worker(worker2)
.add_worker(worker3);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
harness
.logger
.info("[e2e::multi_worker] WORKERS: 3 workers with equal priority and slots");
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let response = send_socket_request(
&socket_path,
"GET /select-worker?project=sticky-project&cores=1",
)
.unwrap();
let first_worker = extract_json_body(&response)
.and_then(parse_selection_response)
.and_then(|json| {
json.get("worker")
.and_then(|w| w.get("id"))
.and_then(|id| id.as_str())
.map(|s| s.to_string())
});
harness.logger.info(format!(
"[e2e::multi_worker] First selection: {:?}",
first_worker
));
let mut same_worker_count = 0;
let num_repeats = 5;
for _ in 0..num_repeats {
let response = send_socket_request(
&socket_path,
"GET /select-worker?project=sticky-project&cores=1",
)
.unwrap();
if let Some(body) = extract_json_body(&response)
&& let Some(json) = parse_selection_response(body)
&& let Some(worker_id) = json
.get("worker")
.and_then(|w| w.get("id"))
.and_then(|id| id.as_str())
&& first_worker.as_deref() == Some(worker_id)
{
same_worker_count += 1;
}
}
harness.logger.info(format!(
"[e2e::multi_worker] Same worker selected: {}/{}",
same_worker_count, num_repeats
));
harness.logger.info(format!(
"[e2e::multi_worker] Cache locality rate: {}%",
(same_worker_count * 100) / num_repeats
));
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_cached_project_locality");
harness.mark_passed();
}
#[test]
#[serial]
fn test_heterogeneous_workers() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("heterogeneous_workers").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_heterogeneous_workers");
let rust_worker = create_worker("rust-worker", "localhost", 4, 100);
let bun_worker = create_worker("bun-worker", "localhost", 4, 100);
let workers = WorkersFixture::empty()
.add_worker(rust_worker)
.add_worker(bun_worker);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
harness
.logger
.info("[e2e::multi_worker] WORKERS: rust-worker, bun-worker");
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let response = send_socket_request(
&socket_path,
"GET /select-worker?project=rust-project&cores=1&runtime=rust",
)
.unwrap();
assert!(response.contains("200 OK"), "Got: {}", response);
let body = extract_json_body(&response).unwrap();
harness
.logger
.info(format!("[e2e::multi_worker] Rust selection: {}", body));
assert!(body.contains("reason"), "Expected reason in response");
let response = send_socket_request(
&socket_path,
"GET /select-worker?project=bun-project&cores=1&runtime=bun",
)
.unwrap();
let body = extract_json_body(&response).unwrap();
harness
.logger
.info(format!("[e2e::multi_worker] Bun selection: {}", body));
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_heterogeneous_workers");
harness.mark_passed();
}
#[test]
#[serial]
fn test_worker_registration() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("worker_registration").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_worker_registration");
let workers = WorkersFixture::mock_local(1);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let response = send_socket_request(&socket_path, "GET /status").unwrap();
assert!(response.contains("200 OK"), "Got: {}", response);
let body = extract_json_body(&response).unwrap();
harness
.logger
.info(format!("[e2e::multi_worker] Initial status: {}", body));
assert!(body.contains("workers"), "Expected workers in status");
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_worker_registration");
harness.mark_passed();
}
#[test]
#[serial]
fn test_worker_health_monitoring() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("worker_health_monitoring").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_worker_health_monitoring");
let healthy_worker = create_worker("healthy-worker", "localhost", 4, 100);
let unhealthy_worker = create_worker("unhealthy-worker", "unreachable.invalid", 4, 100);
let workers = WorkersFixture::empty()
.add_worker(healthy_worker)
.add_worker(unhealthy_worker);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let response = send_socket_request(&socket_path, "GET /status").unwrap();
assert!(response.contains("200 OK"), "Got: {}", response);
let body = extract_json_body(&response).unwrap();
harness
.logger
.info(format!("[e2e::multi_worker] Health status: {}", body));
assert!(body.contains("workers"), "Expected workers in status");
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_worker_health_monitoring");
harness.mark_passed();
}
#[test]
#[serial]
fn test_selection_avoids_unhealthy() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("selection_avoids_unhealthy").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_selection_avoids_unhealthy");
let healthy_worker = create_worker("healthy-worker", "localhost", 4, 100);
let workers = WorkersFixture::empty().add_worker(healthy_worker);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let response = send_socket_request(
&socket_path,
"GET /select-worker?project=health-test&cores=1",
)
.unwrap();
assert!(response.contains("200 OK"), "Got: {}", response);
let body = extract_json_body(&response).unwrap();
harness
.logger
.info(format!("[e2e::multi_worker] Selection result: {}", body));
if let Some(json) = parse_selection_response(body)
&& let Some(worker) = json.get("worker")
{
harness
.logger
.info(format!("[e2e::multi_worker] Selected worker: {}", worker));
}
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_selection_avoids_unhealthy");
harness.mark_passed();
}
#[test]
#[serial]
fn test_worker_graceful_shutdown() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("worker_graceful_shutdown").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_worker_graceful_shutdown");
let workers = WorkersFixture::mock_local(2);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let response = send_socket_request(&socket_path, "GET /status").unwrap();
assert!(response.contains("200 OK"), "Got: {}", response);
let body = extract_json_body(&response).unwrap();
harness
.logger
.info(format!("[e2e::multi_worker] Initial status: {}", body));
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_worker_graceful_shutdown");
harness.mark_passed();
}
#[test]
#[serial]
fn test_fleet_single_worker_fallback() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("fleet_single_worker_fallback").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_fleet_single_worker_fallback");
let workers = WorkersFixture::mock_local(1);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let response = send_socket_request(
&socket_path,
"GET /select-worker?project=single-fallback&cores=1",
)
.unwrap();
assert!(response.contains("200 OK"), "Got: {}", response);
let body = extract_json_body(&response).unwrap();
harness.logger.info(format!(
"[e2e::multi_worker] Single worker selection: {}",
body
));
assert!(body.contains("reason"), "Expected reason in response");
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_fleet_single_worker_fallback");
harness.mark_passed();
}
#[test]
#[serial]
fn test_no_workers_available() {
let _guard = rch_common::test_guard!();
let harness = create_multi_worker_harness("no_workers_available").unwrap();
harness
.logger
.info("[e2e::multi_worker] TEST START: test_no_workers_available");
let busy_worker = WorkerFixture {
id: "busy-worker".to_string(),
host: "localhost".to_string(),
user: current_user(),
identity_file: "~/.ssh/id_rsa".to_string(),
total_slots: 0,
priority: 100,
};
let workers = WorkersFixture::empty().add_worker(busy_worker);
let socket_path = setup_daemon_with_workers(&harness, &workers).unwrap();
start_daemon_with_socket(&harness, &socket_path, &[]).unwrap();
harness
.wait_for_socket(&socket_path, Duration::from_secs(10))
.unwrap();
let response = send_socket_request(
&socket_path,
"GET /select-worker?project=no-workers&cores=1",
)
.unwrap();
assert!(response.contains("200 OK"), "Got: {}", response);
let body = extract_json_body(&response).unwrap();
harness
.logger
.info(format!("[e2e::multi_worker] No workers result: {}", body));
assert!(
body.contains("reason") || body.contains("null"),
"Expected reason or null worker"
);
harness
.logger
.info("[e2e::multi_worker] TEST PASS: test_no_workers_available");
harness.mark_passed();
}