gonzalo_soak/
replica.rs

1//! A set of real `gonzalod` subprocesses, all backed by the same S3
2//! store — the "N stateless replicas behind a Service" of the k8s HA model.
3//!
4//! Each replica is a `gonzalod` process configured entirely from the environment
5//! (`GONZALO_STORE=s3` + the shared bucket/endpoint/creds + distinct bind ports).
6//! [`ReplicaSet::kill`] SIGKILLs one (a pod death); [`ReplicaSet::respawn`]
7//! restarts it. The `gonzalod` binary is found via `GONZALO_SOAK_GONZALOD_BIN`
8//! or, failing that, the workspace target dir next to the test/soak executable.
9
10use crate::target::S3Target;
11use std::path::PathBuf;
12use std::process::{Child, Command, Stdio};
13use std::time::{Duration, Instant};
14
15/// How a single replica is (re)spawned: its bind ports and full env.
16struct ReplicaSpec {
17    http_port: u16,
18    grpc_port: u16,
19    env: Vec<(String, String)>,
20}
21
22struct Replica {
23    spec: ReplicaSpec,
24    child: Option<Child>,
25}
26
27/// A live set of `gonzalod` replicas over a shared S3 backend.
28pub struct ReplicaSet {
29    gonzalod: PathBuf,
30    replicas: Vec<Replica>,
31}
32
33impl ReplicaSet {
34    /// Spawn `n` replicas over `target` and wait until each answers `/readyz`.
35    pub async fn spawn(n: usize, target: &S3Target) -> Result<Self, String> {
36        assert!(n >= 1, "need at least one replica");
37        let gonzalod = locate_gonzalod()?;
38        let mut set = ReplicaSet {
39            gonzalod,
40            replicas: Vec::with_capacity(n),
41        };
42        for i in 0..n {
43            // Distinct localhost ports per replica; the shared bucket makes them
44            // interchangeable stateless fronts over one durable store.
45            let http_port = 18080 + i as u16;
46            let grpc_port = 18150 + i as u16;
47            let env = replica_env(target, http_port, grpc_port);
48            let spec = ReplicaSpec {
49                http_port,
50                grpc_port,
51                env,
52            };
53            let child = set.launch(&spec)?;
54            set.replicas.push(Replica {
55                spec,
56                child: Some(child),
57            });
58        }
59        for i in 0..n {
60            set.wait_ready(i).await?;
61        }
62        Ok(set)
63    }
64
65    /// The HTTP base URLs of every replica (whether currently live or not).
66    pub fn base_urls(&self) -> Vec<String> {
67        self.replicas
68            .iter()
69            .map(|r| format!("http://127.0.0.1:{}", r.spec.http_port))
70            .collect()
71    }
72
73    /// Number of replicas.
74    pub fn len(&self) -> usize {
75        self.replicas.len()
76    }
77
78    /// True when there are no replicas.
79    pub fn is_empty(&self) -> bool {
80        self.replicas.is_empty()
81    }
82
83    /// SIGKILL replica `idx` (simulates a pod death). Idempotent.
84    pub fn kill(&mut self, idx: usize) {
85        if let Some(mut child) = self.replicas[idx].child.take() {
86            let _ = child.kill();
87            let _ = child.wait();
88        }
89    }
90
91    /// Restart replica `idx` and wait for it to be ready again.
92    pub async fn respawn(&mut self, idx: usize) -> Result<(), String> {
93        if self.replicas[idx].child.is_none() {
94            let spec = ReplicaSpec {
95                http_port: self.replicas[idx].spec.http_port,
96                grpc_port: self.replicas[idx].spec.grpc_port,
97                env: self.replicas[idx].spec.env.clone(),
98            };
99            let child = self.launch(&spec)?;
100            self.replicas[idx].child = Some(child);
101        }
102        self.wait_ready(idx).await
103    }
104
105    fn launch(&self, spec: &ReplicaSpec) -> Result<Child, String> {
106        Command::new(&self.gonzalod)
107            .env_clear()
108            .envs(spec.env.iter().map(|(k, v)| (k.as_str(), v.as_str())))
109            // Keep a minimal PATH so the AWS SDK / TLS can find system bits.
110            .env("PATH", std::env::var("PATH").unwrap_or_default())
111            .stdout(Stdio::null())
112            .stderr(Stdio::inherit())
113            .spawn()
114            .map_err(|e| format!("spawn gonzalod ({}): {e}", self.gonzalod.display()))
115    }
116
117    /// Poll `/readyz` on replica `idx` until 200 OK or a deadline.
118    async fn wait_ready(&self, idx: usize) -> Result<(), String> {
119        let url = format!(
120            "http://127.0.0.1:{}/readyz",
121            self.replicas[idx].spec.http_port
122        );
123        let client = reqwest::Client::new();
124        let deadline = Instant::now() + Duration::from_secs(20);
125        loop {
126            if let Ok(resp) = client
127                .get(&url)
128                .timeout(Duration::from_secs(2))
129                .send()
130                .await
131                && resp.status().is_success()
132            {
133                return Ok(());
134            }
135            if Instant::now() >= deadline {
136                return Err(format!("replica {idx} never became ready at {url}"));
137            }
138            tokio::time::sleep(Duration::from_millis(150)).await;
139        }
140    }
141}
142
143impl Drop for ReplicaSet {
144    fn drop(&mut self) {
145        for r in &mut self.replicas {
146            if let Some(mut child) = r.child.take() {
147                let _ = child.kill();
148                let _ = child.wait();
149            }
150        }
151    }
152}
153
154/// The env a single S3-backed `gonzalod` replica runs with.
155fn replica_env(target: &S3Target, http_port: u16, grpc_port: u16) -> Vec<(String, String)> {
156    let mut env = vec![
157        ("GONZALO_STORE".into(), "s3".into()),
158        ("GONZALO_S3_BUCKET".into(), target.bucket.clone()),
159        ("GONZALO_S3_ENDPOINT".into(), target.endpoint.clone()),
160        ("AWS_ACCESS_KEY_ID".into(), target.access_key.clone()),
161        ("AWS_SECRET_ACCESS_KEY".into(), target.secret_key.clone()),
162        ("GONZALO_HTTP_ADDR".into(), format!("127.0.0.1:{http_port}")),
163        ("GONZALO_GRPC_ADDR".into(), format!("127.0.0.1:{grpc_port}")),
164    ];
165    env.push((
166        "GONZALO_S3_REGION".into(),
167        target.region.clone().unwrap_or_else(|| "us-east-1".into()),
168    ));
169    env
170}
171
172/// Find the `gonzalod` binary: an explicit override, or the workspace target dir
173/// adjacent to the currently-running test/soak executable.
174fn locate_gonzalod() -> Result<PathBuf, String> {
175    if let Some(p) = std::env::var_os("GONZALO_SOAK_GONZALOD_BIN") {
176        let p = PathBuf::from(p);
177        return if p.exists() {
178            Ok(p)
179        } else {
180            Err(format!(
181                "GONZALO_SOAK_GONZALOD_BIN set but not found: {}",
182                p.display()
183            ))
184        };
185    }
186    // current_exe is `.../target/<profile>/deps/<name>` or `.../target/<profile>/<name>`.
187    let exe = std::env::current_exe().map_err(|e| format!("current_exe: {e}"))?;
188    let mut dir = exe.parent().ok_or("exe has no parent")?.to_path_buf();
189    if dir.file_name().is_some_and(|n| n == "deps") {
190        dir = dir.parent().ok_or("deps has no parent")?.to_path_buf();
191    }
192    let bin = dir.join("gonzalod");
193    if bin.exists() {
194        Ok(bin)
195    } else {
196        Err(format!(
197            "gonzalod binary not found at {} — build it first: `cargo build --bin gonzalod` \
198             (or set GONZALO_SOAK_GONZALOD_BIN)",
199            bin.display()
200        ))
201    }
202}
gonzalo_soak/replica.rs

gonzalo_soak/
replica.rs
