open-standards/wai/wai-quantum-service/tests/restart.rs

//! The service killed mid-job and started again on the same store: every
//! result it gave is served again with its receipt, which still verifies;
//! the job it was running runs again, once; the one queued behind it is
//! queued again. Killed while that job runs again, the job comes back
//! interrupted, with the reason, and the queue goes on.

use serde_json::{Value, json};
use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};
use std::path::Path;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};

mod common;

use common::Running;

fn free_port() -> u16 {
    TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port()
}

/// One HTTP/1.1 exchange: the status and the body.
fn http(port: u16, method: &str, path: &str, body: Option<&Value>) -> (u16, String) {
    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(120))).unwrap();
    let text = body.map(Value::to_string).unwrap_or_default();
    write!(s, "{method} {path} HTTP/1.1\r\nHost: 127.0.0.1\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{text}", text.len()).unwrap();
    let mut out = String::new();
    s.read_to_string(&mut out).unwrap();
    let status = out.split_whitespace().nth(1).and_then(|c| c.parse().ok()).unwrap_or(0);
    let body = out.split_once("\r\n\r\n").map(|(_, b)| b.to_string()).unwrap_or_default();
    (status, body)
}

fn json_of(port: u16, method: &str, path: &str, body: Option<&Value>) -> (u16, Value) {
    let (code, text) = http(port, method, path, body);
    (code, serde_json::from_str(&text).unwrap_or_else(|e| panic!("{method} {path}: {e}: {text}")))
}

/// Whether the service on `port` answers its health check. A connection refused,
/// reset or cut short is a no, not a failure of the test: see [`start`].
fn healthy(port: u16) -> bool {
    let Ok(mut s) = TcpStream::connect(("127.0.0.1", port)) else { return false };
    let _ = s.set_read_timeout(Some(Duration::from_secs(10)));
    let mut out = String::new();
    write!(s, "GET /v1/health HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n").is_ok()
        && s.read_to_string(&mut out).is_ok()
        && out.split_whitespace().nth(1) == Some("200")
}

/// The service, serving the store at `dir`, once it answers; started again if it
/// exits before it answers. A service killed a moment ago can leave the store's
/// lock and its listening port held for a moment after it is reaped: on a loaded
/// host the next service then refused the store ("another process holds the job
/// store") and exited, as it must, while the port took the health check's
/// connection and reset it.
fn start(dir: &Path, port: u16) -> Running {
    let t = Instant::now();
    loop {
        let mut child = Running::new(Command::new(env!("CARGO_BIN_EXE_wai-quantum-service"))
            .arg("serve")
            .env("WAI_QSVC_ADDR", format!("127.0.0.1:{port}"))
            .env("WAI_QSVC_STORE", dir.join("store"))
            .env("WAI_QSVC_KEY", dir.join("signing.key"))
            .env("WAI_QSVC_JOB_GB", "2")
            .env("WAI_QSVC_BUDGET_GB", "4")
            .env("WAI_QSVC_WORKERS", "1")
            .env("WAI_QSVC_ENERGY", "0")
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .spawn()
            .expect("the service starts"));
        loop {
            if healthy(port) {
                return child;
            }
            if let Some(status) = child.try_wait().unwrap() {
                eprintln!("the service exited before it answered ({status}); starting it again");
                break;
            }
            assert!(t.elapsed() < Duration::from_secs(60), "the service did not answer");
            std::thread::sleep(Duration::from_millis(50));
        }
        assert!(t.elapsed() < Duration::from_secs(60), "the service did not start");
        std::thread::sleep(Duration::from_millis(100));
    }
}

/// Waits until a job's status is `want`.
fn until(port: u16, id: &str, want: &str) -> Value {
    let t = Instant::now();
    loop {
        let (_, j) = json_of(port, "GET", &format!("/v1/jobs/{id}"), None);
        if j["status"] == want {
            return j;
        }
        assert!(t.elapsed() < Duration::from_secs(120), "{id} never {want}: {j}");
        std::thread::sleep(Duration::from_millis(25));
    }
}

#[test]
fn a_service_killed_mid_job_serves_its_results_again_and_says_what_it_interrupted() {
    let dir = std::env::temp_dir().join(format!("wai-qsvc-restart-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&dir);
    std::fs::create_dir_all(&dir).unwrap();
    let port = free_port();
    let mut svc = start(&dir, port);

    // A job and a set, finished.
    let bell = json!({ "qasm": "OPENQASM 3.0; qubit[2] q; bit[2] c; h q[0]; cx q[0], q[1]; c = measure q;", "shots": 200, "seed": 3 });
    let (code, done) = json_of(port, "POST", "/v1/run", Some(&bell));
    assert_eq!((code, done["status"].as_str()), (200, Some("done")), "{done}");
    let set = json!({ "qasm": "OPENQASM 3.0;\ninput float theta;\nqubit q;\nrx(theta) q;\n", "sweep": { "theta": [0.1, 0.7] }, "observables": ["Z0"] });
    let (code, set_done) = json_of(port, "POST", "/v1/run", Some(&set));
    assert_eq!((code, set_done["status"].as_str()), (200, Some("done")), "{set_done}");

    // A long job running, and one queued behind it (one worker).
    let mut long = String::from("OPENQASM 3.0;\nqubit[25] q;\nh q;\n");
    for k in 0..4000 {
        long += &format!("rx(0.{}) q[{}];\ncx q[{}], q[{}];\n", k % 9 + 1, k % 25, k % 25, (k + 1) % 25);
    }
    let long = json!({ "qasm": long, "observables": ["Z0"] });
    let (code, running) = json_of(port, "POST", "/v1/jobs", Some(&long));
    assert_eq!(code, 202, "{running}");
    let running = running["id"].as_str().unwrap().to_string();
    let (_, queued) = json_of(port, "POST", "/v1/jobs", Some(&bell));
    let queued = queued["id"].as_str().unwrap().to_string();
    until(port, &running, "running");
    assert_eq!(json_of(port, "GET", &format!("/v1/jobs/{queued}"), None).1["status"], "queued");

    // Killed, as a crash or a power cut would stop it, and started again.
    svc.kill().unwrap();
    svc.wait().unwrap();
    let mut svc = start(&dir, port);

    // The finished job, as it was, its receipt still verifying, here and by the service.
    let (code, again) = json_of(port, "GET", &format!("/v1/jobs/{}", done["id"].as_str().unwrap()), None);
    assert_eq!((code, again["status"].as_str()), (200, Some("done")), "{again}");
    assert_eq!(again["result"], done["result"], "served as it was served before");
    assert!(wai_quantum::engine::job::verify(&json!({ "request": again["request"], "result": again["result"] })).unwrap().0);
    let (_, v) = json_of(port, "POST", "/v1/verify", Some(&json!({ "request": again["request"], "result": again["result"] })));
    assert_eq!(v["valid"], true, "{v}");
    let (_, set_again) = json_of(port, "GET", &format!("/v1/jobs/{}", set_done["id"].as_str().unwrap()), None);
    assert_eq!(wai_quantum::engine::sets::verify(&set_again)["valid"], true, "{set_again}");
    assert_eq!(set_again["result"]["items"][1]["point"], json!({ "theta": 0.7 }));

    // The description says what was found: the running job run again, the
    // queued one queued again.
    let restored = |port: u16| {
        let (_, d) = json_of(port, "GET", "/v1", None);
        assert_eq!(d["store"]["durable"], true, "{}", d["store"]);
        let r = &d["store"]["restored"];
        (r["finished"].as_u64(), r["interrupted"].as_u64(), r["requeued"].as_u64(), r["rerun"].as_u64())
    };
    assert_eq!(restored(port), (Some(2), Some(0), Some(1), Some(1)));

    // The job it was running runs again from its request, first in its
    // order, saying when the service stopped.
    let r = until(port, &running, "running");
    assert_eq!(r["interruptions"].as_array().map(Vec::len), Some(1), "{r}");
    assert!(r.get("result").is_none() && r.get("error").is_none() && r["request"] == long, "{r}");
    assert_eq!(json_of(port, "GET", &format!("/v1/jobs/{queued}"), None).1["status"], "queued");

    // Killed again while it runs again: interrupted, with the reason, its
    // request kept, and not run a third time.
    svc.kill().unwrap();
    svc.wait().unwrap();
    let mut svc = start(&dir, port);
    assert_eq!(restored(port), (Some(2), Some(1), Some(1), Some(0)));
    let (_, r) = json_of(port, "GET", &format!("/v1/jobs/{running}"), None);
    assert_eq!(r["status"], "interrupted", "{r}");
    assert!(r["error"].as_str().unwrap().contains("the service stopped while this job ran, each of the 2 times it ran"), "{r}");
    assert_eq!(r["interruptions"].as_array().map(Vec::len), Some(2), "{r}");
    assert!(r.get("result").is_none() && r["request"] == long);
    // The one queued behind it runs now.
    let q = until(port, &queued, "done");
    assert!(wai_quantum::engine::job::verify(&json!({ "request": q["request"], "result": q["result"] })).unwrap().0);
    assert_eq!(json_of(port, "GET", &format!("/v1/jobs/{running}"), None).1["status"], "interrupted");

    svc.kill().unwrap();
    svc.wait().unwrap();
    let _ = std::fs::remove_dir_all(&dir);
}