gonzalo_soak/
workload.rs

1//! The concurrent write workload and result collection for the [`oracle`].
2//!
3//! Each writer runs two op streams against the replica [`Dispatcher`]:
4//!
5//! - **Contended RMW** on a small set of shared keys — read the current record,
6//!   append this op's globally-unique id to the record's comma-separated set,
7//!   and conditionally `put` against the revision just read. On `Conflict`,
8//!   re-read and retry (bounded). This is the arbitration proof: if conditional
9//!   writes are correct, every *committed* op-id survives in the final set.
10//! - **Unique-key writes** — disjoint keys written once, later read back to prove
11//!   no acked write is lost across replica kills.
12//!
13//! After all writers finish, [`run`] reads every key's final state and returns a
14//! [`SoakStats`] for [`oracle::check`].
15//!
16//! [`oracle`]: crate::oracle
17//! [`oracle::check`]: crate::oracle::check
18
19use crate::dispatch::Dispatcher;
20use crate::oracle::{FinalContended, FinalUnique, OpRecord, OpResult, SoakStats};
21use gonzalo_core::{Body, Identity, Meta, PutResult, Record, RecordKey, RecordKind, Revision};
22use std::collections::BTreeMap;
23use std::sync::Arc;
24use std::sync::atomic::{AtomicU64, Ordering};
25
26/// Workload shape. The bounded gate and the deep soak both use this; they differ
27/// only in the magnitudes and (for the deep soak) running against a duration.
28#[derive(Debug, Clone)]
29pub struct WorkloadConfig {
30    pub namespace: String,
31    pub collection: String,
32    pub writers: usize,
33    pub shared_keys: usize,
34    pub ops_per_writer: usize,
35    pub unique_per_writer: usize,
36    pub max_conflict_retries: usize,
37}
38
39impl Default for WorkloadConfig {
40    fn default() -> Self {
41        Self {
42            namespace: "soak".into(),
43            collection: "ha".into(),
44            writers: 8,
45            shared_keys: 4,
46            ops_per_writer: 25,
47            unique_per_writer: 4,
48            max_conflict_retries: 50,
49        }
50    }
51}
52
53/// Run the workload to completion and collect stats for the oracle. `dispatcher`
54/// fans ops across the live replicas; chaos (replica kills) is driven separately
55/// by the caller while this runs.
56pub async fn run(dispatcher: Arc<Dispatcher>, cfg: WorkloadConfig) -> SoakStats {
57    let op_ids = Arc::new(AtomicU64::new(1));
58    let mut handles = Vec::new();
59    for w in 0..cfg.writers {
60        let d = dispatcher.clone();
61        let c = cfg.clone();
62        let ids = op_ids.clone();
63        handles.push(tokio::spawn(async move { writer(w, d, c, ids).await }));
64    }
65
66    let mut ops = Vec::new();
67    let mut unique_acked: Vec<(String, Vec<u8>)> = Vec::new();
68    let mut conflicts_observed = 0u64;
69    let mut writers_completed = 0;
70    for h in handles {
71        if let Ok(res) = h.await {
72            ops.extend(res.ops);
73            unique_acked.extend(res.unique_acked);
74            conflicts_observed += res.conflicts;
75            writers_completed += 1;
76        }
77    }
78
79    let contended = collect_contended(&dispatcher, &cfg).await;
80    let unique = collect_unique(&dispatcher, &cfg, &unique_acked).await;
81
82    SoakStats {
83        ops,
84        contended,
85        unique,
86        conflicts_observed,
87        writers_completed,
88        writers_total: cfg.writers,
89    }
90}
91
92struct WriterResult {
93    ops: Vec<OpRecord>,
94    unique_acked: Vec<(String, Vec<u8>)>,
95    conflicts: u64,
96}
97
98async fn writer(
99    w: usize,
100    d: Arc<Dispatcher>,
101    cfg: WorkloadConfig,
102    op_ids: Arc<AtomicU64>,
103) -> WriterResult {
104    let mut ops = Vec::with_capacity(cfg.ops_per_writer);
105    let mut conflicts = 0u64;
106    for _ in 0..cfg.ops_per_writer {
107        let op_id = op_ids.fetch_add(1, Ordering::Relaxed);
108        let key_id = format!("shared-{}", (op_id as usize) % cfg.shared_keys.max(1));
109        let (result, seen) = rmw_append(&d, &cfg, &key_id, op_id).await;
110        conflicts += seen;
111        ops.push(OpRecord {
112            key: key_id,
113            op_id,
114            result,
115        });
116    }
117
118    let mut unique_acked = Vec::with_capacity(cfg.unique_per_writer);
119    for i in 0..cfg.unique_per_writer {
120        let key_id = format!("unique-{w}-{i}");
121        let value = key_id.clone().into_bytes();
122        if create(&d, &cfg, &key_id, &value).await {
123            unique_acked.push((key_id, value));
124        }
125    }
126
127    WriterResult {
128        ops,
129        unique_acked,
130        conflicts,
131    }
132}
133
134/// Read-modify-write: append `op_id` to a shared key's set under a conditional
135/// put, retrying on `Conflict`. Returns the op's final outcome and the number of
136/// transient conflicts observed while racing to commit it.
137async fn rmw_append(
138    d: &Dispatcher,
139    cfg: &WorkloadConfig,
140    key_id: &str,
141    op_id: u64,
142) -> (OpResult, u64) {
143    let key = RecordKey::new(&cfg.namespace, &cfg.collection, key_id);
144    let mut conflicts = 0u64;
145    for _ in 0..=cfg.max_conflict_retries {
146        let current = match d.get(&key).await {
147            Ok(c) => c,
148            Err(_) => return (OpResult::Failed, conflicts),
149        };
150        let (mut members, expected) = match &current {
151            Some(rec) => (parse_members(rec.body.bytes()), Some(rec.revision.clone())),
152            None => (Vec::new(), None),
153        };
154        if !members.contains(&op_id) {
155            members.push(op_id);
156        }
157        let record = build_record(&key, &encode_members(&members), expected.clone());
158        match d.put(record, expected).await {
159            Ok(PutResult::Committed(_)) => return (OpResult::Committed, conflicts),
160            Ok(PutResult::Conflict(_)) => {
161                conflicts += 1;
162                continue;
163            }
164            Err(_) => return (OpResult::Failed, conflicts),
165        }
166    }
167    // Exhausted the retry budget without winning the race — did not commit.
168    (OpResult::Conflict, conflicts)
169}
170
171/// Create a unique key once. Returns `true` if the write was acked (`Committed`).
172async fn create(d: &Dispatcher, cfg: &WorkloadConfig, key_id: &str, value: &[u8]) -> bool {
173    let key = RecordKey::new(&cfg.namespace, &cfg.collection, key_id);
174    let record = build_record(&key, value, None);
175    matches!(d.put(record, None).await, Ok(PutResult::Committed(_)))
176}
177
178async fn collect_contended(d: &Dispatcher, cfg: &WorkloadConfig) -> Vec<FinalContended> {
179    let mut out = Vec::new();
180    for i in 0..cfg.shared_keys {
181        let key_id = format!("shared-{i}");
182        let key = RecordKey::new(&cfg.namespace, &cfg.collection, &key_id);
183        let members = match d.get(&key).await {
184            Ok(Some(rec)) => parse_members(rec.body.bytes()),
185            _ => Vec::new(),
186        };
187        out.push(FinalContended {
188            key: key_id,
189            members,
190        });
191    }
192    out
193}
194
195async fn collect_unique(
196    d: &Dispatcher,
197    cfg: &WorkloadConfig,
198    acked: &[(String, Vec<u8>)],
199) -> Vec<FinalUnique> {
200    let mut out = Vec::new();
201    for (key_id, value) in acked {
202        let key = RecordKey::new(&cfg.namespace, &cfg.collection, key_id);
203        let readable_with_value = matches!(
204            d.get(&key).await,
205            Ok(Some(rec)) if rec.body.bytes() == value.as_slice()
206        );
207        out.push(FinalUnique {
208            key: key_id.clone(),
209            acked: true,
210            readable_with_value,
211        });
212    }
213    out
214}
215
216fn build_record(key: &RecordKey, body_bytes: &[u8], parent: Option<Revision>) -> Record {
217    let body = Body::Inline(body_bytes.to_vec());
218    let revision = match &parent {
219        Some(p) => p.next(body.bytes()),
220        None => Revision::initial(body.bytes()),
221    };
222    Record {
223        revision,
224        parent,
225        body,
226        kind: RecordKind::Topic,
227        meta: Meta {
228            author: Identity::new("soak"),
229            origin_system: "soak".into(),
230            created: 0,
231            updated: 0,
232            labels: BTreeMap::new(),
233        },
234        links: Vec::new(),
235        key: key.clone(),
236    }
237}
238
239fn parse_members(bytes: &[u8]) -> Vec<u64> {
240    std::str::from_utf8(bytes)
241        .unwrap_or("")
242        .split(',')
243        .filter_map(|s| s.trim().parse::<u64>().ok())
244        .collect()
245}
246
247fn encode_members(members: &[u64]) -> Vec<u8> {
248    members
249        .iter()
250        .map(u64::to_string)
251        .collect::<Vec<_>>()
252        .join(",")
253        .into_bytes()
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259    use gonzalo_core::Store;
260    use gonzalo_store_fs::FsStore;
261
262    #[test]
263    fn parse_round_trips() {
264        assert_eq!(parse_members(&encode_members(&[3, 1, 2])), vec![3, 1, 2]);
265        assert_eq!(parse_members(b""), Vec::<u64>::new());
266    }
267
268    /// End-to-end workload against three in-process `FsStore` replicas over one
269    /// shared directory — real concurrency + real conditional-write arbitration,
270    /// no external S3 backend. Proves the RMW/oracle/dispatch integration holds the invariant.
271    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
272    async fn in_process_fs_replicas_hold_the_invariant() {
273        let dir = tempfile::tempdir().unwrap();
274        let replicas: Vec<Arc<dyn Store>> = (0..3)
275            .map(|_| Arc::new(FsStore::new(dir.path())) as Arc<dyn Store>)
276            .collect();
277        let dispatcher = Arc::new(Dispatcher::new(replicas));
278
279        let cfg = WorkloadConfig {
280            writers: 6,
281            shared_keys: 3,
282            ops_per_writer: 30,
283            unique_per_writer: 3,
284            ..Default::default()
285        };
286        let stats = run(dispatcher, cfg).await;
287
288        let violations = crate::oracle::check(&stats);
289        assert!(violations.is_empty(), "invariant violated: {violations:?}");
290    }
291}
gonzalo_soak/workload.rs

gonzalo_soak/
workload.rs
