common_procedure/local/
recovery.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::sync::Arc;
16
17use common_telemetry::{error, info};
18
19use crate::error;
20use crate::local::{LocalManager, RootSubmissionOrigin};
21use crate::procedure::InitProcedureState;
22use crate::store::ProcedureMessage;
23use crate::{ProcedureId, ProcedureState, ProcedureWithId};
24
25/// A loaded procedure ready to resume from its checkpoint.
26struct RecoveredProcedure {
27    procedure: ProcedureWithId,
28    state: ProcedureState,
29    step: u32,
30}
31
32impl LocalManager {
33    /// Loads all eligible procedures before starting any recovered runner.
34    pub(crate) fn submit_recovered_messages(
35        &self,
36        messages: impl IntoIterator<Item = (ProcedureId, ProcedureMessage, InitProcedureState)>,
37    ) {
38        // Load both rollback and execution checkpoints before spawning. Loaders
39        // rebuild trackers that parents consult before submitting new work.
40        let procedures: Vec<_> = messages
41            .into_iter()
42            .filter_map(|(id, message, state)| self.load_recovered_procedure(id, message, state))
43            .collect();
44
45        for RecoveredProcedure {
46            procedure,
47            state,
48            step,
49        } in procedures
50        {
51            let procedure_id = procedure.id;
52            // Detached children use the existing root lifecycle and cleanup.
53            if let Err(e) = self.submit_root(
54                procedure_id,
55                state,
56                step,
57                procedure.procedure,
58                procedure.context,
59                RootSubmissionOrigin::Recovery,
60            ) {
61                error!(e; "Failed to recover procedure {}", procedure_id);
62            }
63        }
64    }
65
66    fn load_recovered_procedure(
67        &self,
68        procedure_id: ProcedureId,
69        message: ProcedureMessage,
70        init_state: InitProcedureState,
71    ) -> Option<RecoveredProcedure> {
72        if message.parent_id.is_some() && message.subprocedure_recovery.is_discard() {
73            return None;
74        }
75        // The loader logs failures; skip them without adding a duplicate error.
76        let mut loaded = self
77            .manager_ctx
78            .load_one_procedure_from_message(procedure_id, &message)?;
79        match message.parent_id {
80            Some(parent_id) => info!(
81                "Recover detached subprocedure {}-{}, former parent: {}, step: {}",
82                loaded.procedure.type_name(),
83                procedure_id,
84                parent_id,
85                loaded.step
86            ),
87            None => info!(
88                "Recover root procedure {}-{}, step: {}",
89                loaded.procedure.type_name(),
90                procedure_id,
91                loaded.step
92            ),
93        }
94        let state = match init_state {
95            InitProcedureState::Running => ProcedureState::Running,
96            InitProcedureState::RollingBack => ProcedureState::RollingBack {
97                error: Arc::new(
98                    error::RollbackProcedureRecoveredSnafu {
99                        error: message.error.unwrap_or_else(|| "Unknown error".to_string()),
100                    }
101                    .build(),
102                ),
103            },
104        };
105        if let Err(e) = loaded.procedure.recover() {
106            error!(e; "Failed to recover procedure {}", procedure_id);
107        }
108        Some(RecoveredProcedure {
109            procedure: ProcedureWithId {
110                id: procedure_id,
111                procedure: loaded.procedure,
112                context: message.context,
113            },
114            state,
115            step: loaded.step,
116        })
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use std::sync::atomic::{AtomicUsize, Ordering};
123    use std::time::Duration;
124
125    use async_trait::async_trait;
126    use common_test_util::temp_dir::create_temp_dir;
127    use tokio::sync::Notify;
128    use tokio::time::timeout;
129
130    use super::*;
131    use crate::error::Result;
132    use crate::local::{ManagerConfig, test_util};
133    use crate::store::state_store::ObjectStateStore;
134    use crate::test_util::InMemoryPoisonStore;
135    use crate::{
136        Context, LockKey, Procedure, ProcedureContext, ProcedureManager, Status,
137        SubprocedureRecovery,
138    };
139
140    struct MockParentProcedure {
141        child: Option<ProcedureWithId>,
142        recovery: SubprocedureRecovery,
143    }
144
145    #[async_trait]
146    impl Procedure for MockParentProcedure {
147        fn type_name(&self) -> &str {
148            "MockParentProcedure"
149        }
150
151        async fn execute(&mut self, _: &Context) -> Result<Status> {
152            let Some(child) = self.child.take() else {
153                return Ok(Status::done());
154            };
155            Ok(match self.recovery {
156                SubprocedureRecovery::Discard => Status::suspended(vec![child], false),
157                SubprocedureRecovery::Detach => Status::Suspended {
158                    subprocedures: vec![child],
159                    persist: false,
160                    recovery: SubprocedureRecovery::Detach,
161                },
162            })
163        }
164
165        fn dump(&self) -> Result<String> {
166            unreachable!("The fixture deliberately does not checkpoint the parent")
167        }
168
169        fn lock_key(&self) -> LockKey {
170            LockKey::default()
171        }
172    }
173
174    #[derive(Clone)]
175    struct MockChildProcedure {
176        checkpointed: bool,
177        init_state: InitProcedureState,
178        entered: Arc<Notify>,
179        release: Arc<Notify>,
180        recoveries: Arc<AtomicUsize>,
181        rollbacks: Arc<AtomicUsize>,
182    }
183
184    #[async_trait]
185    impl Procedure for MockChildProcedure {
186        fn type_name(&self) -> &str {
187            "MockChildProcedure"
188        }
189
190        async fn execute(&mut self, _: &Context) -> Result<Status> {
191            if matches!(self.init_state, InitProcedureState::RollingBack) {
192                return error::UnexpectedSnafu {
193                    err_msg: "Child needs rollback",
194                }
195                .fail();
196            }
197            if !self.checkpointed {
198                self.checkpointed = true;
199                return Ok(Status::executing(true));
200            }
201            self.entered.notify_one();
202            self.release.notified().await;
203            Ok(Status::done())
204        }
205
206        fn recover(&mut self) -> Result<()> {
207            self.recoveries.fetch_add(1, Ordering::Relaxed);
208            // The next execution checkpoint must persist the detached root identity.
209            self.checkpointed = false;
210            Ok(())
211        }
212
213        fn dump(&self) -> Result<String> {
214            Ok(self.checkpointed.to_string())
215        }
216
217        fn lock_key(&self) -> LockKey {
218            LockKey::default()
219        }
220
221        fn rollback_supported(&self) -> bool {
222            true
223        }
224
225        async fn rollback(&mut self, _: &Context) -> Result<()> {
226            self.rollbacks.fetch_add(1, Ordering::Relaxed);
227            self.entered.notify_one();
228            self.release.notified().await;
229            Ok(())
230        }
231    }
232
233    #[tokio::test]
234    async fn test_detach_child_after_restart() {
235        check_child_recovery(SubprocedureRecovery::Detach, InitProcedureState::Running).await;
236    }
237
238    #[tokio::test]
239    async fn test_detach_rollback_child_after_restart() {
240        check_child_recovery(
241            SubprocedureRecovery::Detach,
242            InitProcedureState::RollingBack,
243        )
244        .await;
245    }
246
247    #[tokio::test]
248    async fn test_discard_child_after_restart() {
249        check_child_recovery(SubprocedureRecovery::Discard, InitProcedureState::Running).await;
250    }
251
252    #[tokio::test]
253    async fn test_discard_rollback_child_after_restart() {
254        check_child_recovery(
255            SubprocedureRecovery::Discard,
256            InitProcedureState::RollingBack,
257        )
258        .await;
259    }
260
261    #[tokio::test]
262    async fn test_load_all_procedures_before_spawning() {
263        // Cover both checkpoint kinds and both cross-kind parent/child orderings.
264        for (parent_rollback, child_rollback) in
265            [(false, false), (false, true), (true, false), (true, true)]
266        {
267            let dir = create_temp_dir("load-before-spawn");
268            let manager = LocalManager::new(
269                ManagerConfig::default(),
270                Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir))),
271                Arc::new(InMemoryPoisonStore::new()),
272                None,
273                None,
274            );
275            let manager_ctx = manager.manager_ctx.clone();
276            let loads = Arc::new(AtomicUsize::new(0));
277            let recoveries = Arc::new(AtomicUsize::new(0));
278            let loaded_count = loads.clone();
279            let recovered_count = recoveries.clone();
280            manager
281                .register_loader(
282                    "MockChildProcedure",
283                    Box::new(move |data| {
284                        // A loader must never observe an already spawned parent.
285                        // This check is deterministic even if its runner has not run yet.
286                        assert!(manager_ctx.runner_tasks.lock().unwrap().is_empty());
287                        loaded_count.fetch_add(1, Ordering::Relaxed);
288                        Ok(Box::new(MockChildProcedure {
289                            checkpointed: data.parse().unwrap(),
290                            init_state: InitProcedureState::Running,
291                            entered: Arc::new(Notify::new()),
292                            release: Arc::new(Notify::new()),
293                            recoveries: recovered_count.clone(),
294                            rollbacks: Arc::new(AtomicUsize::new(0)),
295                        }))
296                    }),
297                )
298                .unwrap();
299            let parent_id = ProcedureId::random();
300            let child_id = ProcedureId::random();
301            for (id, parent, rollback) in [
302                (parent_id, None, parent_rollback),
303                (child_id, Some(parent_id), child_rollback),
304            ] {
305                let message = ProcedureMessage {
306                    type_name: "MockChildProcedure".to_string(),
307                    data: "true".to_string(),
308                    parent_id: parent,
309                    subprocedure_recovery: SubprocedureRecovery::Detach,
310                    step: 7,
311                    error: rollback.then(|| "Needs rollback".to_string()),
312                    context: ProcedureContext::default(),
313                };
314                if rollback {
315                    manager
316                        .procedure_store
317                        .rollback_procedure(id, message)
318                        .await
319                        .unwrap();
320                } else {
321                    manager
322                        .procedure_store
323                        .store_procedure_message(id, &message)
324                        .await
325                        .unwrap();
326                }
327            }
328
329            manager.start().await.unwrap();
330            assert_eq!(loads.load(Ordering::Relaxed), 2);
331            assert_eq!(recoveries.load(Ordering::Relaxed), 2);
332            assert_eq!(manager.manager_ctx.runner_tasks.lock().unwrap().len(), 2);
333            manager.stop().await.unwrap();
334        }
335    }
336
337    #[tokio::test]
338    async fn test_recover_roots_and_detached_children_over_limit() {
339        let dir = create_temp_dir("detached-recovery-limit");
340        let manager = LocalManager::new(
341            ManagerConfig {
342                max_running_procedures: 1,
343                ..Default::default()
344            },
345            Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir))),
346            Arc::new(InMemoryPoisonStore::new()),
347            None,
348            None,
349        );
350        manager
351            .register_loader(
352                "MockChildProcedure",
353                Box::new(|data| {
354                    Ok(Box::new(MockChildProcedure {
355                        checkpointed: data.parse().unwrap(),
356                        init_state: InitProcedureState::Running,
357                        entered: Arc::new(Notify::new()),
358                        release: Arc::new(Notify::new()),
359                        recoveries: Arc::new(AtomicUsize::new(0)),
360                        rollbacks: Arc::new(AtomicUsize::new(0)),
361                    }))
362                }),
363            )
364            .unwrap();
365        let root_id = ProcedureId::random();
366        let child_ids = [ProcedureId::random(), ProcedureId::random()];
367        let root = ProcedureMessage {
368            type_name: "MockChildProcedure".to_string(),
369            data: "true".to_string(),
370            parent_id: None,
371            subprocedure_recovery: SubprocedureRecovery::Discard,
372            step: 7,
373            error: None,
374            context: ProcedureContext::default(),
375        };
376        manager
377            .procedure_store
378            .store_procedure_message(root_id, &root)
379            .await
380            .unwrap();
381        for child_id in child_ids {
382            let child = ProcedureMessage {
383                parent_id: Some(root_id),
384                subprocedure_recovery: SubprocedureRecovery::Detach,
385                ..root.clone()
386            };
387            manager
388                .procedure_store
389                .store_procedure_message(child_id, &child)
390                .await
391                .unwrap();
392        }
393
394        manager.start().await.unwrap();
395        {
396            let procedures = manager.manager_ctx.procedures.read().unwrap();
397            assert_eq!(procedures.len(), 3);
398            for id in [root_id, child_ids[0], child_ids[1]] {
399                let meta = &procedures[&id];
400                assert_eq!(meta.parent_id, None);
401                assert_eq!(meta.num_children(), 0);
402            }
403        }
404        // Wait for detached children to checkpoint as roots, then verify that
405        // recovery retained their nonzero step rather than restarting from zero.
406        timeout(Duration::from_secs(5), async {
407            loop {
408                let stored = manager.procedure_store.load_messages().await.unwrap();
409                if child_ids.iter().all(|id| {
410                    stored
411                        .messages
412                        .get(id)
413                        .is_some_and(|message| message.parent_id.is_none())
414                }) {
415                    for child_id in child_ids {
416                        assert_eq!(stored.messages[&child_id].step, 7);
417                    }
418                    break;
419                }
420                tokio::task::yield_now().await;
421            }
422        })
423        .await
424        .unwrap();
425        manager.stop().await.unwrap();
426    }
427
428    async fn check_child_recovery(policy: SubprocedureRecovery, init_state: InitProcedureState) {
429        let dir = create_temp_dir("child-recovery");
430        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
431        let manager = LocalManager::new(
432            ManagerConfig {
433                max_running_procedures: 1,
434                ..Default::default()
435            },
436            state_store,
437            Arc::new(InMemoryPoisonStore::new()),
438            None,
439            None,
440        );
441        let child_id = ProcedureId::random();
442        let entered = Arc::new(Notify::new());
443        let release = Arc::new(Notify::new());
444        let recoveries = Arc::new(AtomicUsize::new(0));
445        let rollbacks = Arc::new(AtomicUsize::new(0));
446        let child = MockChildProcedure {
447            checkpointed: false,
448            init_state: init_state.clone(),
449            entered: entered.clone(),
450            release: release.clone(),
451            recoveries: recoveries.clone(),
452            rollbacks: rollbacks.clone(),
453        };
454        let loaded_child = child.clone();
455        manager
456            .register_loader(
457                "MockChildProcedure",
458                Box::new(move |data| {
459                    assert_eq!(
460                        policy,
461                        SubprocedureRecovery::Detach,
462                        "Discarded children must not be loaded"
463                    );
464                    let mut child = loaded_child.clone();
465                    child.checkpointed = data.parse().unwrap();
466                    Ok(Box::new(child))
467                }),
468            )
469            .unwrap();
470        manager.start().await.unwrap();
471        let parent = ProcedureWithId::with_random_id(Box::new(MockParentProcedure {
472            child: Some(ProcedureWithId {
473                id: child_id,
474                procedure: Box::new(child),
475                context: ProcedureContext::default(),
476            }),
477            recovery: policy,
478        }))
479        .with_context(ProcedureContext {
480            actor: Some("alice".to_string()),
481            ..Default::default()
482        });
483        let parent_id = parent.id;
484        manager.submit(parent).await.unwrap();
485        timeout(Duration::from_secs(5), entered.notified())
486            .await
487            .unwrap();
488
489        // The first child checkpoint carries the submission policy, even though
490        // Status::Suspended.persist was false and the parent has no checkpoint.
491        let stored = manager.procedure_store.load_messages().await.unwrap();
492        let checkpoint = match init_state {
493            InitProcedureState::Running => &stored.messages[&child_id],
494            InitProcedureState::RollingBack => &stored.rollback_messages[&child_id],
495        }
496        .clone();
497        assert_eq!(checkpoint.parent_id, Some(parent_id));
498        assert_eq!(checkpoint.subprocedure_recovery, policy);
499        assert_eq!(checkpoint.context.actor.as_deref(), Some("alice"));
500        assert_eq!(checkpoint.step, 0);
501        if matches!(init_state, InitProcedureState::RollingBack) {
502            assert!(!stored.messages.contains_key(&child_id));
503            assert_eq!(checkpoint.data, "false");
504        }
505        assert!(!stored.messages.contains_key(&parent_id));
506        manager.stop().await.unwrap();
507        manager.start().await.unwrap();
508
509        if policy == SubprocedureRecovery::Discard {
510            assert!(manager.procedure_state(child_id).await.unwrap().is_none());
511            assert_eq!(recoveries.load(Ordering::Relaxed), 0);
512            let stored = manager.procedure_store.load_messages().await.unwrap();
513            let retained = match init_state {
514                InitProcedureState::Running => &stored.messages[&child_id],
515                InitProcedureState::RollingBack => &stored.rollback_messages[&child_id],
516            };
517            assert_eq!(retained, &checkpoint);
518            manager.stop().await.unwrap();
519            return;
520        }
521
522        // The original parent is absent. The child owns its cleanup as a root,
523        // and remains recoverable across a second restart.
524        for restart in 1..=2 {
525            timeout(Duration::from_secs(5), entered.notified())
526                .await
527                .unwrap();
528            assert!(manager.procedure_state(parent_id).await.unwrap().is_none());
529            let meta = manager.manager_ctx.procedures.read().unwrap()[&child_id].clone();
530            assert_eq!(meta.parent_id, None);
531            assert_eq!(meta.context, checkpoint.context);
532            assert_eq!(recoveries.load(Ordering::Relaxed), restart);
533            let stored = manager.procedure_store.load_messages().await.unwrap();
534            match init_state {
535                InitProcedureState::Running => {
536                    let recovered = &stored.messages[&child_id];
537                    assert_eq!(recovered.parent_id, None);
538                    assert_eq!(recovered.step, checkpoint.step);
539                    assert_eq!(recovered.context, checkpoint.context);
540                }
541                InitProcedureState::RollingBack => {
542                    assert!(meta.state().is_rolling_back());
543                    assert_eq!(stored.rollback_messages[&child_id], checkpoint);
544                }
545            }
546            if restart == 1 {
547                manager.stop().await.unwrap();
548                manager.start().await.unwrap();
549            }
550        }
551        release.notify_one();
552        timeout(Duration::from_secs(5), async {
553            while manager
554                .manager_ctx
555                .runner_tasks
556                .lock()
557                .unwrap()
558                .contains_key(&child_id)
559            {
560                tokio::task::yield_now().await;
561            }
562        })
563        .await
564        .unwrap();
565        let stored = manager.procedure_store.load_messages().await.unwrap();
566        assert!(stored.messages.is_empty());
567        assert!(stored.rollback_messages.is_empty());
568        assert!(stored.finished_ids.is_empty());
569        let state = manager.procedure_state(child_id).await.unwrap().unwrap();
570        match init_state {
571            InitProcedureState::Running => {
572                assert!(state.is_done());
573                assert_eq!(rollbacks.load(Ordering::Relaxed), 0);
574            }
575            InitProcedureState::RollingBack => {
576                assert!(state.is_failed());
577                assert_eq!(rollbacks.load(Ordering::Relaxed), 3);
578            }
579        }
580        manager.stop().await.unwrap();
581    }
582}
common_procedure/local/recovery.rs

common_procedure/local/
recovery.rs
