common_procedure/

local.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod runner;

use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};

use async_trait::async_trait;
use backon::ExponentialBuilder;
use common_error::ext::BoxedError;
use common_event_recorder::EventRecorderRef;
use common_runtime::{JoinHandle, RepeatedTask, TaskFunction};
use common_telemetry::tracing_context::{FutureExt, TracingContext};
use common_telemetry::{error, info, tracing};
use snafu::{OptionExt, ResultExt, ensure};
use tokio::sync::watch::{self, Receiver, Sender};
use tokio::sync::{Mutex as TokioMutex, Notify};

use crate::error::{
    self, CheckStatusSnafu, DuplicateProcedureSnafu, Error, LoaderConflictSnafu,
    ManagerNotStartSnafu, ManagerPasuedSnafu, PoisonKeyNotDefinedSnafu, ProcedureNotFoundSnafu,
    Result, StartRemoveOutdatedMetaTaskSnafu, StopRemoveOutdatedMetaTaskSnafu,
    TooManyRunningProceduresSnafu,
};
use crate::local::runner::Runner;
use crate::procedure::{BoxedProcedureLoader, InitProcedureState, PoisonKeys, ProcedureInfo};
use crate::rwlock::{KeyRwLock, OwnedKeyRwLockGuard};
use crate::store::poison_store::PoisonStoreRef;
use crate::store::{ProcedureMessage, ProcedureMessages, ProcedureStore, StateStoreRef};
use crate::{
    BoxedProcedure, ContextProvider, EventTrigger, LockKey, PoisonKey, ProcedureContext,
    ProcedureId, ProcedureManager, ProcedureState, ProcedureWithId, StringKey, Watcher,
};

/// The expired time of a procedure's metadata.
const META_TTL: Duration = Duration::from_secs(60 * 10);

#[derive(Clone, Copy)]
enum RootSubmissionOrigin {
    Fresh,
    Recovery,
}

impl RootSubmissionOrigin {
    fn event_trigger(self) -> EventTrigger {
        match self {
            Self::Fresh => EventTrigger::Submitted,
            Self::Recovery => EventTrigger::Recovered,
        }
    }
}

/// Shared metadata of a procedure.
///
/// # Note
/// [Notify] is not a condition variable, we can't guarantee the waiters are notified
/// if they didn't call `notified()` before we signal the notify. So we
/// 1. use dedicated notify for each condition, such as waiting for a lock, waiting
///    for children;
/// 2. always use `notify_one` and ensure there are only one waiter.
#[derive(Debug)]
pub(crate) struct ProcedureMeta {
    /// Id of this procedure.
    id: ProcedureId,
    /// Type name of this procedure.
    type_name: String,
    /// Parent procedure id.
    parent_id: Option<ProcedureId>,
    /// Context associated with the root submission.
    context: ProcedureContext,
    /// Notify to wait for subprocedures.
    child_notify: Notify,
    /// Lock required by this procedure.
    lock_key: LockKey,
    /// Poison keys that may cause this procedure to become poisoned during execution.
    poison_keys: PoisonKeys,
    /// Sender to notify the procedure state.
    state_sender: Sender<ProcedureState>,
    /// Receiver to watch the procedure state.
    state_receiver: Receiver<ProcedureState>,
    /// Id of child procedures.
    children: Mutex<Vec<ProcedureId>>,
    /// Start execution time of this procedure.
    start_time_ms: AtomicI64,
    /// End execution time of this procedure.
    end_time_ms: AtomicI64,
}

impl ProcedureMeta {
    #[allow(clippy::too_many_arguments)]
    fn new(
        id: ProcedureId,
        procedure_state: ProcedureState,
        parent_id: Option<ProcedureId>,
        context: ProcedureContext,
        lock_key: LockKey,
        poison_keys: PoisonKeys,
        type_name: &str,
    ) -> ProcedureMeta {
        let (state_sender, state_receiver) = watch::channel(procedure_state);
        ProcedureMeta {
            id,
            parent_id,
            context,
            child_notify: Notify::new(),
            lock_key,
            poison_keys,
            state_sender,
            state_receiver,
            children: Mutex::new(Vec::new()),
            start_time_ms: AtomicI64::new(0),
            end_time_ms: AtomicI64::new(0),
            type_name: type_name.to_string(),
        }
    }

    /// Returns current [ProcedureState].
    fn state(&self) -> ProcedureState {
        self.state_receiver.borrow().clone()
    }

    /// Update current [ProcedureState].
    fn set_state(&self, state: ProcedureState) {
        // Safety: ProcedureMeta also holds the receiver, so `send()` should never fail.
        self.state_sender.send(state).unwrap();
    }

    /// Push `procedure_id` of the subprocedure to the metadata.
    fn push_child(&self, procedure_id: ProcedureId) {
        let mut children = self.children.lock().unwrap();
        children.push(procedure_id);
    }

    /// Append subprocedures to given `buffer`.
    fn list_children(&self, buffer: &mut Vec<ProcedureId>) {
        let children = self.children.lock().unwrap();
        buffer.extend_from_slice(&children);
    }

    /// Returns the number of subprocedures.
    fn num_children(&self) -> usize {
        self.children.lock().unwrap().len()
    }

    /// update the start time of the procedure.
    fn set_start_time_ms(&self) {
        self.start_time_ms
            .store(common_time::util::current_time_millis(), Ordering::Relaxed);
    }

    /// update the end time of the procedure.
    fn set_end_time_ms(&self) {
        self.end_time_ms
            .store(common_time::util::current_time_millis(), Ordering::Relaxed);
    }
}

/// Reference counted pointer to [ProcedureMeta].
type ProcedureMetaRef = Arc<ProcedureMeta>;

/// Procedure loaded from store.
struct LoadedProcedure {
    procedure: BoxedProcedure,
    step: u32,
}

/// The dynamic lock for procedure execution.
///
/// Unlike the procedure-level locks, these locks are acquired dynamically by the procedure
/// during execution. They are only held when the procedure specifically needs these keys
/// and are released as soon as the procedure no longer needs them.
/// This allows for more fine-grained concurrency control during procedure execution.
pub(crate) type DynamicKeyLock = Arc<KeyRwLock<String>>;

/// Acquires a dynamic key lock for the given key.
///
/// This function takes a reference to the dynamic key lock and a pointer to the key.
/// It then matches the key type and acquires the appropriate lock.
pub async fn acquire_dynamic_key_lock(
    lock: &DynamicKeyLock,
    key: &StringKey,
) -> DynamicKeyLockGuard {
    match key {
        StringKey::Share(key) => {
            let guard = lock.read(key.clone()).await;
            DynamicKeyLockGuard {
                guard: Some(OwnedKeyRwLockGuard::from(guard)),
                key: key.clone(),
                lock: lock.clone(),
            }
        }
        StringKey::Exclusive(key) => {
            let guard = lock.write(key.clone()).await;
            DynamicKeyLockGuard {
                guard: Some(OwnedKeyRwLockGuard::from(guard)),
                key: key.clone(),
                lock: lock.clone(),
            }
        }
    }
}
/// A guard for the dynamic key lock.
///
/// This guard is used to release the lock when the procedure no longer needs it.
/// It also ensures that the lock is cleaned up when the guard is dropped.
pub struct DynamicKeyLockGuard {
    guard: Option<OwnedKeyRwLockGuard>,
    key: String,
    lock: DynamicKeyLock,
}

impl Drop for DynamicKeyLockGuard {
    fn drop(&mut self) {
        if let Some(guard) = self.guard.take() {
            drop(guard);
        }
        self.lock.clean_keys(std::slice::from_ref(&self.key));
    }
}

/// Shared context of the manager.
pub(crate) struct ManagerContext {
    /// Procedure loaders. The key is the type name of the procedure which the loader returns.
    loaders: Mutex<HashMap<String, BoxedProcedureLoader>>,
    /// The key lock for the procedure.
    ///
    /// The lock keys are defined in `Procedure::lock_key()`.
    /// These locks are acquired before the procedure starts and released after the procedure finishes.
    /// They ensure exclusive access to resources throughout the entire procedure lifecycle.
    key_lock: KeyRwLock<String>,
    /// The dynamic lock for procedure execution.
    ///
    /// Unlike the procedure-level locks, these locks are acquired dynamically by the procedure
    /// during execution. They are only held when the procedure specifically needs these keys
    /// and are released as soon as the procedure no longer needs them.
    /// This allows for more fine-grained concurrency control during procedure execution.
    dynamic_key_lock: DynamicKeyLock,
    /// Procedures in the manager.
    procedures: RwLock<HashMap<ProcedureId, ProcedureMetaRef>>,
    /// Running procedures.
    running_procedures: Mutex<HashSet<ProcedureId>>,
    /// Ids and finished time of finished procedures.
    finished_procedures: Mutex<VecDeque<(ProcedureId, Instant)>>,
    /// Runner tasks of procedures.
    runner_tasks: Mutex<HashMap<ProcedureId, JoinHandle<()>>>,
    /// Running flag.
    running: Arc<AtomicBool>,
    /// Poison manager.
    poison_manager: PoisonStoreRef,
}

#[async_trait]
impl ContextProvider for ManagerContext {
    async fn procedure_state(&self, procedure_id: ProcedureId) -> Result<Option<ProcedureState>> {
        Ok(self.state(procedure_id))
    }

    async fn procedure_state_receiver(
        &self,
        procedure_id: ProcedureId,
    ) -> Result<Option<Receiver<ProcedureState>>> {
        Ok(self.state_receiver(procedure_id))
    }

    async fn try_put_poison(&self, key: &PoisonKey, procedure_id: ProcedureId) -> Result<()> {
        {
            // validate the procedure exists
            let procedures = self.procedures.read().unwrap();
            let procedure = procedures
                .get(&procedure_id)
                .context(ProcedureNotFoundSnafu { procedure_id })?;

            // validate the poison key is defined
            ensure!(
                procedure.poison_keys.contains(key),
                PoisonKeyNotDefinedSnafu {
                    key: key.clone(),
                    procedure_id
                }
            );
        }
        let key = key.to_string();
        let procedure_id = procedure_id.to_string();
        self.poison_manager.try_put_poison(key, procedure_id).await
    }

    async fn acquire_lock(&self, key: &StringKey) -> DynamicKeyLockGuard {
        acquire_dynamic_key_lock(&self.dynamic_key_lock, key).await
    }
}

impl ManagerContext {
    /// Returns a new [ManagerContext].
    fn new(poison_manager: PoisonStoreRef) -> ManagerContext {
        ManagerContext {
            key_lock: KeyRwLock::new(),
            dynamic_key_lock: Arc::new(KeyRwLock::new()),
            loaders: Mutex::new(HashMap::new()),
            procedures: RwLock::new(HashMap::new()),
            running_procedures: Mutex::new(HashSet::new()),
            finished_procedures: Mutex::new(VecDeque::new()),
            runner_tasks: Mutex::new(HashMap::new()),
            running: Arc::new(AtomicBool::new(false)),
            poison_manager,
        }
    }

    #[cfg(test)]
    pub(crate) fn set_running(&self) {
        self.running.store(true, Ordering::Relaxed);
    }

    /// Set the running flag.
    pub(crate) fn start(&self) {
        self.running.store(true, Ordering::Relaxed);
    }

    pub(crate) fn stop(&self) {
        self.running.store(false, Ordering::Relaxed);
    }

    fn reset_runtime_state(&self) {
        self.procedures.write().unwrap().clear();
        self.running_procedures.lock().unwrap().clear();
        self.finished_procedures.lock().unwrap().clear();
        for handle in self
            .runner_tasks
            .lock()
            .unwrap()
            .drain()
            .map(|(_, handle)| handle)
        {
            handle.abort();
        }
        self.key_lock.clear();
        self.dynamic_key_lock.clear();
    }

    fn spawn_runner_task<F>(&self, procedure_id: ProcedureId, spawn: F) -> bool
    where
        F: FnOnce() -> JoinHandle<()>,
    {
        let mut tasks = self.runner_tasks.lock().unwrap();
        if !self.running() {
            return false;
        }

        let handle = spawn();
        let _ = tasks.insert(procedure_id, handle);
        true
    }

    fn remove_procedure(&self, procedure_id: ProcedureId) {
        self.procedures.write().unwrap().remove(&procedure_id);
        self.running_procedures
            .lock()
            .unwrap()
            .remove(&procedure_id);
    }

    pub(crate) fn remove_runner_task(&self, procedure_id: ProcedureId) {
        let _ = self.runner_tasks.lock().unwrap().remove(&procedure_id);
    }

    fn take_runner_tasks(&self) -> Vec<JoinHandle<()>> {
        self.runner_tasks
            .lock()
            .unwrap()
            .drain()
            .map(|(_, handle)| handle)
            .collect()
    }

    async fn abort_runner_tasks(&self) {
        let handles = self.take_runner_tasks();

        for handle in &handles {
            handle.abort();
        }

        for handle in handles {
            if let Err(e) = handle.await
                && !e.is_cancelled()
            {
                error!(
                    e; "Procedure runner task exits unexpectedly during stop",
                );
            }
        }
    }

    /// Return `ProcedureManager` is running.
    pub(crate) fn running(&self) -> bool {
        self.running.load(Ordering::Relaxed)
    }

    /// Returns true if the procedure with specific `procedure_id` exists.
    fn contains_procedure(&self, procedure_id: ProcedureId) -> bool {
        let procedures = self.procedures.read().unwrap();
        procedures.contains_key(&procedure_id)
    }

    /// Returns the number of running procedures.
    fn num_running_procedures(&self) -> usize {
        self.running_procedures.lock().unwrap().len()
    }

    /// Try to insert the `procedure` to the context if there is no procedure
    /// with same [ProcedureId].
    ///
    /// Returns `false` if there is already a procedure using the same [ProcedureId].
    fn try_insert_procedure(&self, meta: ProcedureMetaRef) -> bool {
        let procedure_id = meta.id;
        let mut procedures = self.procedures.write().unwrap();
        match procedures.entry(procedure_id) {
            Entry::Occupied(_) => return false,
            Entry::Vacant(vacant_entry) => {
                vacant_entry.insert(meta);
            }
        }

        let mut running_procedures = self.running_procedures.lock().unwrap();
        running_procedures.insert(procedure_id);

        true
    }

    /// Returns the [ProcedureState] of specific `procedure_id`.
    fn state(&self, procedure_id: ProcedureId) -> Option<ProcedureState> {
        let procedures = self.procedures.read().unwrap();
        procedures.get(&procedure_id).map(|meta| meta.state())
    }

    /// Returns the [Receiver<ProcedureState>] of specific `procedure_id`.
    fn state_receiver(&self, procedure_id: ProcedureId) -> Option<Receiver<ProcedureState>> {
        let procedures = self.procedures.read().unwrap();
        procedures
            .get(&procedure_id)
            .map(|meta| meta.state_receiver.clone())
    }

    /// Returns the [ProcedureMeta] of all procedures.
    fn list_procedure(&self) -> Vec<ProcedureInfo> {
        let procedures = self.procedures.read().unwrap();
        procedures
            .values()
            .map(|meta| ProcedureInfo {
                id: meta.id,
                type_name: meta.type_name.clone(),
                start_time_ms: meta.start_time_ms.load(Ordering::Relaxed),
                end_time_ms: meta.end_time_ms.load(Ordering::Relaxed),
                state: meta.state(),
                lock_keys: meta.lock_key.get_keys(),
            })
            .collect()
    }

    /// Returns the [Watcher] of specific `procedure_id`.
    fn watcher(&self, procedure_id: ProcedureId) -> Option<Watcher> {
        let procedures = self.procedures.read().unwrap();
        procedures
            .get(&procedure_id)
            .map(|meta| meta.state_receiver.clone())
    }

    /// Notify a suspended parent procedure with specific `procedure_id` by its subprocedure.
    fn notify_by_subprocedure(&self, procedure_id: ProcedureId) {
        let procedures = self.procedures.read().unwrap();
        if let Some(meta) = procedures.get(&procedure_id) {
            meta.child_notify.notify_one();
        }
    }

    /// Load procedure from specific [ProcedureMessage].
    fn load_one_procedure_from_message(
        &self,
        procedure_id: ProcedureId,
        message: &ProcedureMessage,
    ) -> Option<LoadedProcedure> {
        let loaders = self.loaders.lock().unwrap();
        let loader = loaders.get(&message.type_name).or_else(|| {
            error!(
                "Loader not found, procedure_id: {}, type_name: {}",
                procedure_id, message.type_name
            );
            None
        })?;

        let procedure = loader(&message.data)
            .map_err(|e| {
                error!(
                    "Failed to load procedure data, key: {}, source: {:?}",
                    procedure_id, e
                );
                e
            })
            .ok()?;

        Some(LoadedProcedure {
            procedure,
            step: message.step,
        })
    }

    /// Returns all procedures in the tree (including given `root` procedure).
    ///
    /// If callers need a consistent view of the tree, they must ensure no new
    /// procedure is added to the tree during using this method.
    fn procedures_in_tree(&self, root: &ProcedureMetaRef) -> Vec<ProcedureId> {
        let sub_num = root.num_children();
        // Reserve capacity for the root procedure and its children.
        let mut procedures = Vec::with_capacity(1 + sub_num);

        let mut queue = VecDeque::with_capacity(1 + sub_num);
        // Push the root procedure to the queue.
        queue.push_back(root.clone());

        let mut children_ids = Vec::with_capacity(sub_num);
        let mut children = Vec::with_capacity(sub_num);
        while let Some(meta) = queue.pop_front() {
            procedures.push(meta.id);

            // Find metadatas of children.
            children_ids.clear();
            meta.list_children(&mut children_ids);
            self.find_procedures(&children_ids, &mut children);

            // Traverse children later.
            for child in children.drain(..) {
                queue.push_back(child);
            }
        }

        procedures
    }

    /// Finds procedures by given `procedure_ids`.
    ///
    /// Ignores the id if corresponding procedure is not found.
    fn find_procedures(&self, procedure_ids: &[ProcedureId], metas: &mut Vec<ProcedureMetaRef>) {
        let procedures = self.procedures.read().unwrap();
        for procedure_id in procedure_ids {
            if let Some(meta) = procedures.get(procedure_id) {
                metas.push(meta.clone());
            }
        }
    }

    /// Clean resources of finished procedures.
    fn on_procedures_finish(&self, procedure_ids: &[ProcedureId]) {
        // Since users need to query the procedure state, so we can't remove the
        // meta of the procedure directly.
        let now = Instant::now();
        let mut finished_procedures = self.finished_procedures.lock().unwrap();
        finished_procedures.extend(procedure_ids.iter().map(|id| (*id, now)));

        // Remove the procedures from the running set.
        let mut running_procedures = self.running_procedures.lock().unwrap();
        for procedure_id in procedure_ids {
            running_procedures.remove(procedure_id);
        }
    }

    /// Remove metadata of outdated procedures.
    fn remove_outdated_meta(&self, ttl: Duration) {
        let ids = {
            let mut finished_procedures = self.finished_procedures.lock().unwrap();
            if finished_procedures.is_empty() {
                return;
            }

            let mut ids_to_remove = Vec::new();
            while let Some((id, finish_time)) = finished_procedures.front() {
                if finish_time.elapsed() > ttl {
                    ids_to_remove.push(*id);
                    let _ = finished_procedures.pop_front();
                } else {
                    // The rest procedures are finished later, so we can break
                    // the loop.
                    break;
                }
            }
            ids_to_remove
        };

        if ids.is_empty() {
            return;
        }

        let mut procedures = self.procedures.write().unwrap();
        for id in ids {
            let _ = procedures.remove(&id);
        }
    }
}

/// Config for [LocalManager].
#[derive(Debug)]
pub struct ManagerConfig {
    pub parent_path: String,
    pub max_retry_times: usize,
    pub retry_delay: Duration,
    pub remove_outdated_meta_task_interval: Duration,
    pub remove_outdated_meta_ttl: Duration,
    pub max_running_procedures: usize,
}

impl Default for ManagerConfig {
    fn default() -> Self {
        Self {
            parent_path: String::default(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            remove_outdated_meta_task_interval: Duration::from_secs(60 * 10),
            remove_outdated_meta_ttl: META_TTL,
            max_running_procedures: 128,
        }
    }
}

type PauseAwareRef = Arc<dyn PauseAware>;

struct EventRecorderConfig {
    recorder: Option<EventRecorderRef>,
}

/// A delayed configuration handle for procedure lifecycle event recording.
#[derive(Clone)]
pub struct EventRecorderHandle(Arc<Mutex<EventRecorderConfig>>);

impl EventRecorderHandle {
    fn new(recorder: Option<EventRecorderRef>) -> Self {
        Self(Arc::new(Mutex::new(EventRecorderConfig { recorder })))
    }

    /// Installs the recorder used by subsequently submitted procedures.
    pub fn install(&self, recorder: EventRecorderRef) {
        let mut config = self.0.lock().unwrap();
        config.recorder = Some(recorder);
    }
}

#[async_trait]
pub trait PauseAware: Send + Sync {
    /// Returns true if the procedure manager is paused.
    async fn is_paused(&self) -> std::result::Result<bool, BoxedError>;
}

/// A [ProcedureManager] that maintains procedure states locally.
pub struct LocalManager {
    manager_ctx: Arc<ManagerContext>,
    procedure_store: Arc<ProcedureStore>,
    max_retry_times: usize,
    retry_delay: Duration,
    /// GC task.
    remove_outdated_meta_task: TokioMutex<Option<RepeatedTask<Error>>>,
    config: ManagerConfig,
    pause_aware: Option<PauseAwareRef>,
    event_recorder: EventRecorderHandle,
}

impl LocalManager {
    /// Create a new [LocalManager] with specific `config`.
    pub fn new(
        config: ManagerConfig,
        state_store: StateStoreRef,
        poison_store: PoisonStoreRef,
        pause_aware: Option<PauseAwareRef>,
        event_recorder: Option<EventRecorderRef>,
    ) -> LocalManager {
        let manager_ctx = Arc::new(ManagerContext::new(poison_store));

        LocalManager {
            manager_ctx,
            procedure_store: Arc::new(ProcedureStore::new(&config.parent_path, state_store)),
            max_retry_times: config.max_retry_times,
            retry_delay: config.retry_delay,
            remove_outdated_meta_task: TokioMutex::new(None),
            config,
            pause_aware,
            event_recorder: EventRecorderHandle::new(event_recorder),
        }
    }

    /// Returns the handle used to configure procedure lifecycle event recording.
    pub fn event_recorder_handle(&self) -> EventRecorderHandle {
        self.event_recorder.clone()
    }

    /// Build remove outedated meta task
    pub fn build_remove_outdated_meta_task(&self) -> RepeatedTask<Error> {
        RepeatedTask::new(
            self.config.remove_outdated_meta_task_interval,
            Box::new(RemoveOutdatedMetaFunction {
                manager_ctx: self.manager_ctx.clone(),
                ttl: self.config.remove_outdated_meta_ttl,
            }),
        )
    }

    /// Submit a root procedure with given `procedure_id`.
    fn submit_root(
        &self,
        procedure_id: ProcedureId,
        procedure_state: ProcedureState,
        step: u32,
        procedure: BoxedProcedure,
        context: ProcedureContext,
        origin: RootSubmissionOrigin,
    ) -> Result<Watcher> {
        ensure!(self.manager_ctx.running(), ManagerNotStartSnafu);

        let meta = Arc::new(ProcedureMeta::new(
            procedure_id,
            procedure_state,
            None,
            context,
            procedure.lock_key(),
            procedure.poison_keys(),
            procedure.type_name(),
        ));
        let event_recorder = self.event_recorder.0.lock().unwrap();
        let runner = Runner {
            meta: meta.clone(),
            procedure,
            manager_ctx: self.manager_ctx.clone(),
            step,
            exponential_builder: ExponentialBuilder::default()
                .with_min_delay(self.retry_delay)
                .with_max_times(self.max_retry_times),
            store: self.procedure_store.clone(),
            rolling_back: false,
            event_recorder: event_recorder.recorder.clone(),
            execute_retry_attempt: 0,
            rollback_retry_attempt: 0,
        };

        let watcher = meta.state_receiver.clone();

        ensure!(
            self.manager_ctx.num_running_procedures() < self.config.max_running_procedures,
            TooManyRunningProceduresSnafu {
                max_running_procedures: self.config.max_running_procedures,
            }
        );

        // Inserts meta into the manager before actually spawnd the runner.
        ensure!(
            self.manager_ctx.try_insert_procedure(meta),
            DuplicateProcedureSnafu { procedure_id },
        );

        let initial_event = runner.build_event(origin.event_trigger());

        let tracing_context = TracingContext::from_current_span();

        ensure!(
            self.manager_ctx.spawn_runner_task(procedure_id, || {
                common_runtime::spawn_global(async move {
                    let span = tracing_context.attach(tracing::info_span!(
                    "LocalManager::submit_root_procedure",
                        procedure_name = %runner.meta.type_name,
                        procedure_id = %runner.meta.id,
                    ));
                    // Run the root procedure.
                    // The task was moved to another runtime for execution.
                    // In order not to interrupt tracing, a span needs to be created to continue tracing the current task.
                    runner.run().trace(span).await;
                })
            }),
            {
                self.manager_ctx.remove_procedure(procedure_id);
                ManagerNotStartSnafu
            }
        );

        if let Some(event) = initial_event
            && let Some(recorder) = event_recorder.recorder.as_ref()
        {
            recorder.record(Box::new(event));
        }

        Ok(watcher)
    }

    fn submit_recovered_messages(
        &self,
        messages: HashMap<ProcedureId, ProcedureMessage>,
        init_state: InitProcedureState,
    ) {
        for (procedure_id, message) in &messages {
            if message.parent_id.is_none() {
                // This is the root procedure. We only submit the root procedure as it will
                // submit sub-procedures to the manager.
                let Some(mut loaded_procedure) = self
                    .manager_ctx
                    .load_one_procedure_from_message(*procedure_id, message)
                else {
                    // Try to load other procedures.
                    continue;
                };

                info!(
                    "Recover root procedure {}-{}, step: {}",
                    loaded_procedure.procedure.type_name(),
                    procedure_id,
                    loaded_procedure.step
                );

                let procedure_state = match init_state {
                    InitProcedureState::RollingBack => ProcedureState::RollingBack {
                        error: Arc::new(
                            error::RollbackProcedureRecoveredSnafu {
                                error: message.error.clone().unwrap_or("Unknown error".to_string()),
                            }
                            .build(),
                        ),
                    },
                    InitProcedureState::Running => ProcedureState::Running,
                };

                if let Err(e) = loaded_procedure.procedure.recover() {
                    error!(e; "Failed to recover procedure {}", procedure_id);
                }

                if let Err(e) = self.submit_root(
                    *procedure_id,
                    procedure_state,
                    loaded_procedure.step,
                    loaded_procedure.procedure,
                    message.context.clone(),
                    RootSubmissionOrigin::Recovery,
                ) {
                    error!(e; "Failed to recover procedure {}", procedure_id);
                }
            }
        }
    }

    /// Recovers unfinished procedures and reruns them.
    async fn recover(&self) -> Result<()> {
        info!("LocalManager start to recover");
        let recover_start = Instant::now();

        let ProcedureMessages {
            messages,
            rollback_messages,
            finished_ids,
        } = self.procedure_store.load_messages().await?;
        // Submits recovered messages first.
        self.submit_recovered_messages(rollback_messages, InitProcedureState::RollingBack);
        self.submit_recovered_messages(messages, InitProcedureState::Running);

        if !finished_ids.is_empty() {
            info!(
                "LocalManager try to clean finished procedures, num: {}",
                finished_ids.len()
            );

            for procedure_id in finished_ids {
                if let Err(e) = self.procedure_store.delete_procedure(procedure_id).await {
                    error!(e; "Failed to delete procedure {}", procedure_id);
                }
            }
        }

        info!(
            "LocalManager finish recovery, cost: {}ms",
            recover_start.elapsed().as_millis()
        );

        Ok(())
    }

    #[cfg(any(test, feature = "testing"))]
    /// Returns true if contains a specified loader.
    pub fn contains_loader(&self, name: &str) -> bool {
        let loaders = self.manager_ctx.loaders.lock().unwrap();
        loaders.contains_key(name)
    }

    async fn check_status(&self) -> Result<()> {
        if let Some(pause_aware) = self.pause_aware.as_ref() {
            ensure!(
                !pause_aware.is_paused().await.context(CheckStatusSnafu)?,
                ManagerPasuedSnafu
            );
        }

        Ok(())
    }
}

#[async_trait]
impl ProcedureManager for LocalManager {
    fn register_loader(&self, name: &str, loader: BoxedProcedureLoader) -> Result<()> {
        let mut loaders = self.manager_ctx.loaders.lock().unwrap();
        ensure!(!loaders.contains_key(name), LoaderConflictSnafu { name });

        let _ = loaders.insert(name.to_string(), loader);

        Ok(())
    }

    async fn start(&self) -> Result<()> {
        let mut task = self.remove_outdated_meta_task.lock().await;

        if task.is_some() {
            return Ok(());
        }

        let task_inner = self.build_remove_outdated_meta_task();

        task_inner
            .start(common_runtime::global_runtime())
            .context(StartRemoveOutdatedMetaTaskSnafu)?;

        *task = Some(task_inner);

        self.manager_ctx.reset_runtime_state();
        self.manager_ctx.start();

        info!("LocalManager is start.");

        self.recover().await
    }

    async fn stop(&self) -> Result<()> {
        self.manager_ctx.stop();

        let mut task = self.remove_outdated_meta_task.lock().await;
        if let Some(task) = task.take()
            && let Err(e) = task.stop().await.context(StopRemoveOutdatedMetaTaskSnafu)
        {
            error!(e; "Failed to stop remove outdated meta task");
        };

        self.manager_ctx.abort_runner_tasks().await;
        self.manager_ctx.reset_runtime_state();

        info!("LocalManager is stopped.");

        Ok(())
    }

    async fn submit(&self, procedure: ProcedureWithId) -> Result<Watcher> {
        let procedure_id = procedure.id;
        ensure!(
            !self.manager_ctx.contains_procedure(procedure_id),
            DuplicateProcedureSnafu { procedure_id }
        );
        self.check_status().await?;

        self.submit_root(
            procedure.id,
            ProcedureState::Running,
            0,
            procedure.procedure,
            procedure.context,
            RootSubmissionOrigin::Fresh,
        )
    }

    async fn procedure_state(&self, procedure_id: ProcedureId) -> Result<Option<ProcedureState>> {
        Ok(self.manager_ctx.state(procedure_id))
    }

    fn procedure_watcher(&self, procedure_id: ProcedureId) -> Option<Watcher> {
        self.manager_ctx.watcher(procedure_id)
    }

    async fn list_procedures(&self) -> Result<Vec<ProcedureInfo>> {
        Ok(self.manager_ctx.list_procedure())
    }

    async fn has_unfinished_procedure(&self, type_names: &[&str]) -> Result<bool> {
        let messages = self.procedure_store.load_messages().await?;
        Ok(messages
            .messages
            .values()
            .chain(messages.rollback_messages.values())
            .any(|message| type_names.contains(&message.type_name.as_str())))
    }
}

struct RemoveOutdatedMetaFunction {
    manager_ctx: Arc<ManagerContext>,
    ttl: Duration,
}

#[async_trait::async_trait]
impl TaskFunction<Error> for RemoveOutdatedMetaFunction {
    fn name(&self) -> &str {
        "ProcedureManager-remove-outdated-meta-task"
    }

    async fn call(&mut self) -> Result<()> {
        self.manager_ctx.remove_outdated_meta(self.ttl);
        Ok(())
    }
}

/// Create a new [ProcedureMeta] for test purpose.
#[cfg(test)]
pub(crate) mod test_util {
    use common_test_util::temp_dir::TempDir;
    use object_store::ObjectStore;
    use object_store::services::Fs as Builder;

    use super::*;

    pub(crate) fn procedure_meta_for_test() -> ProcedureMeta {
        ProcedureMeta::new(
            ProcedureId::random(),
            ProcedureState::Running,
            None,
            ProcedureContext::default(),
            LockKey::default(),
            PoisonKeys::default(),
            "ProcedureAdapter",
        )
    }

    pub(crate) fn new_object_store(dir: &TempDir) -> ObjectStore {
        let store_dir = dir.path().to_str().unwrap();
        let builder = Builder::default();
        ObjectStore::new(builder.root(store_dir)).unwrap()
    }
}

#[cfg(test)]
mod tests {
    use std::assert_matches;
    use std::collections::HashSet;
    use std::sync::Mutex;
    use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering};

    use common_error::mock::MockError;
    use common_error::status_code::StatusCode;
    use common_event_recorder::{
        Event, EventRecorder, EventTypeFilter, EventTypeFilterRef, PersistentEventContext,
        TriggerReason,
    };
    use common_test_util::temp_dir::create_temp_dir;
    use tokio::sync::oneshot;
    use tokio::time::{sleep, timeout};

    use super::*;
    use crate::error::{self, Error};
    use crate::store::state_store::{ObjectStateStore, StateStore};
    use crate::test_util::InMemoryPoisonStore;
    use crate::{
        ChildSubmissionOutcome, Context, EventContext, EventTrigger, Procedure, ProcedureEvent,
        Status,
    };

    fn new_test_manager_context() -> ManagerContext {
        let poison_manager = Arc::new(InMemoryPoisonStore::default());
        ManagerContext::new(poison_manager)
    }

    #[derive(Debug)]
    struct CapturingEventRecorder {
        events: Mutex<Vec<Box<dyn Event>>>,
        event_type_filter: EventTypeFilterRef,
    }

    impl CapturingEventRecorder {
        fn with_event_type_filter(event_type_filter: EventTypeFilterRef) -> Self {
            Self {
                events: Mutex::new(vec![]),
                event_type_filter,
            }
        }

        fn triggers(&self) -> Vec<EventTrigger> {
            self.events
                .lock()
                .unwrap()
                .iter()
                .map(|event| {
                    event
                        .as_any()
                        .downcast_ref::<ProcedureEvent>()
                        .unwrap()
                        .trigger
                        .clone()
                })
                .collect()
        }

        fn procedure_events(&self) -> Vec<(ProcedureId, EventTrigger)> {
            self.events
                .lock()
                .unwrap()
                .iter()
                .map(|event| {
                    let event = event.as_any().downcast_ref::<ProcedureEvent>().unwrap();
                    (event.procedure_id, event.trigger.clone())
                })
                .collect()
        }

        fn procedure_context(&self, procedure_id: ProcedureId) -> Option<ProcedureContext> {
            self.events.lock().unwrap().iter().find_map(|event| {
                let event = event.as_any().downcast_ref::<ProcedureEvent>().unwrap();
                (event.procedure_id == procedure_id).then(|| event.context.clone())
            })
        }
    }

    impl EventRecorder for CapturingEventRecorder {
        fn record(&self, event: Box<dyn Event>) {
            self.events.lock().unwrap().push(event);
        }

        fn event_type_filter(&self) -> EventTypeFilterRef {
            self.event_type_filter.clone()
        }

        fn close(&self) {}
    }

    impl Default for CapturingEventRecorder {
        fn default() -> Self {
            Self::with_event_type_filter(Arc::new(EventTypeFilter::All))
        }
    }

    async fn wait_for_trigger(event_recorder: &CapturingEventRecorder, trigger: EventTrigger) {
        timeout(Duration::from_secs(1), async {
            while !event_recorder.triggers().contains(&trigger) {
                sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .unwrap();
    }

    #[derive(Debug)]
    struct TestProcedureEvent;

    impl Event for TestProcedureEvent {
        fn event_type(&self) -> &str {
            "test_procedure"
        }

        fn as_any(&self) -> &dyn std::any::Any {
            self
        }
    }

    #[test]
    fn test_manager_context() {
        let ctx = new_test_manager_context();
        let meta = Arc::new(test_util::procedure_meta_for_test());

        assert!(!ctx.contains_procedure(meta.id));
        assert!(ctx.state(meta.id).is_none());

        assert!(ctx.try_insert_procedure(meta.clone()));
        assert!(ctx.contains_procedure(meta.id));

        assert!(ctx.state(meta.id).unwrap().is_running());
        meta.set_state(ProcedureState::Done { output: None });
        assert!(ctx.state(meta.id).unwrap().is_done());
    }

    #[test]
    fn test_reset_runtime_state() {
        let ctx = new_test_manager_context();
        ctx.set_running();
        let mut meta = test_util::procedure_meta_for_test();
        meta.lock_key = LockKey::single_exclusive("test.reset_runtime_state");
        let meta = Arc::new(meta);
        let procedure_id = meta.id;

        assert!(ctx.try_insert_procedure(meta.clone()));
        ctx.finished_procedures
            .lock()
            .unwrap()
            .push_back((procedure_id, Instant::now()));
        ctx.spawn_runner_task(procedure_id, || {
            common_runtime::spawn_global(std::future::pending::<()>())
        });

        drop(
            ctx.key_lock
                .try_write("test.reset_runtime_state".to_string()),
        );
        drop(
            ctx.dynamic_key_lock
                .try_write("test.reset_runtime_state.dynamic".to_string()),
        );
        assert!(ctx.contains_procedure(procedure_id));
        assert_eq!(1, ctx.running_procedures.lock().unwrap().len());
        assert_eq!(1, ctx.finished_procedures.lock().unwrap().len());
        assert_eq!(1, ctx.runner_tasks.lock().unwrap().len());
        assert_eq!(1, ctx.key_lock.len());
        assert_eq!(1, ctx.dynamic_key_lock.len());

        ctx.reset_runtime_state();

        assert!(!ctx.contains_procedure(procedure_id));
        assert!(ctx.running_procedures.lock().unwrap().is_empty());
        assert!(ctx.finished_procedures.lock().unwrap().is_empty());
        assert!(ctx.runner_tasks.lock().unwrap().is_empty());
        assert!(ctx.key_lock.is_empty());
        assert!(ctx.dynamic_key_lock.is_empty());
    }

    #[test]
    fn test_spawn_runner_task_not_started_after_stop() {
        let ctx = new_test_manager_context();
        let procedure_id = ProcedureId::random();

        let spawned = Arc::new(AtomicBool::new(false));
        let spawned_in_task = spawned.clone();
        let started = ctx.spawn_runner_task(procedure_id, || {
            common_runtime::spawn_global(async move {
                spawned_in_task.store(true, AtomicOrdering::Relaxed);
            })
        });

        assert!(!started);
        assert!(!spawned.load(AtomicOrdering::Relaxed));
        assert!(ctx.runner_tasks.lock().unwrap().is_empty());
    }

    #[test]
    fn test_manager_context_insert_duplicate() {
        let ctx = new_test_manager_context();
        let meta = Arc::new(test_util::procedure_meta_for_test());

        assert!(ctx.try_insert_procedure(meta.clone()));
        assert!(!ctx.try_insert_procedure(meta));
    }

    fn new_child(parent_id: ProcedureId, ctx: &ManagerContext) -> ProcedureMetaRef {
        let mut child = test_util::procedure_meta_for_test();
        child.parent_id = Some(parent_id);
        let child = Arc::new(child);
        assert!(ctx.try_insert_procedure(child.clone()));

        let mut parent = Vec::new();
        ctx.find_procedures(&[parent_id], &mut parent);
        parent[0].push_child(child.id);

        child
    }

    #[test]
    fn test_procedures_in_tree() {
        let ctx = new_test_manager_context();
        let root = Arc::new(test_util::procedure_meta_for_test());
        assert!(ctx.try_insert_procedure(root.clone()));

        assert_eq!(1, ctx.procedures_in_tree(&root).len());

        let child1 = new_child(root.id, &ctx);
        let child2 = new_child(root.id, &ctx);

        let child3 = new_child(child1.id, &ctx);
        let child4 = new_child(child1.id, &ctx);

        let child5 = new_child(child2.id, &ctx);

        let expect = vec![
            root.id, child1.id, child2.id, child3.id, child4.id, child5.id,
        ];
        assert_eq!(expect, ctx.procedures_in_tree(&root));
    }

    #[derive(Debug)]
    struct ProcedureToLoad {
        content: String,
        lock_key: LockKey,
        poison_keys: PoisonKeys,
    }

    #[async_trait]
    impl Procedure for ProcedureToLoad {
        fn type_name(&self) -> &str {
            "ProcedureToLoad"
        }

        async fn execute(&mut self, _ctx: &Context) -> Result<Status> {
            Ok(Status::done())
        }

        fn dump(&self) -> Result<String> {
            Ok(self.content.clone())
        }

        fn lock_key(&self) -> LockKey {
            self.lock_key.clone()
        }

        fn poison_keys(&self) -> PoisonKeys {
            self.poison_keys.clone()
        }

        fn event(&self, _ctx: &EventContext<'_>) -> Option<Box<dyn Event>> {
            Some(Box::new(TestProcedureEvent))
        }
    }

    impl ProcedureToLoad {
        fn new(content: &str) -> ProcedureToLoad {
            ProcedureToLoad {
                content: content.to_string(),
                lock_key: LockKey::default(),
                poison_keys: PoisonKeys::default(),
            }
        }

        fn loader() -> BoxedProcedureLoader {
            let f = |json: &str| {
                let procedure = ProcedureToLoad::new(json);
                Ok(Box::new(procedure) as _)
            };
            Box::new(f)
        }
    }

    #[derive(Debug)]
    struct ParentWithChildProcedure {
        child_id: ProcedureId,
        child_submitted: bool,
    }

    #[async_trait]
    impl Procedure for ParentWithChildProcedure {
        fn type_name(&self) -> &str {
            "ParentWithChildProcedure"
        }

        async fn execute(&mut self, _ctx: &Context) -> Result<Status> {
            if self.child_submitted {
                return Ok(Status::done());
            }

            self.child_submitted = true;
            Ok(Status::suspended(
                vec![ProcedureWithId {
                    id: self.child_id,
                    procedure: Box::new(ProcedureToLoad::new("child")),
                    context: ProcedureContext::default(),
                }],
                false,
            ))
        }

        fn dump(&self) -> Result<String> {
            Ok(String::new())
        }

        fn lock_key(&self) -> LockKey {
            LockKey::default()
        }

        fn event(&self, _ctx: &EventContext<'_>) -> Option<Box<dyn Event>> {
            Some(Box::new(TestProcedureEvent))
        }
    }

    struct ParentWithBlockedChildProcedure {
        child_id: ProcedureId,
        ready: Option<oneshot::Sender<()>>,
        proceed: Option<oneshot::Receiver<()>>,
        child_submitted: bool,
    }

    #[async_trait]
    impl Procedure for ParentWithBlockedChildProcedure {
        fn type_name(&self) -> &str {
            "ParentWithBlockedChildProcedure"
        }

        async fn execute(&mut self, _ctx: &Context) -> Result<Status> {
            if self.child_submitted {
                return Ok(Status::done());
            }

            self.child_submitted = true;
            self.ready.take().unwrap().send(()).unwrap();
            self.proceed.take().unwrap().await.unwrap();
            Ok(Status::suspended(
                vec![ProcedureWithId {
                    id: self.child_id,
                    procedure: Box::new(ProcedureToLoad::new("child")),
                    context: ProcedureContext::default(),
                }],
                false,
            ))
        }

        fn dump(&self) -> Result<String> {
            Ok(String::new())
        }

        fn lock_key(&self) -> LockKey {
            LockKey::default()
        }

        fn event(&self, _ctx: &EventContext<'_>) -> Option<Box<dyn Event>> {
            Some(Box::new(TestProcedureEvent))
        }
    }

    struct FilterCapturingProcedure {
        captured_filter: Arc<Mutex<Option<EventTypeFilterRef>>>,
    }

    #[async_trait]
    impl Procedure for FilterCapturingProcedure {
        fn type_name(&self) -> &str {
            "FilterCapturingProcedure"
        }

        async fn execute(&mut self, _: &Context) -> Result<Status> {
            Ok(Status::done())
        }

        fn dump(&self) -> Result<String> {
            Ok(String::new())
        }

        fn lock_key(&self) -> LockKey {
            LockKey::default()
        }

        fn event(&self, ctx: &EventContext<'_>) -> Option<Box<dyn Event>> {
            *self.captured_filter.lock().unwrap() = Some(ctx.event_type_filter.clone());
            Some(Box::new(TestProcedureEvent))
        }
    }

    #[tokio::test]
    async fn test_event_filter_is_shared_with_procedure_context() {
        let dir = create_temp_dir("shared_event_filter");
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let event_type_filter = Arc::new(EventTypeFilter::Only(HashSet::from([String::from(
            "test_procedure",
        )])));
        let event_recorder = Arc::new(CapturingEventRecorder::with_event_type_filter(
            event_type_filter.clone(),
        ));
        let captured_filter = Arc::new(Mutex::new(None));
        let manager = LocalManager::new(
            ManagerConfig::default(),
            state_store,
            poison_manager,
            None,
            None,
        );
        manager.event_recorder_handle().install(event_recorder);
        manager.manager_ctx.start();

        manager
            .submit(ProcedureWithId {
                id: ProcedureId::random(),
                procedure: Box::new(FilterCapturingProcedure {
                    captured_filter: captured_filter.clone(),
                }),
                context: ProcedureContext::default(),
            })
            .await
            .unwrap();

        let captured_filter = timeout(Duration::from_secs(1), async {
            loop {
                if let Some(captured_filter) = captured_filter.lock().unwrap().clone() {
                    return captured_filter;
                }
                sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .unwrap();
        assert!(Arc::ptr_eq(&event_type_filter, &captured_filter));
        assert!(captured_filter.allows("test_procedure"));
        assert!(!captured_filter.allows("other_procedure"));
    }

    #[tokio::test]
    async fn test_fresh_submission_emits_submitted_event() {
        let dir = create_temp_dir("fresh_submission_event");
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let event_recorder = Arc::new(CapturingEventRecorder::default());
        let manager = LocalManager::new(
            ManagerConfig::default(),
            state_store,
            poison_manager,
            None,
            None,
        );
        manager
            .event_recorder_handle()
            .install(event_recorder.clone());
        manager.manager_ctx.start();

        manager
            .submit(ProcedureWithId {
                id: ProcedureId::random(),
                procedure: Box::new(ProcedureToLoad::new("fresh submission")),
                context: ProcedureContext::default(),
            })
            .await
            .unwrap();

        wait_for_trigger(&event_recorder, EventTrigger::Submitted).await;
    }

    #[test]
    fn test_submit_root_does_not_record_event_when_runner_is_not_spawned() {
        let dir = create_temp_dir("submit_root_without_runner");
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let event_recorder = Arc::new(CapturingEventRecorder::default());
        let manager = LocalManager::new(
            ManagerConfig::default(),
            state_store,
            poison_manager,
            None,
            None,
        );
        manager
            .event_recorder_handle()
            .install(event_recorder.clone());
        manager.manager_ctx.start();

        let procedure_id = ProcedureId::random();
        let runner_tasks = manager.manager_ctx.runner_tasks.lock().unwrap();
        std::thread::scope(|scope| {
            let handle = scope.spawn(|| {
                manager.submit_root(
                    procedure_id,
                    ProcedureState::Running,
                    0,
                    Box::new(ProcedureToLoad::new("submit root without runner")),
                    ProcedureContext::default(),
                    RootSubmissionOrigin::Fresh,
                )
            });

            let deadline = Instant::now() + Duration::from_secs(1);
            while !manager.manager_ctx.contains_procedure(procedure_id) {
                assert!(
                    Instant::now() < deadline,
                    "procedure metadata was not inserted"
                );
                std::thread::sleep(Duration::from_millis(10));
            }

            manager.manager_ctx.stop();
            drop(runner_tasks);

            let err = handle.join().unwrap().unwrap_err();
            assert_matches!(err, Error::ManagerNotStart { .. });
        });

        assert!(!manager.manager_ctx.contains_procedure(procedure_id));
        assert!(!event_recorder.triggers().contains(&EventTrigger::Submitted));
    }

    #[tokio::test]
    async fn test_child_submission_emits_submitted_event_for_child() {
        let dir = create_temp_dir("child_submission_event");
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let event_recorder = Arc::new(CapturingEventRecorder::default());
        let manager = LocalManager::new(
            ManagerConfig::default(),
            state_store,
            poison_manager,
            None,
            None,
        );
        manager
            .event_recorder_handle()
            .install(event_recorder.clone());
        manager.manager_ctx.start();

        let parent_id = ProcedureId::random();
        let child_id = ProcedureId::random();
        let mut watcher = manager
            .submit(ProcedureWithId {
                id: parent_id,
                procedure: Box::new(ParentWithChildProcedure {
                    child_id,
                    child_submitted: false,
                }),
                context: ProcedureContext::from_event_context(PersistentEventContext::new(
                    TriggerReason::AutoRepartition,
                )),
            })
            .await
            .unwrap();
        timeout(Duration::from_secs(1), async {
            while !watcher.borrow().is_done() {
                watcher.changed().await.unwrap();
            }
        })
        .await
        .unwrap();

        let procedure_events = event_recorder.procedure_events();
        assert!(procedure_events.contains(&(
            parent_id,
            EventTrigger::ChildSubmitted {
                procedure_id: child_id,
                outcome: ChildSubmissionOutcome::Accepted,
            },
        )));
        assert!(procedure_events.contains(&(child_id, EventTrigger::Submitted)));
        assert_eq!(
            event_recorder
                .procedure_context(child_id)
                .unwrap()
                .event_context
                .unwrap()
                .reason,
            TriggerReason::AutoRepartition
        );
    }

    #[tokio::test]
    async fn test_child_submission_does_not_emit_event_when_spawn_fails() {
        let dir = create_temp_dir("child_submission_spawn_failed_event");
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let event_recorder = Arc::new(CapturingEventRecorder::default());
        let manager = LocalManager::new(
            ManagerConfig::default(),
            state_store,
            poison_manager,
            None,
            None,
        );
        manager
            .event_recorder_handle()
            .install(event_recorder.clone());
        manager.manager_ctx.start();

        let parent_id = ProcedureId::random();
        let child_id = ProcedureId::random();
        let (ready_tx, ready_rx) = oneshot::channel();
        let (proceed_tx, proceed_rx) = oneshot::channel();
        manager
            .submit(ProcedureWithId {
                id: parent_id,
                procedure: Box::new(ParentWithBlockedChildProcedure {
                    child_id,
                    ready: Some(ready_tx),
                    proceed: Some(proceed_rx),
                    child_submitted: false,
                }),
                context: ProcedureContext::default(),
            })
            .await
            .unwrap();

        timeout(Duration::from_secs(1), ready_rx)
            .await
            .unwrap()
            .unwrap();
        let runner_tasks = manager.manager_ctx.runner_tasks.lock().unwrap();
        proceed_tx.send(()).unwrap();
        let deadline = Instant::now() + Duration::from_secs(1);
        while !manager.manager_ctx.contains_procedure(child_id) {
            assert!(
                Instant::now() < deadline,
                "child procedure was not inserted"
            );
            std::thread::yield_now();
        }

        manager.manager_ctx.stop();
        drop(runner_tasks);

        timeout(Duration::from_secs(1), async {
            while !event_recorder.procedure_events().contains(&(
                parent_id,
                EventTrigger::ChildSubmitted {
                    procedure_id: child_id,
                    outcome: ChildSubmissionOutcome::SpawnFailed,
                },
            )) {
                tokio::task::yield_now().await;
            }
        })
        .await
        .unwrap();

        assert!(!manager.manager_ctx.contains_procedure(child_id));
        assert!(
            !event_recorder
                .procedure_events()
                .contains(&(child_id, EventTrigger::Submitted))
        );

        manager.stop().await.unwrap();
    }

    #[tokio::test]
    async fn test_recovery_loads_legacy_procedure_message() {
        let dir = create_temp_dir("legacy_procedure_message_recovery");
        let object_store = test_util::new_object_store(&dir);
        let state_store = Arc::new(ObjectStateStore::new(object_store));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let event_recorder = Arc::new(CapturingEventRecorder::default());
        let manager = LocalManager::new(
            ManagerConfig {
                parent_path: "data/".to_string(),
                ..Default::default()
            },
            state_store.clone(),
            poison_manager,
            None,
            None,
        );
        manager
            .event_recorder_handle()
            .install(event_recorder.clone());
        manager.manager_ctx.start();
        manager
            .register_loader("ProcedureToLoad", ProcedureToLoad::loader())
            .unwrap();

        let procedure = ProcedureToLoad::new("legacy recovered submission");
        let procedure_id = ProcedureId::random();
        let key = format!("data/procedure/{procedure_id}/0000000000.step");
        // This is the persisted JSON shape written before ProcedureMessage had `context`.
        let legacy_message = serde_json::json!({
            "type_name": procedure.type_name(),
            "data": procedure.dump().unwrap(),
            "parent_id": null,
            "step": 0,
        });
        state_store
            .put(&key, serde_json::to_vec(&legacy_message).unwrap())
            .await
            .unwrap();

        manager.recover().await.unwrap();

        assert!(
            manager
                .procedure_state(procedure_id)
                .await
                .unwrap()
                .is_some()
        );
        wait_for_trigger(&event_recorder, EventTrigger::Recovered).await;
        assert_eq!(
            event_recorder.procedure_context(procedure_id),
            Some(ProcedureContext::default())
        );
    }

    #[tokio::test]
    async fn test_recovery_preserves_procedure_context() {
        let dir = create_temp_dir("recovery_submission_event");
        let object_store = test_util::new_object_store(&dir);
        let state_store = Arc::new(ObjectStateStore::new(object_store.clone()));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let event_recorder = Arc::new(CapturingEventRecorder::default());
        let manager = LocalManager::new(
            ManagerConfig {
                parent_path: "data/".to_string(),
                ..Default::default()
            },
            state_store,
            poison_manager,
            None,
            None,
        );
        manager
            .event_recorder_handle()
            .install(event_recorder.clone());
        manager.manager_ctx.start();
        manager
            .register_loader("ProcedureToLoad", ProcedureToLoad::loader())
            .unwrap();

        let procedure = ProcedureToLoad::new("recovered submission");
        let procedure_id = ProcedureId::random();
        let mut event_context =
            PersistentEventContext::new(TriggerReason::AutoRebalance).with_protocol("mysql");
        event_context.extensions.insert(
            "source".to_string(),
            serde_json::Value::String("test".to_string()),
        );
        let procedure_context = ProcedureContext {
            actor: Some("alice".to_string()),
            event_context: Some(event_context),
        };
        ProcedureStore::from_object_store(object_store)
            .store_procedure_with_context(
                procedure_id,
                0,
                procedure.type_name().to_string(),
                procedure.dump().unwrap(),
                None,
                procedure_context.clone(),
            )
            .await
            .unwrap();

        manager.recover().await.unwrap();

        assert!(
            manager
                .procedure_state(procedure_id)
                .await
                .unwrap()
                .is_some()
        );
        wait_for_trigger(&event_recorder, EventTrigger::Recovered).await;
        assert!(!event_recorder.triggers().contains(&EventTrigger::Submitted));
        assert_eq!(
            event_recorder.procedure_context(procedure_id),
            Some(procedure_context)
        );
    }

    #[tokio::test]
    async fn test_event_recorder_handle_installs_after_start() {
        let dir = create_temp_dir("set_event_recorder_after_start");
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(
            ManagerConfig::default(),
            state_store,
            poison_manager,
            None,
            None,
        );

        manager.start().await.unwrap();

        manager
            .event_recorder_handle()
            .install(Arc::new(CapturingEventRecorder::default()));

        manager.stop().await.unwrap();
    }

    #[derive(Debug)]
    struct BlockingProcedure {
        started_tx: Option<oneshot::Sender<()>>,
        dropped: Arc<AtomicBool>,
        lock_key: LockKey,
    }

    impl Drop for BlockingProcedure {
        fn drop(&mut self) {
            self.dropped.store(true, AtomicOrdering::Relaxed);
        }
    }

    #[async_trait]
    impl Procedure for BlockingProcedure {
        fn type_name(&self) -> &str {
            "BlockingProcedure"
        }

        async fn execute(&mut self, _ctx: &Context) -> Result<Status> {
            if let Some(tx) = self.started_tx.take() {
                let _ = tx.send(());
            }
            std::future::pending::<Result<Status>>().await
        }

        fn dump(&self) -> Result<String> {
            Ok(String::new())
        }

        fn lock_key(&self) -> LockKey {
            self.lock_key.clone()
        }
    }

    #[tokio::test]
    async fn test_stop_aborts_runner_and_resets_runtime_state() {
        let dir = create_temp_dir("stop_aborts_runner_and_resets_runtime_state");
        let config = ManagerConfig::default();
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);
        manager.start().await.unwrap();

        let procedure_id = ProcedureId::random();
        let (started_tx, started_rx) = oneshot::channel();
        let dropped = Arc::new(AtomicBool::new(false));
        let procedure = BlockingProcedure {
            started_tx: Some(started_tx),
            dropped: dropped.clone(),
            lock_key: LockKey::single_exclusive("test.stop_aborts_runner"),
        };

        manager
            .submit(ProcedureWithId {
                id: procedure_id,
                procedure: Box::new(procedure),
                context: ProcedureContext::default(),
            })
            .await
            .unwrap();
        timeout(Duration::from_secs(5), started_rx)
            .await
            .unwrap()
            .unwrap();

        assert!(manager.manager_ctx.contains_procedure(procedure_id));
        assert_eq!(
            1,
            manager.manager_ctx.running_procedures.lock().unwrap().len()
        );
        assert_eq!(1, manager.manager_ctx.runner_tasks.lock().unwrap().len());
        assert_eq!(1, manager.manager_ctx.key_lock.len());

        manager.stop().await.unwrap();

        assert!(dropped.load(AtomicOrdering::Relaxed));
        assert!(!manager.manager_ctx.running());
        assert!(!manager.manager_ctx.contains_procedure(procedure_id));
        assert!(
            manager
                .manager_ctx
                .running_procedures
                .lock()
                .unwrap()
                .is_empty()
        );
        assert!(
            manager
                .manager_ctx
                .finished_procedures
                .lock()
                .unwrap()
                .is_empty()
        );
        assert!(manager.manager_ctx.runner_tasks.lock().unwrap().is_empty());
        assert!(manager.manager_ctx.key_lock.is_empty());
        assert!(manager.manager_ctx.dynamic_key_lock.is_empty());
    }

    #[test]
    fn test_register_loader() {
        let dir = create_temp_dir("register");
        let config = ManagerConfig {
            parent_path: "data/".to_string(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            ..Default::default()
        };
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);
        manager.manager_ctx.start();

        manager
            .register_loader("ProcedureToLoad", ProcedureToLoad::loader())
            .unwrap();
        // Register duplicate loader.
        let err = manager
            .register_loader("ProcedureToLoad", ProcedureToLoad::loader())
            .unwrap_err();
        assert!(matches!(err, Error::LoaderConflict { .. }), "{err}");
    }

    #[tokio::test]
    async fn test_recover() {
        let dir = create_temp_dir("recover");
        let object_store = test_util::new_object_store(&dir);
        let config = ManagerConfig {
            parent_path: "data/".to_string(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            ..Default::default()
        };
        let state_store = Arc::new(ObjectStateStore::new(object_store.clone()));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);
        manager.manager_ctx.start();

        manager
            .register_loader("ProcedureToLoad", ProcedureToLoad::loader())
            .unwrap();

        // Prepare data
        let procedure_store = ProcedureStore::from_object_store(object_store.clone());
        let root: BoxedProcedure = Box::new(ProcedureToLoad::new("test recover manager"));
        let root_id = ProcedureId::random();
        // Prepare data for the root procedure.
        for step in 0..3 {
            let type_name = root.type_name().to_string();
            let data = root.dump().unwrap();
            procedure_store
                .store_procedure(root_id, step, type_name, data, None)
                .await
                .unwrap();
        }

        let child: BoxedProcedure = Box::new(ProcedureToLoad::new("a child procedure"));
        let child_id = ProcedureId::random();
        // Prepare data for the child procedure
        for step in 0..2 {
            let type_name = child.type_name().to_string();
            let data = child.dump().unwrap();
            procedure_store
                .store_procedure(child_id, step, type_name, data, Some(root_id))
                .await
                .unwrap();
        }

        // Recover the manager
        manager.recover().await.unwrap();

        // The manager should submit the root procedure.
        let _ = manager.procedure_state(root_id).await.unwrap().unwrap();
        // Since the mocked root procedure actually doesn't submit subprocedures, so there is no
        // related state.
        assert!(manager.procedure_state(child_id).await.unwrap().is_none());
    }

    #[tokio::test]
    async fn test_submit_procedure() {
        let dir = create_temp_dir("submit");
        let config = ManagerConfig {
            parent_path: "data/".to_string(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            ..Default::default()
        };
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);
        manager.manager_ctx.start();

        let procedure_id = ProcedureId::random();
        assert!(
            manager
                .procedure_state(procedure_id)
                .await
                .unwrap()
                .is_none()
        );
        assert!(manager.procedure_watcher(procedure_id).is_none());

        let mut procedure = ProcedureToLoad::new("submit");
        procedure.lock_key = LockKey::single_exclusive("test.submit");
        assert!(
            manager
                .submit(ProcedureWithId {
                    id: procedure_id,
                    procedure: Box::new(procedure),
                    context: ProcedureContext::default(),
                })
                .await
                .is_ok()
        );
        assert!(
            manager
                .procedure_state(procedure_id)
                .await
                .unwrap()
                .is_some()
        );
        // Wait for the procedure done.
        let mut watcher = manager.procedure_watcher(procedure_id).unwrap();
        watcher.changed().await.unwrap();
        assert!(watcher.borrow().is_done());

        // Try to submit procedure with same id again.
        let err = manager
            .submit(ProcedureWithId {
                id: procedure_id,
                procedure: Box::new(ProcedureToLoad::new("submit")),
                context: ProcedureContext::default(),
            })
            .await
            .unwrap_err();
        assert!(matches!(err, Error::DuplicateProcedure { .. }), "{err}");
    }

    #[tokio::test]
    async fn test_state_changed_on_err() {
        let dir = create_temp_dir("on_err");
        let config = ManagerConfig {
            parent_path: "data/".to_string(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            ..Default::default()
        };
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);
        manager.manager_ctx.start();

        #[derive(Debug)]
        struct MockProcedure {
            panic: bool,
        }

        #[async_trait]
        impl Procedure for MockProcedure {
            fn type_name(&self) -> &str {
                "MockProcedure"
            }

            async fn execute(&mut self, _ctx: &Context) -> Result<Status> {
                if self.panic {
                    // Test the runner can set the state to failed even the procedure
                    // panics.
                    panic!();
                } else {
                    Err(Error::external(MockError::new(StatusCode::Unexpected)))
                }
            }

            async fn rollback(&mut self, _: &Context) -> Result<()> {
                Ok(())
            }

            fn rollback_supported(&self) -> bool {
                true
            }

            fn dump(&self) -> Result<String> {
                Ok(String::new())
            }

            fn lock_key(&self) -> LockKey {
                LockKey::single_exclusive("test.submit")
            }

            fn poison_keys(&self) -> PoisonKeys {
                PoisonKeys::default()
            }
        }

        let check_procedure = |procedure| async {
            let procedure_id = ProcedureId::random();
            manager
                .submit(ProcedureWithId {
                    id: procedure_id,
                    procedure: Box::new(procedure),
                    context: ProcedureContext::default(),
                })
                .await
                .unwrap()
        };

        let mut watcher = check_procedure(MockProcedure { panic: false }).await;
        // Wait for the notification.
        watcher.changed().await.unwrap();
        assert!(watcher.borrow().is_prepare_rollback());
        watcher.changed().await.unwrap();
        assert!(watcher.borrow().is_rolling_back());
        watcher.changed().await.unwrap();
        assert!(watcher.borrow().is_failed());
        // The runner won't rollback a panicked procedure.
        let mut watcher = check_procedure(MockProcedure { panic: true }).await;
        watcher.changed().await.unwrap();
        assert!(watcher.borrow().is_failed());
    }

    #[tokio::test]
    async fn test_procedure_manager_stopped() {
        let dir = create_temp_dir("procedure_manager_stopped");
        let config = ManagerConfig {
            parent_path: "data/".to_string(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            ..Default::default()
        };
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);

        let mut procedure = ProcedureToLoad::new("submit");
        procedure.lock_key = LockKey::single_exclusive("test.submit");
        let procedure_id = ProcedureId::random();
        assert_matches!(
            manager
                .submit(ProcedureWithId {
                    id: procedure_id,
                    procedure: Box::new(procedure),
                    context: ProcedureContext::default(),
                })
                .await
                .unwrap_err(),
            error::Error::ManagerNotStart { .. }
        );
    }

    #[tokio::test]
    async fn test_procedure_manager_restart() {
        let dir = create_temp_dir("procedure_manager_restart");
        let config = ManagerConfig {
            parent_path: "data/".to_string(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            ..Default::default()
        };
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);

        manager.start().await.unwrap();
        manager.stop().await.unwrap();
        manager.start().await.unwrap();

        let mut procedure = ProcedureToLoad::new("submit");
        procedure.lock_key = LockKey::single_exclusive("test.submit");
        let procedure_id = ProcedureId::random();
        assert!(
            manager
                .submit(ProcedureWithId {
                    id: procedure_id,
                    procedure: Box::new(procedure),
                    context: ProcedureContext::default(),
                })
                .await
                .is_ok()
        );
        assert!(
            manager
                .procedure_state(procedure_id)
                .await
                .unwrap()
                .is_some()
        );
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_remove_outdated_meta_task() {
        let dir = create_temp_dir("remove_outdated_meta_task");
        let object_store = test_util::new_object_store(&dir);
        let config = ManagerConfig {
            parent_path: "data/".to_string(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            remove_outdated_meta_task_interval: Duration::from_millis(1),
            remove_outdated_meta_ttl: Duration::from_millis(1),
            max_running_procedures: 128,
        };
        let state_store = Arc::new(ObjectStateStore::new(object_store.clone()));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);
        manager.manager_ctx.set_running();

        let mut procedure = ProcedureToLoad::new("submit");
        procedure.lock_key = LockKey::single_exclusive("test.submit");
        let procedure_id = ProcedureId::random();
        assert!(
            manager
                .submit(ProcedureWithId {
                    id: procedure_id,
                    procedure: Box::new(procedure),
                    context: ProcedureContext::default(),
                })
                .await
                .is_ok()
        );
        let mut watcher = manager.procedure_watcher(procedure_id).unwrap();
        watcher.changed().await.unwrap();

        manager.start().await.unwrap();
        tokio::time::sleep(Duration::from_millis(300)).await;
        assert!(
            manager
                .procedure_state(procedure_id)
                .await
                .unwrap()
                .is_none()
        );

        // The remove_outdated_meta method has been stopped, so any procedure meta-data will not be automatically removed.
        manager.stop().await.unwrap();
        let mut procedure = ProcedureToLoad::new("submit");
        procedure.lock_key = LockKey::single_exclusive("test.submit");
        let procedure_id = ProcedureId::random();

        manager.manager_ctx.set_running();
        assert!(
            manager
                .submit(ProcedureWithId {
                    id: procedure_id,
                    procedure: Box::new(procedure),
                    context: ProcedureContext::default(),
                })
                .await
                .is_ok()
        );
        let mut watcher = manager.procedure_watcher(procedure_id).unwrap();
        watcher.changed().await.unwrap();
        tokio::time::sleep(Duration::from_millis(300)).await;
        assert!(
            manager
                .procedure_state(procedure_id)
                .await
                .unwrap()
                .is_some()
        );

        // After restart
        let mut procedure = ProcedureToLoad::new("submit");
        procedure.lock_key = LockKey::single_exclusive("test.submit");
        let procedure_id = ProcedureId::random();
        assert!(
            manager
                .submit(ProcedureWithId {
                    id: procedure_id,
                    procedure: Box::new(procedure),
                    context: ProcedureContext::default(),
                })
                .await
                .is_ok()
        );
        let mut watcher = manager.procedure_watcher(procedure_id).unwrap();
        watcher.changed().await.unwrap();

        manager.start().await.unwrap();
        tokio::time::sleep(Duration::from_millis(300)).await;
        assert!(
            manager
                .procedure_state(procedure_id)
                .await
                .unwrap()
                .is_none()
        );
    }

    #[tokio::test]
    async fn test_too_many_running_procedures() {
        let dir = create_temp_dir("too_many_running_procedures");
        let config = ManagerConfig {
            parent_path: "data/".to_string(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            max_running_procedures: 1,
            ..Default::default()
        };
        let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir)));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);
        manager.start().await.unwrap();

        manager
            .manager_ctx
            .running_procedures
            .lock()
            .unwrap()
            .insert(ProcedureId::random());

        // Submit a new procedure should fail.
        let mut procedure = ProcedureToLoad::new("submit");
        procedure.lock_key = LockKey::single_exclusive("test.submit");
        let procedure_id = ProcedureId::random();
        let err = manager
            .submit(ProcedureWithId {
                id: procedure_id,
                procedure: Box::new(procedure),
                context: ProcedureContext::default(),
            })
            .await
            .unwrap_err();
        assert!(matches!(err, Error::TooManyRunningProcedures { .. }));

        manager
            .manager_ctx
            .running_procedures
            .lock()
            .unwrap()
            .clear();

        // Submit a new procedure should succeed.
        let mut procedure = ProcedureToLoad::new("submit");
        procedure.lock_key = LockKey::single_exclusive("test.submit");
        assert!(
            manager
                .submit(ProcedureWithId {
                    id: procedure_id,
                    procedure: Box::new(procedure),
                    context: ProcedureContext::default(),
                })
                .await
                .is_ok()
        );
        assert!(
            manager
                .procedure_state(procedure_id)
                .await
                .unwrap()
                .is_some()
        );
        // Wait for the procedure done.
        let mut watcher = manager.procedure_watcher(procedure_id).unwrap();
        watcher.changed().await.unwrap();
        assert!(watcher.borrow().is_done());
    }

    #[derive(Debug)]
    struct ProcedureToRecover {
        content: String,
        lock_key: LockKey,
        notify: Option<Arc<Notify>>,
        poison_keys: PoisonKeys,
    }

    #[async_trait]
    impl Procedure for ProcedureToRecover {
        fn type_name(&self) -> &str {
            "ProcedureToRecover"
        }

        async fn execute(&mut self, _ctx: &Context) -> Result<Status> {
            Ok(Status::done())
        }

        fn dump(&self) -> Result<String> {
            Ok(self.content.clone())
        }

        fn lock_key(&self) -> LockKey {
            self.lock_key.clone()
        }

        fn recover(&mut self) -> Result<()> {
            self.notify.as_ref().unwrap().notify_one();
            Ok(())
        }

        fn poison_keys(&self) -> PoisonKeys {
            self.poison_keys.clone()
        }
    }

    impl ProcedureToRecover {
        fn new(content: &str) -> ProcedureToRecover {
            ProcedureToRecover {
                content: content.to_string(),
                lock_key: LockKey::default(),
                poison_keys: PoisonKeys::default(),
                notify: None,
            }
        }

        fn loader(notify: Arc<Notify>) -> BoxedProcedureLoader {
            let f = move |json: &str| {
                let procedure = ProcedureToRecover {
                    content: json.to_string(),
                    lock_key: LockKey::default(),
                    poison_keys: PoisonKeys::default(),
                    notify: Some(notify.clone()),
                };
                Ok(Box::new(procedure) as _)
            };
            Box::new(f)
        }
    }

    #[tokio::test]
    async fn test_procedure_recover() {
        common_telemetry::init_default_ut_logging();
        let dir = create_temp_dir("procedure_recover");
        let object_store = test_util::new_object_store(&dir);
        let config = ManagerConfig {
            parent_path: "data/".to_string(),
            max_retry_times: 3,
            retry_delay: Duration::from_millis(500),
            ..Default::default()
        };
        let state_store = Arc::new(ObjectStateStore::new(object_store.clone()));
        let poison_manager = Arc::new(InMemoryPoisonStore::new());
        let manager = LocalManager::new(config, state_store, poison_manager, None, None);
        manager.manager_ctx.start();

        let notify = Arc::new(Notify::new());
        manager
            .register_loader(
                "ProcedureToRecover",
                ProcedureToRecover::loader(notify.clone()),
            )
            .unwrap();

        // Prepare data
        let procedure_store = ProcedureStore::from_object_store(object_store.clone());
        let root: BoxedProcedure = Box::new(ProcedureToRecover::new("test procedure recovery"));
        let root_id = ProcedureId::random();
        // Prepare data for the root procedure.
        for step in 0..3 {
            let type_name = root.type_name().to_string();
            let data = root.dump().unwrap();
            procedure_store
                .store_procedure(root_id, step, type_name, data, None)
                .await
                .unwrap();
        }

        // Recover the manager
        manager.recover().await.unwrap();
        timeout(Duration::from_secs(10), notify.notified())
            .await
            .unwrap();
    }
}