crafty_client/
two_phase.rs

1//! Cross-shard two-phase commit coordinator (optional Tier 2 increment).
2
3use std::future::Future;
4use std::pin::Pin;
5
6use crafty_core::{TwoPhasePlan, TwoPhasePlanError, validate_two_phase_plan};
7
8use crate::{ClientError, KeyedClient};
9
10/// Extension of [`KeyedClient`] for limited cross-shard 2PC.
11pub trait TwoPhaseClient: KeyedClient {
12    /// Stage a command on the shard for `key` under `tx_id`.
13    fn prepare_keyed(
14        &self,
15        tx_id: Vec<u8>,
16        key: Vec<u8>,
17        command: Vec<u8>,
18    ) -> impl Future<Output = Result<Vec<u8>, ClientError>> + Send;
19
20    /// Commit a previously prepared command.
21    fn commit_keyed(
22        &self,
23        tx_id: Vec<u8>,
24        key: Vec<u8>,
25    ) -> impl Future<Output = Result<Vec<u8>, ClientError>> + Send;
26
27    /// Abort a prepared command and release staging state.
28    fn abort_keyed(
29        &self,
30        tx_id: Vec<u8>,
31        key: Vec<u8>,
32    ) -> impl Future<Output = Result<Vec<u8>, ClientError>> + Send;
33}
34
35/// Why a cross-shard 2PC attempt failed.
36#[derive(Debug, thiserror::Error)]
37pub enum TwoPhaseError {
38    /// The coordinator plan was invalid (duplicate keys, empty steps, etc.).
39    #[error("invalid 2PC plan: {0}")]
40    Plan(#[from] TwoPhasePlanError),
41    /// Durable journal read/write failed before or during the transaction.
42    #[error("2PC journal error: {0}")]
43    Journal(#[from] TwoPhaseJournalError),
44    /// A prepare RPC failed after earlier steps succeeded.
45    #[error("2PC prepare failed at step {step} after {prepared} prepare(s): {source}")]
46    Prepare {
47        /// Zero-based step index that failed.
48        step: usize,
49        /// Prepare steps that succeeded before the failure.
50        prepared: usize,
51        #[source]
52        /// Underlying client / transport error.
53        source: ClientError,
54    },
55    /// A commit RPC failed after earlier steps succeeded.
56    #[error("2PC commit failed at step {step} after {committed} commit(s): {source}")]
57    Commit {
58        /// Zero-based step index that failed.
59        step: usize,
60        /// Commit steps that succeeded before the failure.
61        committed: usize,
62        #[source]
63        /// Underlying client / transport error.
64        source: ClientError,
65    },
66}
67
68/// Journal persistence failure.
69#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
70pub enum TwoPhaseJournalError {
71    /// Journal record encode/decode failed.
72    #[error("journal codec error: {0}")]
73    Codec(String),
74    /// Underlying journal storage backend failed.
75    #[error("journal backend error: {0}")]
76    Backend(String),
77}
78
79/// Lifecycle events for metrics / logging (see cross-shard-transactions ADR).
80#[derive(Debug, Clone, Copy, PartialEq, Eq)]
81pub enum TwoPhaseEvent {
82    /// All prepare steps succeeded (ready for commit phase).
83    Prepared {
84        /// Number of prepared steps.
85        steps: usize,
86    },
87    /// Coordinator stuck after partial progress (abort or commit failed).
88    Stuck {
89        /// Steps prepared before failure.
90        prepared: usize,
91        /// Step index where failure occurred.
92        failed_step: usize,
93    },
94}
95
96/// Client-side durable progress for cross-shard 2PC resume.
97#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
98pub struct TwoPhaseJournalRecord {
99    /// Shared transaction id (matches [`TwoPhasePlan::tx_id`]).
100    pub tx_id: Vec<u8>,
101    /// Count of consecutive prepared steps from step 0.
102    pub prepared_steps: u32,
103    /// Count of consecutive committed steps from step 0.
104    pub committed_steps: u32,
105}
106
107/// Optional journal hook for [`propose_cross_shard_2pc`] / [`resume_cross_shard_2pc`].
108pub trait TwoPhaseJournal: Send + Sync {
109    /// Persist progress after a successful prepare step.
110    fn on_prepared<'a>(
111        &'a self,
112        tx_id: &'a [u8],
113        step: usize,
114        total: usize,
115    ) -> Pin<Box<dyn Future<Output = Result<(), TwoPhaseJournalError>> + Send + 'a>>;
116
117    /// Persist progress after a successful commit step.
118    fn on_committed<'a>(
119        &'a self,
120        tx_id: &'a [u8],
121        step: usize,
122        total: usize,
123    ) -> Pin<Box<dyn Future<Output = Result<(), TwoPhaseJournalError>> + Send + 'a>>;
124
125    /// Mark the transaction fully committed (may delete the journal row).
126    fn on_completed<'a>(
127        &'a self,
128        tx_id: &'a [u8],
129    ) -> Pin<Box<dyn Future<Output = Result<(), TwoPhaseJournalError>> + Send + 'a>>;
130
131    /// Load coordinator progress for resume.
132    fn load<'a>(
133        &'a self,
134        tx_id: &'a [u8],
135    ) -> Pin<
136        Box<
137            dyn Future<Output = Result<Option<TwoPhaseJournalRecord>, TwoPhaseJournalError>>
138                + Send
139                + 'a,
140        >,
141    >;
142}
143
144/// In-memory 2PC journal (tests and single-process coordinators).
145#[derive(Debug, Default)]
146pub struct InMemoryTwoPhaseJournal {
147    records: std::sync::Mutex<Vec<TwoPhaseJournalRecord>>,
148}
149
150impl InMemoryTwoPhaseJournal {
151    /// Snapshot persisted journal records.
152    ///
153    /// # Panics
154    /// Panics if the journal lock is poisoned.
155    #[must_use]
156    pub fn records(&self) -> Vec<TwoPhaseJournalRecord> {
157        self.records.lock().expect("lock").clone()
158    }
159
160    fn upsert(&self, tx_id: &[u8], f: impl FnOnce(&mut TwoPhaseJournalRecord)) {
161        let mut guard = self.records.lock().expect("lock");
162        if let Some(rec) = guard.iter_mut().find(|r| r.tx_id == tx_id) {
163            f(rec);
164            return;
165        }
166        let mut rec = TwoPhaseJournalRecord {
167            tx_id: tx_id.to_vec(),
168            prepared_steps: 0,
169            committed_steps: 0,
170        };
171        f(&mut rec);
172        guard.push(rec);
173    }
174}
175
176impl TwoPhaseJournal for InMemoryTwoPhaseJournal {
177    fn on_prepared<'a>(
178        &'a self,
179        tx_id: &'a [u8],
180        step: usize,
181        _total: usize,
182    ) -> Pin<Box<dyn Future<Output = Result<(), TwoPhaseJournalError>> + Send + 'a>> {
183        Box::pin(async move {
184            self.upsert(tx_id, |rec| {
185                rec.prepared_steps =
186                    (u32::try_from(step).expect("step index fits u32") + 1).max(rec.prepared_steps);
187            });
188            Ok(())
189        })
190    }
191
192    fn on_committed<'a>(
193        &'a self,
194        tx_id: &'a [u8],
195        step: usize,
196        _total: usize,
197    ) -> Pin<Box<dyn Future<Output = Result<(), TwoPhaseJournalError>> + Send + 'a>> {
198        Box::pin(async move {
199            self.upsert(tx_id, |rec| {
200                rec.committed_steps = (u32::try_from(step).expect("step index fits u32") + 1)
201                    .max(rec.committed_steps);
202            });
203            Ok(())
204        })
205    }
206
207    fn on_completed<'a>(
208        &'a self,
209        tx_id: &'a [u8],
210    ) -> Pin<Box<dyn Future<Output = Result<(), TwoPhaseJournalError>> + Send + 'a>> {
211        Box::pin(async move {
212            self.upsert(tx_id, |rec| {
213                rec.committed_steps = rec.prepared_steps;
214            });
215            Ok(())
216        })
217    }
218
219    fn load<'a>(
220        &'a self,
221        tx_id: &'a [u8],
222    ) -> Pin<
223        Box<
224            dyn Future<Output = Result<Option<TwoPhaseJournalRecord>, TwoPhaseJournalError>>
225                + Send
226                + 'a,
227        >,
228    > {
229        Box::pin(async move {
230            Ok(self
231                .records
232                .lock()
233                .expect("lock")
234                .iter()
235                .find(|r| r.tx_id == tx_id)
236                .cloned())
237        })
238    }
239}
240
241/// Hooks for [`propose_cross_shard_2pc`].
242#[derive(Clone, Copy, Default)]
243pub struct RunTwoPhaseOpts<'a> {
244    /// Optional client-side journal for resume after coordinator restart.
245    pub journal: Option<&'a dyn TwoPhaseJournal>,
246    /// Metrics / logging callback.
247    pub on_event: Option<&'a (dyn Fn(TwoPhaseEvent) + Send + Sync)>,
248}
249
250/// Hooks for [`resume_cross_shard_2pc`].
251#[derive(Clone, Copy)]
252pub struct ResumeTwoPhaseOpts<'a> {
253    /// Journal that records prepared/committed prefixes.
254    pub journal: Option<&'a dyn TwoPhaseJournal>,
255    /// When `true`, try `commit_keyed` before `prepare_keyed` for unknown steps.
256    pub probe: bool,
257    /// Metrics / logging callback.
258    pub on_event: Option<&'a (dyn Fn(TwoPhaseEvent) + Send + Sync)>,
259}
260
261impl Default for ResumeTwoPhaseOpts<'_> {
262    fn default() -> Self {
263        Self {
264            journal: None,
265            probe: true,
266            on_event: None,
267        }
268    }
269}
270
271/// Postcard-encode a [`TwoPhaseJournalRecord`] for Meta-Raft / Redis storage.
272///
273/// # Errors
274/// Returns [`TwoPhaseJournalError::Codec`] when postcard encoding fails.
275pub fn encode_two_phase_journal_record(
276    record: &TwoPhaseJournalRecord,
277) -> Result<Vec<u8>, TwoPhaseJournalError> {
278    crafty_proto::encode(record).map_err(|e| TwoPhaseJournalError::Codec(e.to_string()))
279}
280
281/// Postcard-decode a [`TwoPhaseJournalRecord`].
282///
283/// # Errors
284/// Returns [`TwoPhaseJournalError::Codec`] when postcard decoding fails.
285pub fn decode_two_phase_journal_record(
286    bytes: &[u8],
287) -> Result<TwoPhaseJournalRecord, TwoPhaseJournalError> {
288    crafty_proto::decode(bytes).map_err(|e| TwoPhaseJournalError::Codec(e.to_string()))
289}
290
291fn emit_stuck(
292    on_event: Option<&(dyn Fn(TwoPhaseEvent) + Send + Sync)>,
293    prepared: usize,
294    failed_step: usize,
295) {
296    if let Some(on) = on_event {
297        on(TwoPhaseEvent::Stuck {
298            prepared,
299            failed_step,
300        });
301    }
302}
303
304async fn abort_prepared<C: TwoPhaseClient>(
305    client: &C,
306    plan: &TwoPhasePlan,
307    prepared: usize,
308) -> bool {
309    let mut ok = true;
310    for prev in plan.steps.iter().take(prepared).rev() {
311        if client
312            .abort_keyed(plan.tx_id.clone(), prev.key.clone())
313            .await
314            .is_err()
315        {
316            ok = false;
317        }
318    }
319    ok
320}
321
322fn is_no_prepared(err: &ClientError) -> bool {
323    matches!(
324        err,
325        ClientError::Server(msg) if msg.contains("no prepared command for transaction key")
326    )
327}
328
329async fn prepare_step<C: TwoPhaseClient>(
330    client: &C,
331    plan: &TwoPhasePlan,
332    step: usize,
333    journal: Option<&dyn TwoPhaseJournal>,
334) -> Result<(), TwoPhaseError> {
335    let item = &plan.steps[step];
336    client
337        .prepare_keyed(plan.tx_id.clone(), item.key.clone(), item.command.clone())
338        .await
339        .map_err(|source| TwoPhaseError::Prepare {
340            step,
341            prepared: step,
342            source,
343        })?;
344    if let Some(j) = journal {
345        j.on_prepared(&plan.tx_id, step, plan.steps.len()).await?;
346    }
347    Ok(())
348}
349
350async fn commit_step<C: TwoPhaseClient>(
351    client: &C,
352    plan: &TwoPhasePlan,
353    step: usize,
354    journal: Option<&dyn TwoPhaseJournal>,
355) -> Result<Vec<u8>, TwoPhaseError> {
356    let item = &plan.steps[step];
357    let bytes = client
358        .commit_keyed(plan.tx_id.clone(), item.key.clone())
359        .await
360        .map_err(|source| TwoPhaseError::Commit {
361            step,
362            committed: step,
363            source,
364        })?;
365    if let Some(j) = journal {
366        j.on_committed(&plan.tx_id, step, plan.steps.len()).await?;
367    }
368    Ok(bytes)
369}
370
371async fn prepare_or_commit_step<C: TwoPhaseClient>(
372    client: &C,
373    plan: &TwoPhasePlan,
374    step: usize,
375    journal: Option<&dyn TwoPhaseJournal>,
376    probe: bool,
377) -> Result<Vec<u8>, TwoPhaseError> {
378    let item = &plan.steps[step];
379    if probe {
380        match client
381            .commit_keyed(plan.tx_id.clone(), item.key.clone())
382            .await
383        {
384            Ok(bytes) => {
385                if let Some(j) = journal {
386                    j.on_prepared(&plan.tx_id, step, plan.steps.len()).await?;
387                    j.on_committed(&plan.tx_id, step, plan.steps.len()).await?;
388                }
389                return Ok(bytes);
390            }
391            Err(err) if is_no_prepared(&err) => {}
392            Err(source) => {
393                return Err(TwoPhaseError::Commit {
394                    step,
395                    committed: step,
396                    source,
397                });
398            }
399        }
400    }
401    prepare_step(client, plan, step, journal).await?;
402    commit_step(client, plan, step, journal).await
403}
404
405/// Execute prepare-all then commit-all, aborting prepared steps on prepare failure.
406///
407/// # Errors
408/// Returns [`TwoPhaseError::Plan`] when the plan is invalid,
409/// [`TwoPhaseError::Prepare`] or [`TwoPhaseError::Commit`] when a shard RPC fails,
410/// or [`TwoPhaseError::Journal`] when journal persistence fails.
411pub async fn propose_cross_shard_2pc<C: TwoPhaseClient>(
412    client: &C,
413    plan: &TwoPhasePlan,
414    group_for_key: impl Fn(&[u8]) -> Option<u32>,
415) -> Result<Vec<Vec<u8>>, TwoPhaseError> {
416    propose_cross_shard_2pc_with_opts(client, plan, group_for_key, RunTwoPhaseOpts::default()).await
417}
418
419/// Like [`propose_cross_shard_2pc`] with an optional client journal.
420///
421/// # Errors
422/// Same as [`propose_cross_shard_2pc`].
423pub async fn propose_cross_shard_2pc_with_opts<C: TwoPhaseClient>(
424    client: &C,
425    plan: &TwoPhasePlan,
426    group_for_key: impl Fn(&[u8]) -> Option<u32>,
427    opts: RunTwoPhaseOpts<'_>,
428) -> Result<Vec<Vec<u8>>, TwoPhaseError> {
429    validate_two_phase_plan(plan, group_for_key)?;
430    let journal = opts.journal;
431
432    for (step, item) in plan.steps.iter().enumerate() {
433        if let Err(source) = client
434            .prepare_keyed(plan.tx_id.clone(), item.key.clone(), item.command.clone())
435            .await
436        {
437            if !abort_prepared(client, plan, step).await {
438                emit_stuck(opts.on_event, step, step);
439            }
440            return Err(TwoPhaseError::Prepare {
441                step,
442                prepared: step,
443                source,
444            });
445        }
446        if let Some(j) = journal {
447            j.on_prepared(&plan.tx_id, step, plan.steps.len()).await?;
448        }
449    }
450
451    if let Some(on) = opts.on_event {
452        on(TwoPhaseEvent::Prepared {
453            steps: plan.steps.len(),
454        });
455    }
456
457    let mut responses = Vec::with_capacity(plan.steps.len());
458    for (step, _item) in plan.steps.iter().enumerate() {
459        match commit_step(client, plan, step, journal).await {
460            Ok(bytes) => responses.push(bytes),
461            Err(err) => {
462                emit_stuck(opts.on_event, plan.steps.len(), step);
463                return Err(err);
464            }
465        }
466    }
467    if let Some(j) = journal {
468        j.on_completed(&plan.tx_id).await?;
469    }
470    Ok(responses)
471}
472
473/// Continue a cross-shard 2PC after partial progress or coordinator restart.
474///
475/// With a [`TwoPhaseJournal`], skips consecutive prepared/committed prefixes recorded
476/// client-side. With `probe = true` (default), steps without journal state attempt
477/// `commit_keyed` first so a durable server-side prepare can be picked up after restart.
478///
479/// # Errors
480/// Returns [`TwoPhaseError::Plan`] when the plan is invalid,
481/// [`TwoPhaseError::Prepare`] or [`TwoPhaseError::Commit`] when a shard RPC fails,
482/// or [`TwoPhaseError::Journal`] when journal persistence fails.
483pub async fn resume_cross_shard_2pc<C: TwoPhaseClient>(
484    client: &C,
485    plan: &TwoPhasePlan,
486    group_for_key: impl Fn(&[u8]) -> Option<u32>,
487    opts: ResumeTwoPhaseOpts<'_>,
488) -> Result<Vec<Vec<u8>>, TwoPhaseError> {
489    validate_two_phase_plan(plan, group_for_key)?;
490    let journal = opts.journal;
491
492    let (prepared_through, committed_through) = if let Some(j) = journal {
493        match j.load(&plan.tx_id).await? {
494            Some(rec) => (rec.prepared_steps as usize, rec.committed_steps as usize),
495            None => (0, 0),
496        }
497    } else {
498        (0, 0)
499    };
500
501    let mut responses = Vec::with_capacity(plan.steps.len());
502
503    for _step in 0..committed_through.min(plan.steps.len()) {
504        responses.push(Vec::new());
505    }
506
507    for step in committed_through..prepared_through.min(plan.steps.len()) {
508        responses.push(commit_step(client, plan, step, journal).await?);
509    }
510
511    for step in prepared_through..plan.steps.len() {
512        match prepare_or_commit_step(client, plan, step, journal, opts.probe).await {
513            Ok(bytes) => responses.push(bytes),
514            Err(err) => {
515                emit_stuck(opts.on_event, step, step);
516                return Err(err);
517            }
518        }
519    }
520
521    if let Some(j) = journal {
522        j.on_completed(&plan.tx_id).await?;
523    }
524    Ok(responses)
525}
526
527#[cfg(test)]
528mod tests {
529    use std::collections::HashSet;
530    use std::sync::Arc;
531    use std::sync::Mutex;
532
533    use crafty_net::{Route, Transport, TransportError, decode_body, encode_body};
534    use crafty_proto::{ClientRequest, ClientResponse, NodeId};
535
536    use super::*;
537    use crate::{RemoteClient, RetryPolicy};
538
539    type PreparedKeys = HashSet<(Vec<u8>, Vec<u8>)>;
540
541    struct TwoPhaseScript {
542        prepared: Arc<Mutex<PreparedKeys>>,
543    }
544
545    impl Transport for TwoPhaseScript {
546        fn send(
547            &self,
548            _peer: NodeId,
549            _route: Route,
550            body: crafty_net::transport::Body,
551        ) -> crafty_net::transport::BoxFuture<
552            'static,
553            Result<crafty_net::transport::Body, TransportError>,
554        > {
555            let request = match decode_body::<ClientRequest>(&body) {
556                Ok(r) => r,
557                Err(e) => {
558                    return Box::pin(async move { Err(TransportError::Wire(e)) });
559                }
560            };
561            let prepared = Arc::clone(&self.prepared);
562            Box::pin(async move {
563                match request {
564                    ClientRequest::TwoPhasePrepare { tx_id, key, .. } => {
565                        prepared.lock().expect("lock").insert((tx_id, key));
566                        encode_body(&ClientResponse::Ok(Vec::new())).map_err(TransportError::Wire)
567                    }
568                    ClientRequest::TwoPhaseCommit { tx_id, key } => {
569                        if prepared.lock().expect("lock").contains(&(tx_id, key)) {
570                            encode_body(&ClientResponse::Ok(vec![1])).map_err(TransportError::Wire)
571                        } else {
572                            encode_body(&ClientResponse::Error(
573                                "no prepared command for transaction key".into(),
574                            ))
575                            .map_err(TransportError::Wire)
576                        }
577                    }
578                    other => Err(TransportError::Io(format!("unexpected request: {other:?}"))),
579                }
580            })
581        }
582    }
583
584    impl TwoPhaseScript {
585        fn new() -> Self {
586            Self {
587                prepared: Arc::new(Mutex::new(HashSet::new())),
588            }
589        }
590
591        fn mark_prepared(&self, tx_id: &[u8], key: &[u8]) {
592            self.prepared
593                .lock()
594                .expect("lock")
595                .insert((tx_id.to_vec(), key.to_vec()));
596        }
597    }
598
599    fn sample_plan() -> TwoPhasePlan {
600        TwoPhasePlan {
601            tx_id: b"tx-resume".to_vec(),
602            steps: vec![
603                crafty_core::TwoPhaseStep {
604                    key: b"a".to_vec(),
605                    command: vec![1],
606                },
607                crafty_core::TwoPhaseStep {
608                    key: b"b".to_vec(),
609                    command: vec![2],
610                },
611            ],
612        }
613    }
614
615    #[tokio::test]
616    async fn resume_probes_commit_before_prepare() {
617        let script = Arc::new(TwoPhaseScript::new());
618        script.mark_prepared(b"tx-resume", b"a");
619        let client = RemoteClient::new(script, [NodeId(1)]).with_retry(RetryPolicy {
620            max_attempts: 1,
621            ..Default::default()
622        });
623        let plan = sample_plan();
624        let out =
625            resume_cross_shard_2pc(&client, &plan, |_| Some(0), ResumeTwoPhaseOpts::default())
626                .await
627                .expect("resume");
628        assert_eq!(out.len(), 2);
629    }
630
631    #[tokio::test]
632    async fn resume_skips_journal_committed_prefix() {
633        let script = Arc::new(TwoPhaseScript::new());
634        script.mark_prepared(b"tx-resume", b"b");
635        let client = RemoteClient::new(script, [NodeId(1)]).with_retry(RetryPolicy {
636            max_attempts: 1,
637            ..Default::default()
638        });
639        let journal = InMemoryTwoPhaseJournal::default();
640        journal
641            .on_prepared(b"tx-resume", 0, 2)
642            .await
643            .expect("prep 0");
644        journal
645            .on_prepared(b"tx-resume", 1, 2)
646            .await
647            .expect("prep 1");
648        journal
649            .on_committed(b"tx-resume", 0, 2)
650            .await
651            .expect("commit 0");
652
653        let plan = sample_plan();
654        let out = resume_cross_shard_2pc(
655            &client,
656            &plan,
657            |_| Some(0),
658            ResumeTwoPhaseOpts {
659                journal: Some(&journal),
660                probe: false,
661                on_event: None,
662            },
663        )
664        .await
665        .expect("resume");
666        assert_eq!(out.len(), 2);
667        assert_eq!(out[1], vec![1]);
668    }
669}
crafty_client/two_phase.rs

crafty_client/
two_phase.rs
