sqlness_runner/
util.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::io::Read;
16use std::net::SocketAddr;
17use std::path::{Path, PathBuf};
18use std::process::Command;
19use std::time::Duration;
20
21use sha2::{Digest, Sha256};
22use tokio::io::AsyncWriteExt;
23use tokio::net::TcpSocket;
24use tokio::time;
25use tokio_stream::StreamExt;
26
27/// Check port every 0.1 second.
28const PORT_CHECK_INTERVAL: Duration = Duration::from_millis(100);
29
30pub const PROGRAM: &str = "greptime";
31
32fn http_proxy() -> Option<String> {
33    for proxy in ["http_proxy", "HTTP_PROXY", "all_proxy", "ALL_PROXY"] {
34        if let Ok(proxy_addr) = std::env::var(proxy) {
35            println!("Getting Proxy from env var: {}={}", proxy, proxy_addr);
36            return Some(proxy_addr);
37        }
38    }
39    None
40}
41
42fn https_proxy() -> Option<String> {
43    for proxy in ["https_proxy", "HTTPS_PROXY", "all_proxy", "ALL_PROXY"] {
44        if let Ok(proxy_addr) = std::env::var(proxy) {
45            println!("Getting Proxy from env var: {}={}", proxy, proxy_addr);
46            return Some(proxy_addr);
47        }
48    }
49    None
50}
51
52async fn download_files(url: &str, path: &str) {
53    let proxy = if url.starts_with("http://") {
54        http_proxy().map(|proxy| reqwest::Proxy::http(proxy).unwrap())
55    } else if url.starts_with("https://") {
56        https_proxy().map(|proxy| reqwest::Proxy::https(proxy).unwrap())
57    } else {
58        None
59    };
60
61    let client = proxy
62        .map(|proxy| {
63            reqwest::Client::builder()
64                .proxy(proxy)
65                .build()
66                .expect("Failed to build client")
67        })
68        .unwrap_or(reqwest::Client::new());
69
70    let mut file = tokio::fs::File::create(path)
71        .await
72        .unwrap_or_else(|_| panic!("Failed to create file in {path}"));
73    println!("Downloading {}...", url);
74
75    let resp = client
76        .get(url)
77        .send()
78        .await
79        .expect("Failed to send download request");
80    let len = resp.content_length();
81    let mut stream = resp.bytes_stream();
82    let mut size_downloaded = 0;
83
84    while let Some(chunk_result) = stream.next().await {
85        let chunk = chunk_result.unwrap();
86        size_downloaded += chunk.len();
87        if let Some(len) = len {
88            print!("\rDownloading {}/{} bytes", size_downloaded, len);
89        } else {
90            print!("\rDownloaded {} bytes", size_downloaded);
91        }
92
93        file.write_all(&chunk).await.unwrap();
94    }
95
96    file.flush().await.unwrap();
97
98    println!("\nDownloaded {}", url);
99}
100
101fn decompress(archive: &str, dest: &str) {
102    let tar = std::fs::File::open(archive).unwrap();
103    let dec = flate2::read::GzDecoder::new(tar);
104    let mut a = tar::Archive::new(dec);
105    a.unpack(dest).unwrap();
106}
107
108/// Use curl to download the binary from the release page.
109///
110/// # Arguments
111///
112/// * `version` - The version of the binary to download. i.e. "v0.9.5"
113pub async fn pull_binary(version: &str) {
114    let os = std::env::consts::OS;
115    let arch = match std::env::consts::ARCH {
116        "x86_64" => "amd64",
117        "aarch64" => "arm64",
118        _ => panic!("Unsupported arch: {}", std::env::consts::ARCH),
119    };
120    let triple = format!("greptime-{}-{}-{}", os, arch, version);
121    let filename = format!("{triple}.tar.gz");
122
123    let url = format!(
124        "https://github.com/GreptimeTeam/greptimedb/releases/download/{version}/{filename}"
125    );
126    println!("Downloading {version} binary from {}", url);
127
128    // mkdir {version}
129    let _ = std::fs::create_dir(version);
130
131    let archive = Path::new(version).join(filename);
132    let folder_path = Path::new(version);
133
134    // download the binary to the version directory
135    download_files(&url, &archive.to_string_lossy()).await;
136
137    let checksum_file = format!("{triple}.sha256sum");
138    let checksum_url = format!(
139        "https://github.com/GreptimeTeam/greptimedb/releases/download/{version}/{checksum_file}"
140    );
141    download_files(
142        &checksum_url,
143        &PathBuf::from_iter([version, &checksum_file]).to_string_lossy(),
144    )
145    .await;
146
147    // verify the checksum
148    let mut file = std::fs::File::open(&archive).unwrap();
149    let mut sha256 = Sha256::new();
150    std::io::copy(&mut file, &mut sha256).unwrap();
151    let checksum: Vec<u8> = sha256.finalize().to_vec();
152
153    let mut expected_checksum =
154        std::fs::File::open(PathBuf::from_iter([version, &checksum_file])).unwrap();
155    let mut buf = String::new();
156    expected_checksum.read_to_string(&mut buf).unwrap();
157    let expected_checksum = hex::decode(buf.lines().next().unwrap()).unwrap();
158
159    assert_eq!(
160        checksum, expected_checksum,
161        "Checksum mismatched, downloaded file is corrupted"
162    );
163
164    decompress(&archive.to_string_lossy(), &folder_path.to_string_lossy());
165    println!("Downloaded and extracted {version} binary to {folder_path:?}");
166
167    // move the binary to the version directory
168    std::fs::rename(
169        PathBuf::from_iter([version, &triple, "greptime"]),
170        PathBuf::from_iter([version, "greptime"]),
171    )
172    .unwrap();
173
174    // remove the archive and inner folder
175    std::fs::remove_file(&archive).unwrap();
176    std::fs::remove_dir(PathBuf::from_iter([version, &triple])).unwrap();
177}
178
179/// Pull the binary if it does not exist and `pull_version_on_need` is true.
180pub async fn maybe_pull_binary(version: &str, pull_version_on_need: bool) {
181    let exist = Path::new(version).join(PROGRAM).is_file();
182    match (exist, pull_version_on_need) {
183        (true, _) => println!("Binary {version} exists"),
184        (false, false) => panic!(
185            "Binary {version} does not exist, please run with --pull-version-on-need or manually download it"
186        ),
187        (false, true) => {
188            pull_binary(version).await;
189        }
190    }
191}
192
193/// Set up a standalone etcd in docker.
194pub fn setup_etcd(client_ports: Vec<u16>, peer_port: Option<u16>, etcd_version: Option<&str>) {
195    if std::process::Command::new("docker")
196        .args(["-v"])
197        .status()
198        .is_err()
199    {
200        panic!("Docker is not installed");
201    }
202    let peer_port = peer_port.unwrap_or(2380);
203    let exposed_port: Vec<_> = client_ports.iter().chain(Some(&peer_port)).collect();
204    let exposed_port_str = exposed_port
205        .iter()
206        .flat_map(|p| ["-p".to_string(), format!("{p}:{p}")])
207        .collect::<Vec<_>>();
208    let etcd_version = etcd_version.unwrap_or("v3.5.17");
209    let etcd_image = format!("quay.io/coreos/etcd:{etcd_version}");
210    let peer_url = format!("http://0.0.0.0:{peer_port}");
211    let my_local_ip = local_ip_address::local_ip().unwrap();
212
213    let my_local_ip_str = my_local_ip.to_string();
214
215    let mut arg_list = vec![];
216    arg_list.extend([
217        "run",
218        "-d",
219        "-v",
220        "/usr/share/ca-certificates/:/etc/ssl/certs",
221    ]);
222    arg_list.extend(exposed_port_str.iter().map(std::ops::Deref::deref));
223    arg_list.extend([
224        "--name",
225        "etcd",
226        &etcd_image,
227        "etcd",
228        "-name",
229        "etcd0",
230        "-advertise-client-urls",
231    ]);
232
233    let adv_client_urls = client_ports
234        .iter()
235        .map(|p| format!("http://{my_local_ip_str}:{p}"))
236        .collect::<Vec<_>>()
237        .join(",");
238
239    arg_list.push(&adv_client_urls);
240
241    arg_list.extend(["-listen-client-urls"]);
242
243    let client_ports_fmt = client_ports
244        .iter()
245        .map(|p| format!("http://0.0.0.0:{p}"))
246        .collect::<Vec<_>>()
247        .join(",");
248
249    arg_list.push(&client_ports_fmt);
250
251    arg_list.push("-initial-advertise-peer-urls");
252    let advertise_peer_url = format!("http://{my_local_ip_str}:{peer_port}");
253    arg_list.push(&advertise_peer_url);
254
255    arg_list.extend(["-listen-peer-urls", &peer_url]);
256
257    arg_list.extend(["-initial-cluster-token", "etcd-cluster-1"]);
258
259    arg_list.push("-initial-cluster");
260
261    let init_cluster_url = format!("etcd0=http://{my_local_ip_str}:{peer_port}");
262
263    arg_list.push(&init_cluster_url);
264
265    arg_list.extend(["-initial-cluster-state", "new"]);
266
267    let mut cmd = std::process::Command::new("docker");
268
269    cmd.args(arg_list);
270
271    println!("Starting etcd with command: {:?}", cmd);
272
273    let status = cmd.status();
274    if status.is_err() {
275        panic!("Failed to start etcd: {:?}", status);
276    } else if let Ok(status) = status {
277        if status.success() {
278            println!(
279                "Started etcd with client ports {:?} and peer port {}, statues:{status:?}",
280                client_ports, peer_port
281            );
282        } else {
283            panic!("Failed to start etcd: {:?}", status);
284        }
285    }
286}
287
288/// Stop and remove the etcd container, failing if it cannot be confirmed absent.
289pub fn stop_rm_etcd_checked() -> Result<(), String> {
290    let status = std::process::Command::new("docker")
291        .args(["container", "rm", "--force", "etcd"])
292        .status()
293        .map_err(|error| format!("Failed to run Docker while removing etcd: {error}"))?;
294    if status.success() {
295        println!("Removed etcd container");
296        return Ok(());
297    }
298
299    let listed = std::process::Command::new("docker")
300        .args([
301            "container",
302            "ls",
303            "--all",
304            "--filter",
305            "name=^/etcd$",
306            "--format",
307            "{{.ID}}",
308        ])
309        .output()
310        .map_err(|error| {
311            format!(
312                "Docker failed to remove etcd ({status}) and could not verify its absence: {error}"
313            )
314        })?;
315    if !listed.status.success() {
316        return Err(format!(
317            "Docker failed to remove etcd ({status}) and listing containers failed: {}",
318            String::from_utf8_lossy(&listed.stderr).trim()
319        ));
320    }
321    if String::from_utf8_lossy(&listed.stdout).trim().is_empty() {
322        println!("Etcd container is already absent");
323        return Ok(());
324    }
325
326    Err(format!(
327        "Docker failed to remove etcd container (status {status}); the container still exists"
328    ))
329}
330
331/// Stop and remove the etcd container.
332///
333/// Legacy callers retain panic-on-cleanup-failure behavior. Compatibility
334/// profiles use [`stop_rm_etcd_checked`] to report the failure structurally.
335pub fn stop_rm_etcd() {
336    stop_rm_etcd_checked().unwrap_or_else(|error| panic!("{error}"));
337}
338
339/// Set up a PostgreSQL server in docker.
340pub fn setup_pg(pg_port: u16, pg_version: Option<&str>) {
341    if std::process::Command::new("docker")
342        .args(["-v"])
343        .status()
344        .is_err()
345    {
346        panic!("Docker is not installed");
347    }
348
349    let pg_image = if let Some(pg_version) = pg_version {
350        format!("postgres:{pg_version}")
351    } else {
352        "postgres:latest".to_string()
353    };
354    let pg_password = "admin";
355    let pg_user = "greptimedb";
356
357    let mut arg_list = vec![];
358    arg_list.extend(["run", "-d"]);
359
360    let pg_password_env = format!("POSTGRES_PASSWORD={pg_password}");
361    let pg_user_env = format!("POSTGRES_USER={pg_user}");
362    let pg_port_forward = format!("{pg_port}:5432");
363    arg_list.extend(["-e", &pg_password_env, "-e", &pg_user_env]);
364    arg_list.extend(["-p", &pg_port_forward]);
365
366    arg_list.extend(["--name", "greptimedb_pg", &pg_image]);
367
368    let mut cmd = std::process::Command::new("docker");
369
370    cmd.args(arg_list);
371
372    println!("Starting PostgreSQL with command: {:?}", cmd);
373
374    let status = cmd.status();
375    if status.is_err() {
376        panic!("Failed to start PostgreSQL: {:?}", status);
377    } else if let Ok(status) = status {
378        if status.success() {
379            println!("Started PostgreSQL with port {}", pg_port);
380        } else {
381            panic!("Failed to start PostgreSQL: {:?}", status);
382        }
383    }
384}
385
386/// Set up a MySql server in docker.
387pub fn setup_mysql(mysql_port: u16, mysql_version: Option<&str>) {
388    if std::process::Command::new("docker")
389        .args(["-v"])
390        .status()
391        .is_err()
392    {
393        panic!("Docker is not installed");
394    }
395
396    let mysql_image = if let Some(mysql_version) = mysql_version {
397        format!("greptime/mysql:{mysql_version}")
398    } else {
399        "greptime/mysql:5.7".to_string()
400    };
401    let mysql_password = "admin";
402    let mysql_user = "greptimedb";
403
404    let mut arg_list = vec![];
405    arg_list.extend(["run", "-d"]);
406
407    let mysql_password_env = format!("MYSQL_PASSWORD={mysql_password}");
408    let mysql_user_env = format!("MYSQL_USER={mysql_user}");
409    let mysql_root_password_env = format!("MYSQL_ROOT_PASSWORD={mysql_password}");
410    let mysql_port_forward = format!("{mysql_port}:3306");
411    arg_list.extend([
412        "-e",
413        &mysql_password_env,
414        "-e",
415        &mysql_user_env,
416        "-e",
417        &mysql_root_password_env,
418        "-e",
419        "MYSQL_DATABASE=mysql",
420    ]);
421    arg_list.extend(["-p", &mysql_port_forward]);
422
423    arg_list.extend(["--name", "greptimedb_mysql", &mysql_image]);
424
425    let mut cmd = std::process::Command::new("docker");
426
427    cmd.args(arg_list);
428
429    println!("Starting MySQL with command: {:?}", cmd);
430
431    let status = cmd.status();
432    if status.is_err() {
433        panic!("Failed to start MySQL: {:?}", status);
434    } else if let Ok(status) = status {
435        if status.success() {
436            println!("Started MySQL with port {}", mysql_port);
437        } else {
438            panic!("Failed to start MySQL: {:?}", status);
439        }
440    }
441}
442
443/// Get the dir of test cases. This function only works when the runner is run
444/// under the project's dir because it depends on some envs set by cargo.
445pub fn get_case_dir(case_dir: Option<PathBuf>) -> String {
446    let runner_path = match case_dir {
447        Some(path) => path,
448        None => {
449            // retrieve the manifest runner (./tests/runner)
450            let mut runner_crate_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
451            // change directory to cases' dir from runner's (should be runner/../cases)
452            let _ = runner_crate_path.pop();
453            runner_crate_path.push("cases");
454            runner_crate_path
455        }
456    };
457
458    runner_path.into_os_string().into_string().unwrap()
459}
460
461/// Get the dir that contains workspace manifest (the top-level Cargo.toml).
462pub fn get_workspace_root() -> String {
463    // retrieve the manifest runner (./tests/runner)
464    let mut runner_crate_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
465
466    // change directory to workspace's root (runner/../..)
467    let _ = runner_crate_path.pop();
468    let _ = runner_crate_path.pop();
469
470    runner_crate_path.into_os_string().into_string().unwrap()
471}
472
473pub fn get_binary_dir(mode: &str) -> PathBuf {
474    // first go to the workspace root.
475    let mut workspace_root = PathBuf::from(get_workspace_root());
476
477    // change directory to target dir (workspace/target/<build mode>/)
478    workspace_root.push("target");
479    workspace_root.push(mode);
480
481    workspace_root
482}
483
484/// Spin-waiting a socket address is available, or timeout.
485/// Returns whether the addr is up.
486pub async fn check_port(ip_addr: SocketAddr, timeout: Duration) -> bool {
487    let check_task = async {
488        loop {
489            let socket = TcpSocket::new_v4().expect("Cannot create v4 socket");
490            match socket.connect(ip_addr).await {
491                Ok(mut stream) => {
492                    let _ = stream.shutdown().await;
493                    break;
494                }
495                Err(_) => time::sleep(PORT_CHECK_INTERVAL).await,
496            }
497        }
498    };
499
500    tokio::time::timeout(timeout, check_task).await.is_ok()
501}
502
503/// Get the path of sqlness config dir `tests/conf`.
504pub fn sqlness_conf_path() -> PathBuf {
505    let mut sqlness_root_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
506    sqlness_root_path.pop();
507    sqlness_root_path.push("conf");
508    sqlness_root_path
509}
510
511/// Start kafka cluster if needed. Config file is `conf/kafka-cluster.yml`.
512///
513/// ```shell
514/// docker compose -f kafka-cluster.yml up kafka -d --wait
515/// ```
516pub fn setup_wal() {
517    let mut sqlness_root_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
518    sqlness_root_path.pop();
519    sqlness_root_path.push("conf");
520
521    Command::new("docker")
522        .current_dir(sqlness_conf_path())
523        .args([
524            "compose",
525            "-f",
526            "kafka-cluster.yml",
527            "up",
528            "kafka",
529            "-d",
530            "--wait",
531        ])
532        .output()
533        .expect("Failed to start kafka cluster");
534
535    println!("kafka cluster is up");
536}
537
538/// Stop kafka cluster if needed. Config file is `conf/kafka-cluster.yml`.
539///
540/// ```shell
541/// docker compose -f docker-compose-standalone.yml down kafka
542/// ```
543pub fn teardown_wal() {
544    let mut sqlness_root_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
545    sqlness_root_path.pop();
546    sqlness_root_path.push("conf");
547
548    Command::new("docker")
549        .current_dir(sqlness_conf_path())
550        .args(["compose", "-f", "kafka-cluster.yml", "down", "kafka"])
551        .output()
552        .expect("Failed to stop kafka cluster");
553
554    println!("kafka cluster is down");
555}
556
557/// Get a random available port by binding to port 0
558pub fn get_random_port() -> u16 {
559    use std::net::TcpListener;
560    let listener = TcpListener::bind("127.0.0.1:0").expect("Failed to bind to random port");
561    listener
562        .local_addr()
563        .expect("Failed to get local address")
564        .port()
565}
566
567/// Retry to execute the function until success or the maximum number of retries is reached.
568pub async fn retry_with_backoff<T, E, Fut, F>(
569    mut fut: F,
570    max_retry: usize,
571    init_backoff: Duration,
572) -> Result<T, E>
573where
574    F: FnMut() -> Fut,
575    Fut: Future<Output = Result<T, E>>,
576{
577    let mut backoff = init_backoff;
578    for attempt in 0..max_retry {
579        match fut().await {
580            Ok(res) => return Ok(res),
581            Err(err) if attempt + 1 == max_retry => return Err(err),
582            Err(_) => {
583                tokio::time::sleep(backoff).await;
584                backoff *= 2;
585            }
586        }
587    }
588
589    unreachable!("loop should have returned before here")
590}
sqlness_runner/util.rs

sqlness_runner/
util.rs
