Skip to main content

crafty_client/
batch.rs

1//! Cross-shard keyed propose helpers (Tier 1 — best-effort batch, not atomic).
2
3use crate::{ClientError, KeyedClient};
4
5/// One keyed write in a [`propose_keyed_batch`].
6#[derive(Debug, Clone)]
7pub struct KeyedBatchStep {
8    /// Shard routing key.
9    pub key: Vec<u8>,
10    /// Application-encoded command bytes.
11    pub payload: Vec<u8>,
12}
13
14/// Why a multi-shard batch stopped early.
15#[derive(Debug, thiserror::Error)]
16pub enum BatchError {
17    /// Step `step` failed after `completed` prior steps succeeded.
18    #[error("step {step} failed after {completed} successful step(s): {source}")]
19    Partial {
20        /// Index of the failing step.
21        step: usize,
22        /// Number of steps that completed successfully.
23        completed: usize,
24        /// The underlying client error.
25        #[source]
26        source: ClientError,
27    },
28}
29
30/// Propose each step to its owning Raft group in order. **Not atomic** across
31/// groups — if a later step fails, earlier writes remain committed (saga-style
32/// callers must compensate explicitly).
33///
34/// # Errors
35/// Returns [`BatchError::Partial`] on the first failing step.
36pub async fn propose_keyed_batch<C: KeyedClient>(
37    client: &C,
38    steps: &[KeyedBatchStep],
39) -> Result<Vec<Vec<u8>>, BatchError> {
40    let mut out = Vec::with_capacity(steps.len());
41    for (step, item) in steps.iter().enumerate() {
42        match client
43            .propose_keyed(item.key.clone(), item.payload.clone())
44            .await
45        {
46            Ok(bytes) => out.push(bytes),
47            Err(source) => {
48                return Err(BatchError::Partial {
49                    step,
50                    completed: step,
51                    source,
52                });
53            }
54        }
55    }
56    Ok(out)
57}
58
59#[cfg(test)]
60mod tests {
61    use std::sync::{Arc, Mutex};
62
63    use crafty_net::{Route, Transport, TransportError, encode_body};
64    use crafty_proto::{ClientResponse, NodeId};
65
66    use super::*;
67    use crate::{RemoteClient, RetryPolicy};
68
69    struct ScriptTransport {
70        ok: Mutex<u32>,
71    }
72
73    impl Transport for ScriptTransport {
74        fn send(
75            &self,
76            _peer: NodeId,
77            _route: Route,
78            _body: crafty_net::transport::Body,
79        ) -> crafty_net::transport::BoxFuture<
80            'static,
81            Result<crafty_net::transport::Body, TransportError>,
82        > {
83            let remaining = {
84                let mut n = self.ok.lock().expect("lock");
85                if *n == 0 {
86                    return Box::pin(async move { Err(TransportError::Unreachable(NodeId(1))) });
87                }
88                *n -= 1;
89                *n
90            };
91            let _ = remaining;
92            Box::pin(async move {
93                encode_body(&ClientResponse::Ok(b"ok".to_vec())).map_err(TransportError::Wire)
94            })
95        }
96    }
97
98    #[tokio::test]
99    async fn batch_stops_on_first_failure_and_reports_partial_progress() {
100        let client =
101            RemoteClient::new(Arc::new(ScriptTransport { ok: Mutex::new(1) }), [NodeId(1)])
102                .with_retry(RetryPolicy {
103                    max_attempts: 1,
104                    ..RetryPolicy::default()
105                });
106        let err = propose_keyed_batch(
107            &client,
108            &[
109                KeyedBatchStep {
110                    key: b"a".to_vec(),
111                    payload: vec![1],
112                },
113                KeyedBatchStep {
114                    key: b"b".to_vec(),
115                    payload: vec![2],
116                },
117            ],
118        )
119        .await
120        .unwrap_err();
121        assert!(matches!(
122            err,
123            BatchError::Partial {
124                step: 1,
125                completed: 1,
126                ..
127            }
128        ));
129    }
130}