|
| 1 | +use toydb::raft::NodeID; |
| 2 | +use toydb::Client; |
| 3 | + |
| 4 | +use rand::Rng; |
| 5 | +use std::collections::BTreeMap; |
| 6 | +use std::error::Error; |
| 7 | +use std::fmt::Write as _; |
| 8 | +use std::path::Path; |
| 9 | +use std::time::Duration; |
| 10 | + |
| 11 | +/// Timeout for command responses and node readiness. |
| 12 | +const TIMEOUT: Duration = Duration::from_secs(5); |
| 13 | + |
| 14 | +/// The base SQL port (+id). |
| 15 | +const SQL_BASE_PORT: u16 = 19600; |
| 16 | + |
| 17 | +/// The base Raft port (+id). |
| 18 | +const RAFT_BASE_PORT: u16 = 19700; |
| 19 | + |
| 20 | +/// Runs a toyDB cluster using the built binary in a temporary directory. The |
| 21 | +/// cluster will be killed and removed when dropped. |
| 22 | +/// |
| 23 | +/// This runs the cluster as child processes using the built binary instead of |
| 24 | +/// spawning in-memory threads for a couple of reasons: it avoids having to |
| 25 | +/// gracefully shut down the server (which is complicated by e.g. |
| 26 | +/// TcpListener::accept() not being interruptable), and it tests the entire |
| 27 | +/// server (and eventually the toySQL client) end-to-end. |
| 28 | +pub struct TestCluster { |
| 29 | + servers: BTreeMap<NodeID, TestServer>, |
| 30 | + #[allow(dead_code)] |
| 31 | + dir: tempfile::TempDir, // deleted when dropped |
| 32 | +} |
| 33 | + |
| 34 | +type NodePorts = BTreeMap<NodeID, (u16, u16)>; // raft,sql on localhost |
| 35 | + |
| 36 | +impl TestCluster { |
| 37 | + /// Runs and returns a test cluster. It keeps running until dropped. |
| 38 | + pub fn run(nodes: u8) -> Result<Self, Box<dyn Error>> { |
| 39 | + // Create temporary directory. |
| 40 | + let dir = tempfile::TempDir::with_prefix("toydb")?; |
| 41 | + |
| 42 | + // Allocate port numbers for nodes. |
| 43 | + let ports: NodePorts = (1..=nodes) |
| 44 | + .map(|id| (id, (RAFT_BASE_PORT + id as u16, SQL_BASE_PORT + id as u16))) |
| 45 | + .collect(); |
| 46 | + |
| 47 | + // Start nodes. |
| 48 | + let mut servers = BTreeMap::new(); |
| 49 | + for id in 1..=nodes { |
| 50 | + let dir = dir.path().join(format!("toydb{id}")); |
| 51 | + servers.insert(id, TestServer::run(id, &dir, &ports)?); |
| 52 | + } |
| 53 | + |
| 54 | + // Wait for the nodes to be ready, by fetching the server status. |
| 55 | + let started = std::time::Instant::now(); |
| 56 | + for server in servers.values_mut() { |
| 57 | + while let Err(error) = server.connect().and_then(|mut c| Ok(c.status()?)) { |
| 58 | + server.assert_alive(); |
| 59 | + if started.elapsed() >= TIMEOUT { |
| 60 | + return Err(error); |
| 61 | + } |
| 62 | + std::thread::sleep(Duration::from_millis(200)); |
| 63 | + } |
| 64 | + } |
| 65 | + |
| 66 | + Ok(Self { servers, dir }) |
| 67 | + } |
| 68 | + |
| 69 | + /// Connects to a random cluster node using the regular client. |
| 70 | + #[allow(dead_code)] |
| 71 | + pub fn connect(&self) -> Result<Client, Box<dyn Error>> { |
| 72 | + let id = rand::thread_rng().gen_range(1..=self.servers.len()) as NodeID; |
| 73 | + self.servers.get(&id).unwrap().connect() |
| 74 | + } |
| 75 | + |
| 76 | + /// Connects to a random cluster node using the toysql binary. |
| 77 | + pub fn connect_toysql(&self) -> Result<TestClient, Box<dyn Error>> { |
| 78 | + let id = rand::thread_rng().gen_range(1..=self.servers.len()) as NodeID; |
| 79 | + self.servers.get(&id).unwrap().connect_toysql() |
| 80 | + } |
| 81 | +} |
| 82 | + |
| 83 | +/// A toyDB server. |
| 84 | +pub struct TestServer { |
| 85 | + id: NodeID, |
| 86 | + child: std::process::Child, |
| 87 | + sql_port: u16, |
| 88 | +} |
| 89 | + |
| 90 | +impl TestServer { |
| 91 | + /// Runs a toyDB server. |
| 92 | + fn run(id: NodeID, dir: &Path, ports: &NodePorts) -> Result<Self, Box<dyn Error>> { |
| 93 | + // Build and write the configuration file. |
| 94 | + let configfile = dir.join("toydb.yaml"); |
| 95 | + std::fs::create_dir_all(dir)?; |
| 96 | + std::fs::write(&configfile, Self::build_config(id, dir, ports)?)?; |
| 97 | + |
| 98 | + // Build the binary. |
| 99 | + // TODO: this may contribute to slow start times, consider building once |
| 100 | + // and passing it in. |
| 101 | + let build = escargot::CargoBuild::new().bin("toydb").run()?; |
| 102 | + |
| 103 | + // Spawn process. Discard output. |
| 104 | + let child = build |
| 105 | + .command() |
| 106 | + .args(["-c", &configfile.to_string_lossy()]) |
| 107 | + .stdout(std::process::Stdio::null()) |
| 108 | + .stderr(std::process::Stdio::null()) |
| 109 | + .spawn()?; |
| 110 | + |
| 111 | + let (_, sql_port) = ports.get(&id).copied().expect("node not in ports"); |
| 112 | + Ok(Self { id, child, sql_port }) |
| 113 | + } |
| 114 | + |
| 115 | + /// Generates a config file for the given node. |
| 116 | + fn build_config(id: NodeID, dir: &Path, ports: &NodePorts) -> Result<String, Box<dyn Error>> { |
| 117 | + let (raft_port, sql_port) = ports.get(&id).expect("node not in ports"); |
| 118 | + let mut cfg = String::new(); |
| 119 | + writeln!(cfg, "id: {id}")?; |
| 120 | + writeln!(cfg, "data_dir: {}", dir.to_string_lossy())?; |
| 121 | + writeln!(cfg, "listen_raft: localhost:{raft_port}")?; |
| 122 | + writeln!(cfg, "listen_sql: localhost:{sql_port}")?; |
| 123 | + writeln!(cfg, "peers: {{")?; |
| 124 | + for (peer_id, (peer_raft_port, _)) in ports.iter().filter(|(peer, _)| **peer != id) { |
| 125 | + writeln!(cfg, " '{peer_id}': localhost:{peer_raft_port},")?; |
| 126 | + } |
| 127 | + writeln!(cfg, "}}")?; |
| 128 | + Ok(cfg) |
| 129 | + } |
| 130 | + |
| 131 | + /// Asserts that the server is still running. |
| 132 | + fn assert_alive(&mut self) { |
| 133 | + if let Some(status) = self.child.try_wait().expect("failed to check exit status") { |
| 134 | + panic!("node {id} exited with status {status}", id = self.id) |
| 135 | + } |
| 136 | + } |
| 137 | + |
| 138 | + /// Connects to the server using a regular client. |
| 139 | + fn connect(&self) -> Result<Client, Box<dyn Error>> { |
| 140 | + Ok(Client::connect(("localhost", self.sql_port))?) |
| 141 | + } |
| 142 | + |
| 143 | + /// Connects to the server using the toysql binary. |
| 144 | + pub fn connect_toysql(&self) -> Result<TestClient, Box<dyn Error>> { |
| 145 | + TestClient::connect(self.sql_port) |
| 146 | + } |
| 147 | +} |
| 148 | + |
| 149 | +impl Drop for TestServer { |
| 150 | + // Kills the child process when dropped. |
| 151 | + fn drop(&mut self) { |
| 152 | + self.child.kill().expect("failed to kill node"); |
| 153 | + self.child.wait().expect("failed to wait for node to terminate"); |
| 154 | + } |
| 155 | +} |
| 156 | + |
| 157 | +/// A toySQL client using the toysql binary. |
| 158 | +pub struct TestClient { |
| 159 | + session: rexpect::session::PtySession, |
| 160 | +} |
| 161 | + |
| 162 | +impl TestClient { |
| 163 | + /// Connects to a toyDB server at the given SQL port number, using |
| 164 | + /// the toysql binary. |
| 165 | + fn connect(port: u16) -> Result<Self, Box<dyn Error>> { |
| 166 | + // Build the binary. |
| 167 | + let build = escargot::CargoBuild::new().bin("toysql").run()?; |
| 168 | + |
| 169 | + // Run it, using rexpect to manage stdin/stdout. |
| 170 | + let mut command = build.command(); |
| 171 | + command.args(["-p", &port.to_string()]); |
| 172 | + let session = rexpect::spawn_with_options( |
| 173 | + command, |
| 174 | + rexpect::reader::Options { |
| 175 | + timeout_ms: Some(TIMEOUT.as_millis() as u64), |
| 176 | + strip_ansi_escape_codes: true, |
| 177 | + }, |
| 178 | + )?; |
| 179 | + |
| 180 | + // Wait for the initial prompt. |
| 181 | + let mut client = Self { session }; |
| 182 | + client.read_until_prompt()?; |
| 183 | + Ok(client) |
| 184 | + } |
| 185 | + |
| 186 | + /// Executes a command, returning it and the resulting toysql prompt. |
| 187 | + pub fn execute(&mut self, command: &str) -> Result<(String, String), Box<dyn Error>> { |
| 188 | + let mut command = command.to_string(); |
| 189 | + if !command.ends_with(';') && !command.starts_with('!') { |
| 190 | + command = format!("{command};"); |
| 191 | + } |
| 192 | + self.session.send_line(&command)?; |
| 193 | + self.session.exp_string(&command)?; // wait for echo |
| 194 | + self.read_until_prompt() |
| 195 | + } |
| 196 | + |
| 197 | + /// Reads output until the next prompt, returning both. |
| 198 | + fn read_until_prompt(&mut self) -> Result<(String, String), Box<dyn Error>> { |
| 199 | + static UNTIL: std::sync::OnceLock<rexpect::ReadUntil> = std::sync::OnceLock::new(); |
| 200 | + let until = UNTIL.get_or_init(|| { |
| 201 | + let re = regex::Regex::new(r"toydb(:\d+|@\d+)?>\s+").expect("invalid regex"); |
| 202 | + rexpect::ReadUntil::Regex(re) |
| 203 | + }); |
| 204 | + let (mut output, mut prompt) = self.session.reader.read_until(until)?; |
| 205 | + output = output.trim().replace("\r\n", "\n"); |
| 206 | + prompt = prompt.trim().to_string(); |
| 207 | + Ok((output, prompt)) |
| 208 | + } |
| 209 | +} |
0 commit comments