use crate::metrics::CloudInfo;
use std::time::Duration;
use ureq::config::Config as UreqConfig;
mod alicloud;
mod aws;
mod azure;
mod gcp;
mod hetzner;
mod ovh;
mod upcloud;
mod vultr;
const IMDS_TIMEOUT: Duration = Duration::from_secs(1);
fn new_imds_agent() -> ureq::Agent {
UreqConfig::builder()
.timeout_connect(Some(IMDS_TIMEOUT))
.timeout_recv_response(Some(IMDS_TIMEOUT))
.build()
.new_agent()
}
fn imds_get(agent: &ureq::Agent, url: &str) -> Option<String> {
imds_get_headers(agent, url, &[])
}
fn imds_get_headers(agent: &ureq::Agent, url: &str, headers: &[(&str, &str)]) -> Option<String> {
let mut req = agent.get(url);
for (k, v) in headers {
req = req.header(*k, *v);
}
req.call()
.ok()
.and_then(|mut r| r.body_mut().read_to_string().ok())
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
}
const PROBES: &[fn() -> Option<CloudInfo>] = &[
aws::probe,
gcp::probe,
azure::probe,
hetzner::probe,
upcloud::probe,
vultr::probe,
alicloud::probe,
ovh::probe,
];
pub fn probe_cloud() -> CloudInfo {
let mut handles = Vec::new();
let mut deferred = Vec::new();
for &p in PROBES {
match crate::thread_util::spawn_named("cloud-probe", p) {
Some(h) => handles.push(h),
None => deferred.push(p),
}
}
for handle in handles {
if let Ok(Some(info)) = handle.join() {
return info;
}
}
for p in deferred {
if let Some(info) = p() {
return info;
}
}
CloudInfo::default()
}
pub fn spawn_cloud_discovery() -> Option<std::sync::mpsc::Receiver<CloudInfo>> {
let (tx, rx) = std::sync::mpsc::channel();
crate::thread_util::spawn_named("cloud-discovery", move || {
let _ = tx.send(probe_cloud());
})?;
Some(rx)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_spawn_cloud_discovery_joins_without_panic() {
let cloud = match spawn_cloud_discovery() {
Some(rx) => rx.recv().unwrap_or_default(),
None => probe_cloud(),
};
let _cloud = cloud;
}
#[cfg(target_os = "linux")]
fn thread_count() -> usize {
std::fs::read_to_string("/proc/self/status")
.unwrap_or_default()
.lines()
.find(|l| l.starts_with("Threads:"))
.and_then(|l| l.split_whitespace().nth(1))
.and_then(|n| n.parse().ok())
.unwrap_or(0)
}
#[cfg(target_os = "linux")]
fn slow_mock_server(delay: std::time::Duration) -> u16 {
use std::io::Write;
use std::net::TcpListener;
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let port = listener.local_addr().unwrap().port();
std::thread::spawn(move || {
if let Ok((mut s, _)) = listener.accept() {
std::thread::sleep(delay);
let _ = s.write_all(b"HTTP/1.0 200 OK\r\nContent-Length: 0\r\n\r\n");
}
});
port
}
#[cfg(target_os = "linux")]
#[test]
fn test_per_phase_timeout_no_helper_thread() {
use std::sync::mpsc;
let port = slow_mock_server(Duration::from_millis(200));
let url = format!("http://127.0.0.1:{port}");
let baseline = thread_count();
let (done_tx, done_rx) = mpsc::channel::<()>();
std::thread::spawn(move || {
let agent = new_imds_agent();
let _ = agent.get(&url).call();
let _ = done_tx.send(());
});
std::thread::sleep(Duration::from_millis(50));
let during = thread_count();
done_rx.recv().unwrap();
let pid = std::process::id();
let thread_names: Vec<String> = std::fs::read_dir(format!("/proc/{pid}/task"))
.into_iter()
.flatten()
.filter_map(|e| e.ok())
.filter_map(|e| std::fs::read_to_string(e.path().join("comm")).ok())
.map(|s| s.trim().to_string())
.collect();
assert!(
during <= baseline + 1,
"ureq spawned extra thread(s) under per-phase timeout: \
baseline={baseline} during={during} threads={thread_names:?}"
);
}
}