common_procedure/
store.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::fmt;
17
18use common_telemetry::{debug, error, info, warn};
19use futures::TryStreamExt;
20use serde::{Deserialize, Serialize};
21use snafu::ResultExt;
22
23use crate::error::{Result, ToJsonSnafu};
24pub(crate) use crate::store::state_store::StateStoreRef;
25use crate::{ProcedureContext, ProcedureId, SubprocedureRecovery};
26
27pub mod poison_store;
28pub mod state_store;
29pub mod util;
30
31/// Key prefix of procedure store.
32const PROC_PATH: &str = "procedure/";
33
34/// Constructs a path for procedure store.
35macro_rules! proc_path {
36    ($store: expr, $fmt:expr) => { format!("{}{}", $store.proc_path(), format_args!($fmt)) };
37    ($store: expr, $fmt:expr, $($args:tt)*) => { format!("{}{}", $store.proc_path(), format_args!($fmt, $($args)*)) };
38}
39
40#[cfg(test)]
41pub(crate) use proc_path;
42
43/// Serialized data of a procedure.
44#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
45pub struct ProcedureMessage {
46    /// Type name of the procedure. The procedure framework also use the type name to
47    /// find a loader to load the procedure.
48    pub type_name: String,
49    /// The data of the procedure.
50    pub data: String,
51    /// Parent procedure id.
52    pub parent_id: Option<ProcedureId>,
53    /// Recovery policy for a child submission, ignored for root procedures.
54    #[serde(default, skip_serializing_if = "SubprocedureRecovery::is_discard")]
55    pub subprocedure_recovery: SubprocedureRecovery,
56    /// Current step.
57    pub step: u32,
58    /// Errors raised during the procedure.
59    #[serde(default, skip_serializing_if = "Option::is_none")]
60    pub error: Option<String>,
61    /// Context associated with the root submission.
62    #[serde(default, skip_serializing_if = "ProcedureContext::is_empty")]
63    pub context: ProcedureContext,
64}
65
66/// A collection of all procedures' messages.
67#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
68pub struct ProcedureMessages {
69    /// A map of uncommitted procedures
70    pub messages: HashMap<ProcedureId, ProcedureMessage>,
71    /// A map of rolling back procedures
72    pub rollback_messages: HashMap<ProcedureId, ProcedureMessage>,
73    /// A list of finished procedures' ids
74    pub finished_ids: Vec<ProcedureId>,
75}
76
77/// Procedure storage layer.
78pub(crate) struct ProcedureStore {
79    proc_path: String,
80    store: StateStoreRef,
81}
82
83impl ProcedureStore {
84    /// Creates a new [ProcedureStore] from specific [StateStoreRef].
85    pub(crate) fn new(parent_path: &str, store: StateStoreRef) -> ProcedureStore {
86        let proc_path = format!("{}{PROC_PATH}", parent_path);
87        info!("The procedure state store path is: {}", &proc_path);
88        ProcedureStore { proc_path, store }
89    }
90
91    #[inline]
92    pub(crate) fn proc_path(&self) -> &str {
93        &self.proc_path
94    }
95
96    /// Dump the `procedure` to the storage.
97    #[cfg(test)]
98    pub(crate) async fn store_procedure(
99        &self,
100        procedure_id: ProcedureId,
101        step: u32,
102        type_name: String,
103        data: String,
104        parent_id: Option<ProcedureId>,
105    ) -> Result<()> {
106        self.store_procedure_with_context(
107            procedure_id,
108            step,
109            type_name,
110            data,
111            parent_id,
112            ProcedureContext::default(),
113        )
114        .await
115    }
116
117    #[cfg(test)]
118    pub(crate) async fn store_procedure_with_context(
119        &self,
120        procedure_id: ProcedureId,
121        step: u32,
122        type_name: String,
123        data: String,
124        parent_id: Option<ProcedureId>,
125        context: ProcedureContext,
126    ) -> Result<()> {
127        let message = ProcedureMessage {
128            type_name,
129            data,
130            parent_id,
131            subprocedure_recovery: SubprocedureRecovery::Discard,
132            step,
133            error: None,
134            context,
135        };
136        self.store_procedure_message(procedure_id, &message).await
137    }
138
139    /// Persists a procedure checkpoint with its recovery policy and submission context.
140    pub(crate) async fn store_procedure_message(
141        &self,
142        procedure_id: ProcedureId,
143        message: &ProcedureMessage,
144    ) -> Result<()> {
145        let key = ParsedKey {
146            prefix: &self.proc_path,
147            procedure_id,
148            step: message.step,
149            key_type: KeyType::Step,
150        }
151        .to_string();
152        let value = serde_json::to_string(&message).context(ToJsonSnafu)?;
153
154        self.store.put(&key, value.into_bytes()).await?;
155
156        Ok(())
157    }
158
159    /// Write commit flag to the storage.
160    pub(crate) async fn commit_procedure(
161        &self,
162        procedure_id: ProcedureId,
163        step: u32,
164    ) -> Result<()> {
165        let key = ParsedKey {
166            prefix: &self.proc_path,
167            procedure_id,
168            step,
169            key_type: KeyType::Commit,
170        }
171        .to_string();
172        self.store.put(&key, Vec::new()).await?;
173
174        Ok(())
175    }
176
177    /// Write rollback flag to the storage.
178    pub(crate) async fn rollback_procedure(
179        &self,
180        procedure_id: ProcedureId,
181        message: ProcedureMessage,
182    ) -> Result<()> {
183        let key = ParsedKey {
184            prefix: &self.proc_path,
185            procedure_id,
186            step: message.step,
187            key_type: KeyType::Rollback,
188        }
189        .to_string();
190
191        self.store
192            .put(&key, serde_json::to_vec(&message).context(ToJsonSnafu)?)
193            .await?;
194
195        Ok(())
196    }
197
198    /// Delete states of procedure from the storage.
199    pub(crate) async fn delete_procedure(&self, procedure_id: ProcedureId) -> Result<()> {
200        let path = proc_path!(self, "{procedure_id}/");
201        // TODO(yingwen): We can optimize this to avoid reading the value.
202        let mut key_values = self.store.walk_top_down(&path).await?;
203        // 8 should be enough for most procedures.
204        let mut step_keys = Vec::with_capacity(8);
205        let mut finish_keys = Vec::new();
206        while let Some((key_set, _)) = key_values.try_next().await? {
207            let key = key_set.key();
208            let Some(curr_key) = ParsedKey::parse_str(&self.proc_path, key) else {
209                warn!("Unknown key while deleting procedures, key: {}", key);
210                continue;
211            };
212            if curr_key.key_type == KeyType::Step {
213                step_keys.extend(key_set.keys());
214            } else {
215                // .commit or .rollback
216                finish_keys.extend(key_set.keys());
217            }
218        }
219
220        debug!(
221            "Delete keys for procedure {}, step_keys: {:?}, finish_keys: {:?}",
222            procedure_id, step_keys, finish_keys
223        );
224        // We delete all step keys first.
225        self.store.batch_delete(step_keys.as_slice()).await?;
226        // Then we delete the finish keys, to ensure
227        self.store.batch_delete(finish_keys.as_slice()).await?;
228        // Finally we remove the directory itself.
229        self.store.delete(&path).await?;
230        // Maybe we could use procedure_id.commit/rollback as the file name so we could
231        // use remove_all to remove the directory and then remove the commit/rollback file.
232
233        Ok(())
234    }
235
236    /// Load procedures from the storage.
237    pub(crate) async fn load_messages(&self) -> Result<ProcedureMessages> {
238        // Track the key-value pair by procedure id.
239        let mut procedure_key_values: HashMap<_, (ParsedKey, Vec<u8>)> = HashMap::new();
240
241        // Scan all procedures.
242        let mut key_values = self.store.walk_top_down(&self.proc_path).await?;
243        while let Some((key_set, value)) = key_values.try_next().await? {
244            let key = key_set.key();
245            let Some(curr_key) = ParsedKey::parse_str(&self.proc_path, key) else {
246                warn!("Unknown key while loading procedures, key: {}", key);
247                continue;
248            };
249
250            if let Some(entry) = procedure_key_values.get_mut(&curr_key.procedure_id) {
251                if entry.0.step < curr_key.step {
252                    entry.0 = curr_key;
253                    entry.1 = value;
254                }
255            } else {
256                let _ = procedure_key_values.insert(curr_key.procedure_id, (curr_key, value));
257            }
258        }
259
260        let mut messages = HashMap::with_capacity(procedure_key_values.len());
261        let mut rollback_messages = HashMap::new();
262        let mut finished_ids = Vec::new();
263        for (procedure_id, (parsed_key, value)) in procedure_key_values {
264            match parsed_key.key_type {
265                KeyType::Step => {
266                    let Some(message) = self.load_one_message(&parsed_key, &value) else {
267                        // We don't abort the loading process and just ignore errors to ensure all remaining
268                        // procedures are loaded.
269                        continue;
270                    };
271                    let _ = messages.insert(procedure_id, message);
272                }
273                KeyType::Commit => {
274                    finished_ids.push(procedure_id);
275                }
276                KeyType::Rollback => {
277                    let Some(message) = self.load_one_message(&parsed_key, &value) else {
278                        // We don't abort the loading process and just ignore errors to ensure all remaining
279                        // procedures are loaded.
280                        continue;
281                    };
282                    let _ = rollback_messages.insert(procedure_id, message);
283                }
284            }
285        }
286
287        Ok(ProcedureMessages {
288            messages,
289            rollback_messages,
290            finished_ids,
291        })
292    }
293
294    fn load_one_message(&self, key: &ParsedKey, value: &[u8]) -> Option<ProcedureMessage> {
295        serde_json::from_slice(value)
296            .map_err(|e| {
297                // `e` doesn't impl ErrorExt so we print it as normal error.
298                error!("Failed to parse value, key: {:?}, source: {:?}", key, e);
299                e
300            })
301            .ok()
302    }
303}
304
305/// Suffix type of the key.
306#[derive(Debug, PartialEq, Eq)]
307enum KeyType {
308    Step,
309    Commit,
310    Rollback,
311}
312
313impl KeyType {
314    fn as_str(&self) -> &'static str {
315        match self {
316            KeyType::Step => "step",
317            KeyType::Commit => "commit",
318            KeyType::Rollback => "rollback",
319        }
320    }
321
322    fn from_str(s: &str) -> Option<KeyType> {
323        match s {
324            "step" => Some(KeyType::Step),
325            "commit" => Some(KeyType::Commit),
326            "rollback" => Some(KeyType::Rollback),
327            _ => None,
328        }
329    }
330}
331
332/// Key to refer the procedure in the [ProcedureStore].
333#[derive(Debug, PartialEq, Eq)]
334struct ParsedKey<'a> {
335    prefix: &'a str,
336    procedure_id: ProcedureId,
337    step: u32,
338    key_type: KeyType,
339}
340
341impl fmt::Display for ParsedKey<'_> {
342    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
343        write!(
344            f,
345            "{}{}/{:010}.{}",
346            self.prefix,
347            self.procedure_id,
348            self.step,
349            self.key_type.as_str(),
350        )
351    }
352}
353
354impl<'a> ParsedKey<'a> {
355    /// Try to parse the key from specific `input`.
356    fn parse_str(prefix: &'a str, input: &str) -> Option<ParsedKey<'a>> {
357        let input = input.strip_prefix(prefix)?;
358        let mut iter = input.rsplit('/');
359        let name = iter.next()?;
360        let id_str = iter.next()?;
361
362        let procedure_id = ProcedureId::parse_str(id_str).ok()?;
363
364        let mut parts = name.split('.');
365        let step_str = parts.next()?;
366        let suffix = parts.next()?;
367        let key_type = KeyType::from_str(suffix)?;
368        let step = step_str.parse().ok()?;
369
370        Some(ParsedKey {
371            prefix,
372            procedure_id,
373            step,
374            key_type,
375        })
376    }
377}
378
379#[cfg(test)]
380mod tests {
381    use std::sync::Arc;
382
383    use object_store::ObjectStore;
384
385    use crate::BoxedProcedure;
386    use crate::procedure::PoisonKeys;
387    use crate::store::state_store::ObjectStateStore;
388
389    impl ProcedureStore {
390        pub(crate) fn from_object_store(store: ObjectStore) -> ProcedureStore {
391            let state_store = ObjectStateStore::new(store);
392
393            ProcedureStore::new("data/", Arc::new(state_store))
394        }
395    }
396
397    use async_trait::async_trait;
398    use common_event_recorder::{PersistentEventContext, TriggerReason};
399    use common_test_util::temp_dir::{TempDir, create_temp_dir};
400    use object_store::services::Fs as Builder;
401
402    use super::*;
403    use crate::{Context, LockKey, Procedure, Status};
404
405    fn procedure_store_for_test(dir: &TempDir) -> ProcedureStore {
406        let store_dir = dir.path().to_str().unwrap();
407        let builder = Builder::default().root(store_dir);
408        let object_store = ObjectStore::new(builder).unwrap();
409
410        ProcedureStore::from_object_store(object_store)
411    }
412
413    #[test]
414    fn test_parsed_key() {
415        let dir = create_temp_dir("store_procedure");
416        let store = procedure_store_for_test(&dir);
417
418        let procedure_id = ProcedureId::random();
419        let key = ParsedKey {
420            prefix: &store.proc_path,
421            procedure_id,
422            step: 2,
423            key_type: KeyType::Step,
424        };
425        assert_eq!(
426            proc_path!(store, "{procedure_id}/0000000002.step"),
427            key.to_string()
428        );
429        assert_eq!(
430            key,
431            ParsedKey::parse_str(&store.proc_path, &key.to_string()).unwrap()
432        );
433
434        let key = ParsedKey {
435            prefix: &store.proc_path,
436            procedure_id,
437            step: 2,
438            key_type: KeyType::Commit,
439        };
440        assert_eq!(
441            proc_path!(store, "{procedure_id}/0000000002.commit"),
442            key.to_string()
443        );
444        assert_eq!(
445            key,
446            ParsedKey::parse_str(&store.proc_path, &key.to_string()).unwrap()
447        );
448
449        let key = ParsedKey {
450            prefix: &store.proc_path,
451            procedure_id,
452            step: 2,
453            key_type: KeyType::Rollback,
454        };
455        assert_eq!(
456            proc_path!(store, "{procedure_id}/0000000002.rollback"),
457            key.to_string()
458        );
459        assert_eq!(
460            key,
461            ParsedKey::parse_str(&store.proc_path, &key.to_string()).unwrap()
462        );
463    }
464
465    #[test]
466    fn test_parse_invalid_key() {
467        let dir = create_temp_dir("store_procedure");
468        let store = procedure_store_for_test(&dir);
469
470        assert!(ParsedKey::parse_str(&store.proc_path, "").is_none());
471        assert!(ParsedKey::parse_str(&store.proc_path, "invalidprefix").is_none());
472        assert!(ParsedKey::parse_str(&store.proc_path, "procedu/0000000003.step").is_none());
473        assert!(ParsedKey::parse_str(&store.proc_path, "procedure-0000000003.step").is_none());
474
475        let procedure_id = ProcedureId::random();
476        let input = proc_path!(store, "{procedure_id}");
477        assert!(ParsedKey::parse_str(&store.proc_path, &input).is_none());
478
479        let input = proc_path!(store, "{procedure_id}");
480        assert!(ParsedKey::parse_str(&store.proc_path, &input).is_none());
481
482        let input = proc_path!(store, "{procedure_id}/0000000003");
483        assert!(ParsedKey::parse_str(&store.proc_path, &input).is_none());
484
485        let input = proc_path!(store, "{procedure_id}/0000000003.");
486        assert!(ParsedKey::parse_str(&store.proc_path, &input).is_none());
487
488        let input = proc_path!(store, "{procedure_id}/0000000003.other");
489        assert!(ParsedKey::parse_str(&store.proc_path, &input).is_none());
490
491        assert!(ParsedKey::parse_str(&store.proc_path, "12345/0000000003.step").is_none());
492
493        let input = proc_path!(store, "{procedure_id}-0000000003.commit");
494        assert!(ParsedKey::parse_str(&store.proc_path, &input).is_none());
495    }
496
497    #[test]
498    fn test_procedure_message() {
499        let mut message = ProcedureMessage {
500            type_name: "TestMessage".to_string(),
501            data: "no parent id".to_string(),
502            parent_id: None,
503            subprocedure_recovery: SubprocedureRecovery::Discard,
504            step: 4,
505            error: None,
506            context: ProcedureContext::default(),
507        };
508
509        let json = serde_json::to_string(&message).unwrap();
510        assert_eq!(
511            json,
512            r#"{"type_name":"TestMessage","data":"no parent id","parent_id":null,"step":4}"#
513        );
514
515        let procedure_id = ProcedureId::parse_str("9f805a1f-05f7-490c-9f91-bd56e3cc54c1").unwrap();
516        message.parent_id = Some(procedure_id);
517        let json = serde_json::to_string(&message).unwrap();
518        assert_eq!(
519            json,
520            r#"{"type_name":"TestMessage","data":"no parent id","parent_id":"9f805a1f-05f7-490c-9f91-bd56e3cc54c1","step":4}"#
521        );
522
523        message.context = ProcedureContext::from_event_context(PersistentEventContext::new(
524            TriggerReason::AutoRebalance,
525        ));
526        let json = serde_json::to_string(&message).unwrap();
527        assert_eq!(
528            json,
529            r#"{"type_name":"TestMessage","data":"no parent id","parent_id":"9f805a1f-05f7-490c-9f91-bd56e3cc54c1","step":4,"context":{"event_context":{"reason":"auto_rebalance"}}}"#
530        );
531
532        message.context.actor = Some("alice".to_string());
533        let json = serde_json::to_string(&message).unwrap();
534        assert!(json.contains(r#""actor":"alice""#));
535        assert_eq!(
536            serde_json::from_str::<ProcedureMessage>(&json).unwrap(),
537            message
538        );
539
540        let legacy: ProcedureMessage = serde_json::from_str(
541            r#"{"type_name":"TestMessage","data":"legacy","parent_id":null,"step":1}"#,
542        )
543        .unwrap();
544        assert_eq!(legacy.context, ProcedureContext::default());
545    }
546
547    #[test]
548    fn test_subprocedure_recovery_serialization() {
549        let legacy = r#"{"type_name":"TestMessage","data":"checkpoint","parent_id":null,"step":4}"#;
550        let mut message: ProcedureMessage = serde_json::from_str(legacy).unwrap();
551        assert_eq!(message.subprocedure_recovery, SubprocedureRecovery::Discard);
552        assert_eq!(serde_json::to_string(&message).unwrap(), legacy);
553        message.subprocedure_recovery = SubprocedureRecovery::Detach;
554        let json = serde_json::to_string(&message).unwrap();
555        assert_eq!(
556            serde_json::from_str::<ProcedureMessage>(&json).unwrap(),
557            message
558        );
559    }
560
561    struct MockProcedure {
562        data: String,
563    }
564
565    impl MockProcedure {
566        fn new(data: impl Into<String>) -> MockProcedure {
567            MockProcedure { data: data.into() }
568        }
569    }
570
571    #[async_trait]
572    impl Procedure for MockProcedure {
573        fn type_name(&self) -> &str {
574            "MockProcedure"
575        }
576
577        async fn execute(&mut self, _ctx: &Context) -> Result<Status> {
578            unimplemented!()
579        }
580
581        fn dump(&self) -> Result<String> {
582            Ok(self.data.clone())
583        }
584
585        fn lock_key(&self) -> LockKey {
586            LockKey::default()
587        }
588
589        fn poison_keys(&self) -> PoisonKeys {
590            PoisonKeys::default()
591        }
592    }
593
594    #[tokio::test]
595    async fn test_store_procedure() {
596        let dir = create_temp_dir("store_procedure");
597        let store = procedure_store_for_test(&dir);
598
599        let procedure_id = ProcedureId::random();
600        let procedure: BoxedProcedure = Box::new(MockProcedure::new("test store procedure"));
601        let type_name = procedure.type_name().to_string();
602        let data = procedure.dump().unwrap();
603        store
604            .store_procedure(procedure_id, 0, type_name, data, None)
605            .await
606            .unwrap();
607
608        let ProcedureMessages {
609            messages,
610            rollback_messages,
611            finished_ids,
612        } = store.load_messages().await.unwrap();
613        assert_eq!(1, messages.len());
614        assert!(rollback_messages.is_empty());
615        assert!(finished_ids.is_empty());
616        let msg = messages.get(&procedure_id).unwrap();
617        let expect = ProcedureMessage {
618            type_name: "MockProcedure".to_string(),
619            data: "test store procedure".to_string(),
620            parent_id: None,
621            subprocedure_recovery: SubprocedureRecovery::Discard,
622            step: 0,
623            error: None,
624            context: ProcedureContext::default(),
625        };
626        assert_eq!(expect, *msg);
627    }
628
629    #[tokio::test]
630    async fn test_commit_procedure() {
631        let dir = create_temp_dir("commit_procedure");
632        let store = procedure_store_for_test(&dir);
633
634        let procedure_id = ProcedureId::random();
635        let procedure: BoxedProcedure = Box::new(MockProcedure::new("test store procedure"));
636        let type_name = procedure.type_name().to_string();
637        let data = procedure.dump().unwrap();
638        store
639            .store_procedure(procedure_id, 0, type_name, data, None)
640            .await
641            .unwrap();
642        store.commit_procedure(procedure_id, 1).await.unwrap();
643
644        let ProcedureMessages {
645            messages,
646            rollback_messages,
647            finished_ids,
648        } = store.load_messages().await.unwrap();
649        assert!(messages.is_empty());
650        assert!(rollback_messages.is_empty());
651        assert_eq!(&[procedure_id], &finished_ids[..]);
652    }
653
654    #[tokio::test]
655    async fn test_rollback_procedure() {
656        let dir = create_temp_dir("rollback_procedure");
657        let store = procedure_store_for_test(&dir);
658
659        let procedure_id = ProcedureId::random();
660        let procedure: BoxedProcedure = Box::new(MockProcedure::new("test store procedure"));
661        let type_name = procedure.type_name().to_string();
662        let data = procedure.dump().unwrap();
663        store
664            .store_procedure(procedure_id, 0, type_name.clone(), data.clone(), None)
665            .await
666            .unwrap();
667        let message = ProcedureMessage {
668            type_name,
669            data,
670            parent_id: None,
671            subprocedure_recovery: SubprocedureRecovery::Discard,
672            step: 1,
673            error: None,
674            context: ProcedureContext::default(),
675        };
676        store
677            .rollback_procedure(procedure_id, message)
678            .await
679            .unwrap();
680
681        let ProcedureMessages {
682            messages,
683            rollback_messages,
684            finished_ids,
685        } = store.load_messages().await.unwrap();
686        assert!(messages.is_empty());
687        assert_eq!(1, rollback_messages.len());
688        assert!(finished_ids.is_empty());
689        assert!(rollback_messages.contains_key(&procedure_id));
690    }
691
692    #[tokio::test]
693    async fn test_delete_procedure() {
694        let dir = create_temp_dir("delete_procedure");
695        let store = procedure_store_for_test(&dir);
696
697        let procedure_id = ProcedureId::random();
698        let procedure: BoxedProcedure = Box::new(MockProcedure::new("test store procedure"));
699        let type_name = procedure.type_name().to_string();
700        let data = procedure.dump().unwrap();
701        store
702            .store_procedure(procedure_id, 0, type_name, data, None)
703            .await
704            .unwrap();
705        let type_name = procedure.type_name().to_string();
706        let data = procedure.dump().unwrap();
707        store
708            .store_procedure(procedure_id, 1, type_name, data, None)
709            .await
710            .unwrap();
711
712        store.delete_procedure(procedure_id).await.unwrap();
713
714        let ProcedureMessages {
715            messages,
716            rollback_messages,
717            finished_ids,
718        } = store.load_messages().await.unwrap();
719        assert!(messages.is_empty());
720        assert!(rollback_messages.is_empty());
721        assert!(finished_ids.is_empty());
722    }
723
724    #[tokio::test]
725    async fn test_delete_committed_procedure() {
726        let dir = create_temp_dir("delete_committed");
727        let store = procedure_store_for_test(&dir);
728
729        let procedure_id = ProcedureId::random();
730        let procedure: BoxedProcedure = Box::new(MockProcedure::new("test store procedure"));
731
732        let type_name = procedure.type_name().to_string();
733        let data = procedure.dump().unwrap();
734        store
735            .store_procedure(procedure_id, 0, type_name, data, None)
736            .await
737            .unwrap();
738
739        let type_name = procedure.type_name().to_string();
740        let data = procedure.dump().unwrap();
741        store
742            .store_procedure(procedure_id, 1, type_name, data, None)
743            .await
744            .unwrap();
745        store.commit_procedure(procedure_id, 2).await.unwrap();
746
747        store.delete_procedure(procedure_id).await.unwrap();
748
749        let ProcedureMessages {
750            messages,
751            rollback_messages,
752            finished_ids,
753        } = store.load_messages().await.unwrap();
754        assert!(messages.is_empty());
755        assert!(rollback_messages.is_empty());
756        assert!(finished_ids.is_empty());
757    }
758
759    #[tokio::test]
760    async fn test_load_messages() {
761        let dir = create_temp_dir("load_messages");
762        let store = procedure_store_for_test(&dir);
763
764        // store 3 steps
765        let id0 = ProcedureId::random();
766        let procedure: BoxedProcedure = Box::new(MockProcedure::new("id0-0"));
767        let type_name = procedure.type_name().to_string();
768        let data = procedure.dump().unwrap();
769        store
770            .store_procedure(id0, 0, type_name, data, None)
771            .await
772            .unwrap();
773        let procedure: BoxedProcedure = Box::new(MockProcedure::new("id0-1"));
774        let type_name = procedure.type_name().to_string();
775        let data = procedure.dump().unwrap();
776        store
777            .store_procedure(id0, 1, type_name, data, None)
778            .await
779            .unwrap();
780        let procedure: BoxedProcedure = Box::new(MockProcedure::new("id0-2"));
781        let type_name = procedure.type_name().to_string();
782        let data = procedure.dump().unwrap();
783        store
784            .store_procedure(id0, 2, type_name, data, None)
785            .await
786            .unwrap();
787
788        // store 2 steps and then commit
789        let id1 = ProcedureId::random();
790        let procedure: BoxedProcedure = Box::new(MockProcedure::new("id1-0"));
791        let type_name = procedure.type_name().to_string();
792        let data = procedure.dump().unwrap();
793        store
794            .store_procedure(id1, 0, type_name, data, None)
795            .await
796            .unwrap();
797        let procedure: BoxedProcedure = Box::new(MockProcedure::new("id1-1"));
798        let type_name = procedure.type_name().to_string();
799        let data = procedure.dump().unwrap();
800        store
801            .store_procedure(id1, 1, type_name, data, None)
802            .await
803            .unwrap();
804        store.commit_procedure(id1, 2).await.unwrap();
805
806        // store 1 step
807        let id2 = ProcedureId::random();
808        let procedure: BoxedProcedure = Box::new(MockProcedure::new("id2-0"));
809        let type_name = procedure.type_name().to_string();
810        let data = procedure.dump().unwrap();
811        store
812            .store_procedure(id2, 0, type_name, data, None)
813            .await
814            .unwrap();
815
816        let ProcedureMessages {
817            messages,
818            rollback_messages,
819            finished_ids,
820        } = store.load_messages().await.unwrap();
821        assert_eq!(2, messages.len());
822        assert!(rollback_messages.is_empty());
823        assert_eq!(1, finished_ids.len());
824
825        let msg = messages.get(&id0).unwrap();
826        assert_eq!("id0-2", msg.data);
827        let msg = messages.get(&id2).unwrap();
828        assert_eq!("id2-0", msg.data);
829    }
830}
common_procedure/store.rs

common_procedure/
store.rs
