gonzalo_soak/
dispatch.rs

1//! The replica dispatcher — the harness's stand-in for the k8s Service.
2//!
3//! Holds one [`Store`] handle per `gonzalod` replica, round-robins each op
4//! (`get`, `get_raw`, `put`, `delete`) to a replica, and fails over on a
5//! `Backend` error (transport, dead replica). `Conflict` (including a
6//! [`PutResult::Conflict`], which the caller answers by re-reading and
7//! retrying the RMW), `NotFound` and `Serde` are answers from a live replica,
8//! returned unchanged and never failover triggers. This is exactly the
9//! failover path a k8s agent pod relies on when a `gonzalod` pod dies behind
10//! the Service.
11
12use gonzalo_core::{
13    CoreError, DeleteResult, PutResult, Record, RecordKey, Result, Revision, Store,
14};
15use std::sync::Arc;
16use std::sync::atomic::{AtomicUsize, Ordering};
17
18/// Fans record ops across N replica stores with round-robin + failover.
19pub struct Dispatcher {
20    replicas: Vec<Arc<dyn Store>>,
21    next: AtomicUsize,
22}
23
24impl Dispatcher {
25    /// Build a dispatcher over one or more replica stores.
26    pub fn new(replicas: Vec<Arc<dyn Store>>) -> Self {
27        assert!(
28            !replicas.is_empty(),
29            "dispatcher needs at least one replica"
30        );
31        Self {
32            replicas,
33            next: AtomicUsize::new(0),
34        }
35    }
36
37    /// Number of replica handles (live or not — deadness is discovered per op).
38    pub fn len(&self) -> usize {
39        self.replicas.len()
40    }
41
42    /// Always false; provided to satisfy clippy's `len`-without-`is_empty` lint.
43    pub fn is_empty(&self) -> bool {
44        self.replicas.is_empty()
45    }
46
47    /// Every replica handle, in construction order. Used to read **each**
48    /// replica directly (no round-robin, no failover) once the run has settled.
49    pub fn replicas(&self) -> &[Arc<dyn Store>] {
50        &self.replicas
51    }
52
53    fn start(&self) -> usize {
54        self.next.fetch_add(1, Ordering::Relaxed) % self.replicas.len()
55    }
56
57    /// Try replicas in round-robin order. Only `CoreError::Backend` (transport,
58    /// 5xx, a dead or killed replica) fails over. `NotFound` and `Serde` are
59    /// answers from a live replica over shared storage and return at once, like
60    /// `Conflict`.
61    async fn with_failover<T, F, Fut>(&self, op: F) -> Result<T>
62    where
63        F: Fn(Arc<dyn Store>) -> Fut,
64        Fut: std::future::Future<Output = Result<T>>,
65    {
66        let n = self.replicas.len();
67        let start = self.start();
68        let mut last_err = None;
69        for offset in 0..n {
70            match op(self.replicas[(start + offset) % n].clone()).await {
71                Err(e @ CoreError::Backend(_)) => last_err = Some(e),
72                other => return other,
73            }
74        }
75        Err(last_err.expect("at least one replica was attempted"))
76    }
77
78    /// `get` with round-robin start + failover across all replicas.
79    pub async fn get(&self, key: &RecordKey) -> Result<Option<Record>> {
80        self.with_failover(|s| {
81            let key = key.clone();
82            async move { s.get(&key).await }
83        })
84        .await
85    }
86
87    /// `put` with round-robin start + failover across all replicas. The record is
88    /// cloned per attempt so a failover can re-issue it to another replica. A
89    /// `Conflict` is a valid `Ok` outcome and returns immediately (not a failover).
90    pub async fn put(&self, record: Record, expected: Option<Revision>) -> Result<PutResult> {
91        self.with_failover(|s| {
92            let record = record.clone();
93            let expected = expected.clone();
94            async move { s.put(record, expected).await }
95        })
96        .await
97    }
98
99    /// `delete` with round-robin start + failover across all replicas. A
100    /// `Conflict` is a valid `Ok` outcome and returns immediately (not a failover).
101    /// A failover after a delete that did commit on the dead replica is safe: the
102    /// retry sees a tombstone and returns `Deleted` without writing.
103    pub async fn delete(
104        &self,
105        key: &RecordKey,
106        expected: Option<Revision>,
107    ) -> Result<DeleteResult> {
108        self.with_failover(|s| {
109            let key = key.clone();
110            let expected = expected.clone();
111            async move { s.delete(&key, expected).await }
112        })
113        .await
114    }
115
116    /// `get_raw` (tombstones visible) with round-robin start + failover.
117    pub async fn get_raw(&self, key: &RecordKey) -> Result<Option<Record>> {
118        self.with_failover(|s| {
119            let key = key.clone();
120            async move { s.get_raw(&key).await }
121        })
122        .await
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use async_trait::async_trait;
130    use gonzalo_core::{
131        Body, CoreError, DeleteResult, Identity, KeyPrefix, Meta, RecordKind, store::Conflict,
132    };
133    use std::collections::BTreeMap;
134    use std::sync::atomic::AtomicBool;
135
136    /// A `Store` double whose liveness and canned outcome are controllable.
137    struct MockStore {
138        alive: AtomicBool,
139        calls: AtomicUsize,
140    }
141    impl MockStore {
142        fn alive() -> Arc<Self> {
143            Arc::new(Self {
144                alive: AtomicBool::new(true),
145                calls: AtomicUsize::new(0),
146            })
147        }
148        fn dead() -> Arc<Self> {
149            Arc::new(Self {
150                alive: AtomicBool::new(false),
151                calls: AtomicUsize::new(0),
152            })
153        }
154        fn calls(&self) -> usize {
155            self.calls.load(Ordering::SeqCst)
156        }
157    }
158    #[async_trait]
159    impl Store for MockStore {
160        async fn get(&self, _key: &RecordKey) -> Result<Option<Record>> {
161            self.calls.fetch_add(1, Ordering::SeqCst);
162            if self.alive.load(Ordering::SeqCst) {
163                Ok(None)
164            } else {
165                Err(CoreError::Backend("connection refused".into()))
166            }
167        }
168        async fn put(&self, _record: Record, _expected: Option<Revision>) -> Result<PutResult> {
169            self.calls.fetch_add(1, Ordering::SeqCst);
170            if self.alive.load(Ordering::SeqCst) {
171                Ok(PutResult::Committed(Revision::initial(b"x")))
172            } else {
173                Err(CoreError::Backend("connection refused".into()))
174            }
175        }
176        async fn list(&self, _prefix: &KeyPrefix) -> Result<Vec<RecordKey>> {
177            Ok(Vec::new())
178        }
179        async fn delete_as(
180            &self,
181            _key: &RecordKey,
182            _expected: Option<Revision>,
183            _author: Option<Identity>,
184        ) -> Result<DeleteResult> {
185            self.calls.fetch_add(1, Ordering::SeqCst);
186            if self.alive.load(Ordering::SeqCst) {
187                Ok(DeleteResult::Deleted)
188            } else {
189                Err(CoreError::Backend("connection refused".into()))
190            }
191        }
192        async fn put_raw(&self, record: Record, expected: Option<Revision>) -> Result<PutResult> {
193            <Self as Store>::put(self, record, expected).await
194        }
195        async fn get_raw(&self, key: &RecordKey) -> Result<Option<Record>> {
196            self.get(key).await
197        }
198        async fn list_raw(&self, prefix: &KeyPrefix) -> Result<Vec<RecordKey>> {
199            self.list(prefix).await
200        }
201        async fn purge(&self, key: &RecordKey, expected: Revision) -> Result<DeleteResult> {
202            self.delete(key, Some(expected)).await
203        }
204    }
205
206    fn as_store(m: &Arc<MockStore>) -> Arc<dyn Store> {
207        m.clone()
208    }
209
210    fn rec() -> Record {
211        let body = Body::Inline(b"x".to_vec());
212        Record {
213            revision: Revision::initial(body.bytes()),
214            parent: None,
215            body,
216            kind: RecordKind::Topic,
217            meta: Meta {
218                author: Identity::new("soak"),
219                origin_system: "soak".into(),
220                created: 0,
221                updated: 0,
222                labels: BTreeMap::new(),
223            },
224            links: Vec::new(),
225            key: RecordKey::new("ns", "col", "k"),
226            ancestors: Vec::new(),
227            deleted_at: None,
228            deleted_blob: None,
229            deleted_kind: None,
230        }
231    }
232
233    /// A live replica that answers `put` and `delete_as` with fixed results and
234    /// counts calls. Every other method answers empty.
235    struct Answering {
236        calls: AtomicUsize,
237        put: fn() -> Result<PutResult>,
238        delete: fn() -> Result<DeleteResult>,
239    }
240    impl Answering {
241        fn new(put: fn() -> Result<PutResult>, delete: fn() -> Result<DeleteResult>) -> Arc<Self> {
242            Arc::new(Self {
243                calls: AtomicUsize::new(0),
244                put,
245                delete,
246            })
247        }
248    }
249    #[async_trait]
250    impl Store for Answering {
251        async fn get(&self, _key: &RecordKey) -> Result<Option<Record>> {
252            Ok(None)
253        }
254        async fn put(&self, _record: Record, _expected: Option<Revision>) -> Result<PutResult> {
255            self.calls.fetch_add(1, Ordering::SeqCst);
256            (self.put)()
257        }
258        async fn list(&self, _prefix: &KeyPrefix) -> Result<Vec<RecordKey>> {
259            Ok(Vec::new())
260        }
261        async fn delete_as(
262            &self,
263            _key: &RecordKey,
264            _expected: Option<Revision>,
265            _author: Option<Identity>,
266        ) -> Result<DeleteResult> {
267            self.calls.fetch_add(1, Ordering::SeqCst);
268            (self.delete)()
269        }
270        async fn get_raw(&self, _key: &RecordKey) -> Result<Option<Record>> {
271            Ok(None)
272        }
273        async fn put_raw(&self, _record: Record, _expected: Option<Revision>) -> Result<PutResult> {
274            Ok(PutResult::Committed(Revision::initial(b"x")))
275        }
276        async fn list_raw(&self, _prefix: &KeyPrefix) -> Result<Vec<RecordKey>> {
277            Ok(Vec::new())
278        }
279        async fn purge(&self, _key: &RecordKey, _expected: Revision) -> Result<DeleteResult> {
280            Ok(DeleteResult::Deleted)
281        }
282    }
283
284    #[tokio::test]
285    async fn failover_skips_dead_replica() {
286        let dead = MockStore::dead();
287        let live = MockStore::alive();
288        let d = Dispatcher::new(vec![as_store(&dead), as_store(&live)]);
289        // Round-robin starts at replica 0 (dead) → must fail over to 1 (live).
290        let out = d.get(&RecordKey::new("ns", "col", "k")).await;
291        assert!(
292            out.is_ok(),
293            "should have failed over to the live replica: {out:?}"
294        );
295        assert_eq!(dead.calls(), 1, "dead replica was tried");
296        assert_eq!(live.calls(), 1, "then the live replica served it");
297    }
298
299    #[tokio::test]
300    async fn all_dead_returns_error() {
301        let a = MockStore::dead();
302        let b = MockStore::dead();
303        let d = Dispatcher::new(vec![as_store(&a), as_store(&b)]);
304        assert!(d.get(&RecordKey::new("ns", "col", "k")).await.is_err());
305        assert_eq!(a.calls() + b.calls(), 2, "both replicas were attempted");
306    }
307
308    #[tokio::test]
309    async fn live_replica_short_circuits() {
310        let a = MockStore::alive();
311        let b = MockStore::alive();
312        let c = MockStore::alive();
313        let d = Dispatcher::new(vec![as_store(&a), as_store(&b), as_store(&c)]);
314        let _ = d.get(&RecordKey::new("ns", "col", "k")).await.unwrap();
315        assert_eq!(
316            a.calls() + b.calls() + c.calls(),
317            1,
318            "exactly one replica served the op"
319        );
320    }
321
322    #[tokio::test]
323    async fn round_robin_distributes_load() {
324        let a = MockStore::alive();
325        let b = MockStore::alive();
326        let c = MockStore::alive();
327        let d = Dispatcher::new(vec![as_store(&a), as_store(&b), as_store(&c)]);
328        for _ in 0..3 {
329            d.get(&RecordKey::new("ns", "col", "k")).await.unwrap();
330        }
331        assert_eq!(
332            (a.calls(), b.calls(), c.calls()),
333            (1, 1, 1),
334            "round-robin over 3 replicas"
335        );
336    }
337
338    #[tokio::test]
339    async fn put_fails_over_too() {
340        let dead = MockStore::dead();
341        let live = MockStore::alive();
342        let d = Dispatcher::new(vec![as_store(&dead), as_store(&live)]);
343        let out = d.put(rec(), None).await;
344        assert!(
345            matches!(out, Ok(PutResult::Committed(_))),
346            "put failed over: {out:?}"
347        );
348    }
349
350    #[tokio::test]
351    async fn delete_fails_over_too() {
352        let dead = MockStore::dead();
353        let live = MockStore::alive();
354        let d = Dispatcher::new(vec![as_store(&dead), as_store(&live)]);
355        let out = d.delete(&RecordKey::new("ns", "col", "k"), None).await;
356        assert!(
357            matches!(out, Ok(DeleteResult::Deleted)),
358            "delete failed over: {out:?}"
359        );
360        assert_eq!(dead.calls(), 1, "dead replica was tried");
361        assert_eq!(live.calls(), 1, "then the live replica served it");
362    }
363
364    #[tokio::test]
365    async fn get_raw_fails_over_too() {
366        let dead = MockStore::dead();
367        let live = MockStore::alive();
368        let d = Dispatcher::new(vec![as_store(&dead), as_store(&live)]);
369        let out = d.get_raw(&RecordKey::new("ns", "col", "k")).await;
370        assert!(matches!(out, Ok(None)), "get_raw failed over: {out:?}");
371        assert_eq!(dead.calls() + live.calls(), 2);
372    }
373
374    #[tokio::test]
375    async fn replicas_exposes_every_handle_in_order() {
376        let a = MockStore::alive();
377        let b = MockStore::dead();
378        let d = Dispatcher::new(vec![as_store(&a), as_store(&b)]);
379        assert_eq!(d.replicas().len(), 2);
380        // Reading each handle directly bypasses round-robin and failover.
381        assert!(
382            d.replicas()[0]
383                .get_raw(&RecordKey::new("ns", "col", "k"))
384                .await
385                .is_ok()
386        );
387        assert!(
388            d.replicas()[1]
389                .get_raw(&RecordKey::new("ns", "col", "k"))
390                .await
391                .is_err()
392        );
393    }
394
395    // A conflict from a live replica is returned as-is, never retried elsewhere.
396    #[tokio::test]
397    async fn conflict_is_returned_not_retried() {
398        let c0 = Answering::new(
399            || {
400                Ok(PutResult::Conflict(Box::new(Conflict {
401                    key: RecordKey::new("ns", "col", "k"),
402                    expected: None,
403                    current: rec(),
404                })))
405            },
406            || Ok(DeleteResult::Deleted),
407        );
408        let c1 = Answering::new(
409            || {
410                Ok(PutResult::Conflict(Box::new(Conflict {
411                    key: RecordKey::new("ns", "col", "k"),
412                    expected: None,
413                    current: rec(),
414                })))
415            },
416            || Ok(DeleteResult::Deleted),
417        );
418        let d = Dispatcher::new(vec![
419            c0.clone() as Arc<dyn Store>,
420            c1.clone() as Arc<dyn Store>,
421        ]);
422        let out = d.put(rec(), None).await;
423        assert!(
424            matches!(out, Ok(PutResult::Conflict(_))),
425            "conflict returned: {out:?}"
426        );
427        assert_eq!(
428            c0.calls.load(Ordering::SeqCst) + c1.calls.load(Ordering::SeqCst),
429            1,
430            "a conflict is a definitive answer — not retried on another replica"
431        );
432    }
433
434    // A delete conflict from a live replica is a definitive answer, never retried.
435    #[tokio::test]
436    async fn delete_conflict_is_returned_not_retried() {
437        let c0 = Answering::new(
438            || Ok(PutResult::Committed(Revision::initial(b"x"))),
439            || {
440                Ok(DeleteResult::Conflict(Box::new(Conflict {
441                    key: RecordKey::new("ns", "col", "k"),
442                    expected: Some(Revision::initial(b"stale")),
443                    current: rec(),
444                })))
445            },
446        );
447        let c1 = Answering::new(
448            || Ok(PutResult::Committed(Revision::initial(b"x"))),
449            || {
450                Ok(DeleteResult::Conflict(Box::new(Conflict {
451                    key: RecordKey::new("ns", "col", "k"),
452                    expected: Some(Revision::initial(b"stale")),
453                    current: rec(),
454                })))
455            },
456        );
457        let d = Dispatcher::new(vec![
458            c0.clone() as Arc<dyn Store>,
459            c1.clone() as Arc<dyn Store>,
460        ]);
461        let out = d
462            .delete(
463                &RecordKey::new("ns", "col", "k"),
464                Some(Revision::initial(b"stale")),
465            )
466            .await;
467        assert!(
468            matches!(out, Ok(DeleteResult::Conflict(_))),
469            "conflict returned: {out:?}"
470        );
471        assert_eq!(
472            c0.calls.load(Ordering::SeqCst) + c1.calls.load(Ordering::SeqCst),
473            1,
474            "a delete conflict is not retried on another replica"
475        );
476    }
477
478    // A `NotFound` from a live replica is an answer, not a failover trigger.
479    #[tokio::test]
480    async fn not_found_is_returned_not_retried() {
481        let c0 = Answering::new(
482            || Err(CoreError::NotFound(RecordKey::new("ns", "col", "k"))),
483            || Ok(DeleteResult::Deleted),
484        );
485        let c1 = Answering::new(
486            || Err(CoreError::NotFound(RecordKey::new("ns", "col", "k"))),
487            || Ok(DeleteResult::Deleted),
488        );
489        let d = Dispatcher::new(vec![
490            c0.clone() as Arc<dyn Store>,
491            c1.clone() as Arc<dyn Store>,
492        ]);
493        let out = d.put(rec(), None).await;
494        assert!(
495            matches!(out, Err(CoreError::NotFound(_))),
496            "not found returned: {out:?}"
497        );
498        assert_eq!(
499            c0.calls.load(Ordering::SeqCst) + c1.calls.load(Ordering::SeqCst),
500            1,
501            "a NotFound answer from a live replica is not retried on another replica"
502        );
503    }
504}
gonzalo_soak/dispatch.rs

gonzalo_soak/
dispatch.rs
