common_procedure/
procedure.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::any::Any;
16use std::fmt;
17use std::fmt::Display;
18use std::str::FromStr;
19use std::sync::Arc;
20
21use async_trait::async_trait;
22use common_event_recorder::{Event, EventTypeFilterRef, PersistentEventContext};
23use serde::{Deserialize, Serialize};
24use smallvec::{SmallVec, smallvec};
25use snafu::{ResultExt, Snafu};
26use tokio::sync::watch::Receiver;
27use uuid::Uuid;
28
29use crate::error::{self, Error, Result};
30use crate::local::DynamicKeyLockGuard;
31use crate::watcher::Watcher;
32
33pub type Output = Arc<dyn Any + Send + Sync>;
34
35/// Context attached to a procedure submission and inherited by its children.
36#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
37pub struct ProcedureContext {
38    /// Effective user that submitted the procedure.
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub actor: Option<String>,
41    /// Context describing why and how the procedure was submitted.
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub event_context: Option<PersistentEventContext>,
44}
45
46impl ProcedureContext {
47    pub fn from_event_context(event_context: PersistentEventContext) -> Self {
48        Self {
49            actor: None,
50            event_context: Some(event_context),
51        }
52    }
53
54    pub fn is_empty(&self) -> bool {
55        self.actor.is_none() && self.event_context.is_none()
56    }
57}
58
59/// Procedure execution status.
60#[derive(Debug)]
61pub enum Status {
62    /// The procedure is still executing.
63    Executing {
64        /// Whether the framework needs to persist the procedure.
65        persist: bool,
66        /// Whether the framework needs to clean the poisons.
67        clean_poisons: bool,
68    },
69    /// The procedure has suspended itself and is waiting for subprocedures.
70    Suspended {
71        subprocedures: Vec<ProcedureWithId>,
72        /// Whether the framework needs to persist the procedure.
73        persist: bool,
74    },
75    /// The procedure is poisoned.
76    Poisoned {
77        /// The keys that cause the procedure to be poisoned.
78        keys: PoisonKeys,
79        /// The error that cause the procedure to be poisoned.
80        error: Error,
81    },
82    /// the procedure is done.
83    Done { output: Option<Output> },
84}
85
86impl Status {
87    /// Returns a [Status::Suspended] with given `subprocedures` and `persist` flag.
88    pub fn suspended(subprocedures: Vec<ProcedureWithId>, persist: bool) -> Status {
89        Status::Suspended {
90            subprocedures,
91            persist,
92        }
93    }
94
95    /// Returns a [Status::Poisoned] with given `keys` and `error`.
96    pub fn poisoned(keys: impl IntoIterator<Item = PoisonKey>, error: Error) -> Status {
97        Status::Poisoned {
98            keys: PoisonKeys::new(keys),
99            error,
100        }
101    }
102
103    /// Returns a [Status::Executing] with given `persist` flag.
104    pub fn executing(persist: bool) -> Status {
105        Status::Executing {
106            persist,
107            clean_poisons: false,
108        }
109    }
110
111    /// Returns a [Status::Executing] with given `persist` flag and clean poisons.
112    pub fn executing_with_clean_poisons(persist: bool) -> Status {
113        Status::Executing {
114            persist,
115            clean_poisons: true,
116        }
117    }
118
119    /// Returns a [Status::Done] without output.
120    pub fn done() -> Status {
121        Status::Done { output: None }
122    }
123
124    #[cfg(any(test, feature = "testing"))]
125    /// Downcasts [Status::Done]'s output to &T
126    ///  #Panic:
127    /// - if [Status] is not the [Status::Done].
128    /// - if the output is None.
129    pub fn downcast_output_ref<T: 'static>(&self) -> Option<&T> {
130        if let Status::Done { output } = self {
131            output
132                .as_ref()
133                .expect("Try to downcast the output of Status::Done, but the output is None")
134                .downcast_ref()
135        } else {
136            panic!("Expected the Status::Done, but got: {:?}", self)
137        }
138    }
139
140    /// Returns a [Status::Done] with output.
141    pub fn done_with_output<T: Any + Send + Sync>(output: T) -> Status {
142        Status::Done {
143            output: Some(Arc::new(output)),
144        }
145    }
146    /// Returns `true` if the procedure is done.
147    pub fn is_done(&self) -> bool {
148        matches!(self, Status::Done { .. })
149    }
150
151    /// Returns `true` if the procedure needs the framework to persist its intermediate state.
152    pub fn need_persist(&self) -> bool {
153        match self {
154            // If the procedure is done/poisoned, the framework doesn't need to persist the procedure
155            // anymore. It only needs to mark the procedure as committed.
156            Status::Executing { persist, .. } | Status::Suspended { persist, .. } => *persist,
157            Status::Done { .. } | Status::Poisoned { .. } => false,
158        }
159    }
160
161    /// Returns `true` if the framework needs to clean the poisons.
162    pub fn need_clean_poisons(&self) -> bool {
163        match self {
164            Status::Executing { clean_poisons, .. } => *clean_poisons,
165            Status::Done { .. } => true,
166            _ => false,
167        }
168    }
169}
170
171/// [ContextProvider] provides information about procedures in the [ProcedureManager].
172#[async_trait]
173pub trait ContextProvider: Send + Sync {
174    /// Query the procedure state.
175    async fn procedure_state(&self, procedure_id: ProcedureId) -> Result<Option<ProcedureState>>;
176
177    async fn procedure_state_receiver(
178        &self,
179        procedure_id: ProcedureId,
180    ) -> Result<Option<Receiver<ProcedureState>>>;
181
182    /// Try to put a poison key for a procedure.
183    ///
184    /// This method is used to mark a resource as being operated on by a procedure.
185    /// If the poison key already exists with a different value, the operation will fail.
186    async fn try_put_poison(&self, key: &PoisonKey, procedure_id: ProcedureId) -> Result<()>;
187
188    /// Acquires a key lock for the procedure.
189    async fn acquire_lock(&self, key: &StringKey) -> DynamicKeyLockGuard;
190}
191
192/// Reference-counted pointer to [ContextProvider].
193pub type ContextProviderRef = Arc<dyn ContextProvider>;
194
195/// Procedure execution context.
196#[derive(Clone)]
197pub struct Context {
198    /// Id of the procedure.
199    pub procedure_id: ProcedureId,
200    /// [ProcedureManager] context provider.
201    pub provider: ContextProviderRef,
202    /// Event context inherited from the root submission.
203    pub event_context: Option<PersistentEventContext>,
204}
205
206impl Context {
207    /// Returns true if current procedure state is retrying.
208    pub async fn is_retrying(&self) -> Option<bool> {
209        self.provider
210            .procedure_state(self.procedure_id)
211            .await
212            .ok()
213            .flatten()
214            .map(|s| s.is_retrying())
215    }
216}
217
218/// A `Procedure` represents an operation or a set of operations to be performed step-by-step.
219#[async_trait]
220pub trait Procedure: Send {
221    /// Type name of the procedure.
222    fn type_name(&self) -> &str;
223
224    /// Execute the procedure.
225    ///
226    /// The implementation must be idempotent.
227    async fn execute(&mut self, ctx: &Context) -> Result<Status>;
228
229    /// Rollback the failed procedure.
230    ///
231    /// The implementation must be idempotent.
232    async fn rollback(&mut self, _: &Context) -> Result<()> {
233        error::RollbackNotSupportedSnafu {}.fail()
234    }
235
236    /// Indicates whether it supports rolling back the procedure.
237    fn rollback_supported(&self) -> bool {
238        false
239    }
240
241    /// Dump the state of the procedure to a string.
242    fn dump(&self) -> Result<String>;
243
244    /// The hook is called after the procedure recovery.
245    fn recover(&mut self) -> Result<()> {
246        Ok(())
247    }
248
249    /// Returns the [LockKey] that this procedure needs to acquire.
250    fn lock_key(&self) -> LockKey;
251
252    /// Returns the [PoisonKeys] that may cause this procedure to become poisoned during execution.
253    fn poison_keys(&self) -> PoisonKeys {
254        PoisonKeys::default()
255    }
256
257    /// Builds an event for a framework lifecycle trigger.
258    ///
259    /// The hook is called with the current procedure instance, so an event can
260    /// include state that was produced after the procedure was submitted. A
261    /// return value of `None` means that this trigger should not be recorded.
262    /// Events that share an [`Event::event_type`] must return identical
263    /// [`Event::extra_schema`] values. The event recorder batches events by type
264    /// and rejects incompatible schemas; use a distinct event type for a
265    /// different schema.
266    fn event(&self, _ctx: &EventContext<'_>) -> Option<Box<dyn Event>> {
267        None
268    }
269}
270
271/// Framework-owned context supplied when a procedure builds a lifecycle event.
272pub struct EventContext<'a> {
273    /// Id of the procedure associated with the event.
274    pub procedure_id: ProcedureId,
275    /// Current framework state of the procedure.
276    pub lifecycle_state: &'a ProcedureState,
277    /// Lifecycle action that caused the event hook to be called.
278    pub trigger: EventTrigger,
279    /// Event types retained by the configured recorder.
280    pub event_type_filter: EventTypeFilterRef,
281    /// Event context inherited from the root submission.
282    pub event_context: Option<&'a PersistentEventContext>,
283}
284
285/// Lifecycle action that causes the framework to invoke [`Procedure::event`].
286///
287/// It is recorded as a tagged JSON object in the `procedure_trigger` event envelope column.
288#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
289#[serde(tag = "type")]
290pub enum EventTrigger {
291    /// The procedure was submitted to the manager.
292    Submitted,
293    /// The root procedure was recovered from persisted state.
294    Recovered,
295    /// A child submission was attempted.
296    ChildSubmitted {
297        /// The submitted child procedure.
298        procedure_id: ProcedureId,
299        /// The result of the submission attempt.
300        outcome: ChildSubmissionOutcome,
301    },
302    /// Procedure execution is being retried.
303    Retrying {
304        /// Phase in which the retry occurs.
305        phase: RetryPhase,
306        /// Retry attempt within the current runner lifecycle.
307        attempt: u32,
308    },
309    /// Procedure rollback is starting.
310    RollingBack,
311    /// The procedure reached a successful terminal state.
312    Succeeded,
313    /// The procedure reached a failed terminal state.
314    Failed,
315    /// The procedure was poisoned and cannot proceed.
316    Poisoned,
317}
318
319/// Phase of a procedure retry.
320#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
321pub enum RetryPhase {
322    /// Retrying procedure execution.
323    Execute,
324    /// Retrying procedure rollback.
325    Rollback,
326}
327
328/// Outcome of submitting a child procedure.
329#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
330pub enum ChildSubmissionOutcome {
331    Accepted,
332    AlreadyAccepted,
333    ManagerStopped,
334    SpawnFailed,
335}
336
337#[async_trait]
338impl<T: Procedure + ?Sized> Procedure for Box<T> {
339    fn type_name(&self) -> &str {
340        (**self).type_name()
341    }
342
343    async fn execute(&mut self, ctx: &Context) -> Result<Status> {
344        (**self).execute(ctx).await
345    }
346
347    async fn rollback(&mut self, ctx: &Context) -> Result<()> {
348        (**self).rollback(ctx).await
349    }
350
351    fn rollback_supported(&self) -> bool {
352        (**self).rollback_supported()
353    }
354
355    fn dump(&self) -> Result<String> {
356        (**self).dump()
357    }
358
359    fn lock_key(&self) -> LockKey {
360        (**self).lock_key()
361    }
362
363    fn poison_keys(&self) -> PoisonKeys {
364        (**self).poison_keys()
365    }
366
367    fn event(&self, ctx: &EventContext<'_>) -> Option<Box<dyn Event>> {
368        (**self).event(ctx)
369    }
370}
371
372#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
373pub struct PoisonKey(String);
374
375impl Display for PoisonKey {
376    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
377        write!(f, "{}", self.0)
378    }
379}
380
381impl PoisonKey {
382    /// Creates a new [PoisonKey] from a [String].
383    pub fn new(key: impl Into<String>) -> Self {
384        Self(key.into())
385    }
386}
387
388/// A collection of [PoisonKey]s.
389///
390/// This type is used to represent the keys that may cause the procedure to become poisoned during execution.
391#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Default)]
392pub struct PoisonKeys(SmallVec<[PoisonKey; 2]>);
393
394impl PoisonKeys {
395    /// Creates a new [PoisonKeys] from a [String].
396    pub fn single(key: impl Into<String>) -> Self {
397        Self(smallvec![PoisonKey::new(key)])
398    }
399
400    /// Creates a new [PoisonKeys] from a [PoisonKey].
401    pub fn new(keys: impl IntoIterator<Item = PoisonKey>) -> Self {
402        Self(keys.into_iter().collect())
403    }
404
405    /// Returns `true` if the [PoisonKeys] contains the given [PoisonKey].
406    pub fn contains(&self, key: &PoisonKey) -> bool {
407        self.0.contains(key)
408    }
409
410    /// Returns an iterator over the [PoisonKey]s.
411    pub fn iter(&self) -> impl Iterator<Item = &PoisonKey> {
412        self.0.iter()
413    }
414}
415
416#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
417pub enum StringKey {
418    Share(String),
419    Exclusive(String),
420}
421
422/// Keys to identify required locks.
423///
424/// [LockKey] always sorts keys lexicographically so that they can be acquired
425/// in the same order.
426/// Most procedures should only acquire 1 ~ 2 locks so we use smallvec to hold keys.
427#[derive(Clone, Debug, Default, PartialEq, Eq)]
428pub struct LockKey(SmallVec<[StringKey; 2]>);
429
430impl StringKey {
431    pub fn into_string(self) -> String {
432        match self {
433            StringKey::Share(s) => s,
434            StringKey::Exclusive(s) => s,
435        }
436    }
437
438    pub fn as_string(&self) -> &String {
439        match self {
440            StringKey::Share(s) => s,
441            StringKey::Exclusive(s) => s,
442        }
443    }
444}
445
446impl LockKey {
447    /// Returns a new [LockKey] with only one key.
448    pub fn single(key: impl Into<StringKey>) -> LockKey {
449        LockKey(smallvec![key.into()])
450    }
451
452    /// Returns a new [LockKey] with only one key.
453    pub fn single_exclusive(key: impl Into<String>) -> LockKey {
454        LockKey(smallvec![StringKey::Exclusive(key.into())])
455    }
456
457    /// Returns a new [LockKey] with keys from specific `iter`.
458    pub fn new(iter: impl IntoIterator<Item = StringKey>) -> LockKey {
459        let mut vec: SmallVec<_> = iter.into_iter().collect();
460        vec.sort();
461        // Dedup keys to avoid acquiring the same key multiple times.
462        vec.dedup();
463        LockKey(vec)
464    }
465
466    /// Returns a new [LockKey] with keys from specific `iter`.
467    pub fn new_exclusive(iter: impl IntoIterator<Item = String>) -> LockKey {
468        Self::new(iter.into_iter().map(StringKey::Exclusive))
469    }
470
471    /// Returns the keys to lock.
472    pub fn keys_to_lock(&self) -> impl Iterator<Item = &StringKey> {
473        self.0.iter()
474    }
475
476    /// Returns the keys to lock.
477    pub fn get_keys(&self) -> Vec<String> {
478        self.0.iter().map(|key| format!("{:?}", key)).collect()
479    }
480}
481
482/// Boxed [Procedure].
483pub type BoxedProcedure = Box<dyn Procedure>;
484
485/// A procedure with specific id.
486pub struct ProcedureWithId {
487    /// Id of the procedure.
488    pub id: ProcedureId,
489    pub procedure: BoxedProcedure,
490    /// Context associated with this procedure submission.
491    pub context: ProcedureContext,
492}
493
494impl ProcedureWithId {
495    /// Returns a new [ProcedureWithId] that holds specific `procedure`
496    /// and a random [ProcedureId].
497    pub fn with_random_id(procedure: BoxedProcedure) -> ProcedureWithId {
498        ProcedureWithId {
499            id: ProcedureId::random(),
500            procedure,
501            context: ProcedureContext::default(),
502        }
503    }
504
505    /// Attaches a complete persisted context to this procedure submission.
506    pub fn with_context(mut self, context: ProcedureContext) -> Self {
507        self.context = context;
508        self
509    }
510}
511
512impl fmt::Debug for ProcedureWithId {
513    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
514        write!(f, "{}-{}", self.procedure.type_name(), self.id)
515    }
516}
517
518#[derive(Debug, Snafu)]
519pub struct ParseIdError {
520    source: uuid::Error,
521}
522
523/// Unique id for [Procedure].
524#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
525pub struct ProcedureId(Uuid);
526
527impl ProcedureId {
528    /// Returns a new unique [ProcedureId] randomly.
529    pub fn random() -> ProcedureId {
530        ProcedureId(Uuid::new_v4())
531    }
532
533    /// Parses id from string.
534    pub fn parse_str(input: &str) -> std::result::Result<ProcedureId, ParseIdError> {
535        Uuid::parse_str(input)
536            .map(ProcedureId)
537            .context(ParseIdSnafu)
538    }
539}
540
541impl fmt::Display for ProcedureId {
542    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
543        write!(f, "{}", self.0)
544    }
545}
546
547impl FromStr for ProcedureId {
548    type Err = ParseIdError;
549
550    fn from_str(s: &str) -> std::result::Result<ProcedureId, ParseIdError> {
551        ProcedureId::parse_str(s)
552    }
553}
554
555/// Loader to recover the [Procedure] instance from serialized data.
556pub type BoxedProcedureLoader = Box<dyn Fn(&str) -> Result<BoxedProcedure> + Send>;
557
558/// State of a submitted procedure.
559#[derive(Debug, Default, Clone)]
560pub enum ProcedureState {
561    /// The procedure is running.
562    #[default]
563    Running,
564    /// The procedure is finished.
565    Done { output: Option<Output> },
566    /// The procedure is failed and can be retried.
567    Retrying { error: Arc<Error> },
568    /// The procedure is failed and commits state before rolling back the procedure.
569    PrepareRollback { error: Arc<Error> },
570    /// The procedure is failed and can be rollback.
571    RollingBack { error: Arc<Error> },
572    /// The procedure is failed and cannot proceed anymore.
573    Failed { error: Arc<Error> },
574    /// The procedure is poisoned.
575    Poisoned { keys: PoisonKeys, error: Arc<Error> },
576}
577
578impl ProcedureState {
579    /// Returns a [ProcedureState] with failed state.
580    pub fn failed(error: Arc<Error>) -> ProcedureState {
581        ProcedureState::Failed { error }
582    }
583
584    /// Returns a [ProcedureState] with prepare rollback state.
585    pub fn prepare_rollback(error: Arc<Error>) -> ProcedureState {
586        ProcedureState::PrepareRollback { error }
587    }
588
589    /// Returns a [ProcedureState] with rolling back state.
590    pub fn rolling_back(error: Arc<Error>) -> ProcedureState {
591        ProcedureState::RollingBack { error }
592    }
593
594    /// Returns a [ProcedureState] with retrying state.
595    pub fn retrying(error: Arc<Error>) -> ProcedureState {
596        ProcedureState::Retrying { error }
597    }
598
599    /// Returns a [ProcedureState] with poisoned state.
600    pub fn poisoned(keys: PoisonKeys, error: Arc<Error>) -> ProcedureState {
601        ProcedureState::Poisoned { keys, error }
602    }
603
604    /// Returns true if the procedure state is running.
605    pub fn is_running(&self) -> bool {
606        matches!(self, ProcedureState::Running)
607    }
608
609    /// Returns true if the procedure state is done.
610    pub fn is_done(&self) -> bool {
611        matches!(self, ProcedureState::Done { .. })
612    }
613
614    /// Returns true if the procedure state is poisoned.
615    pub fn is_poisoned(&self) -> bool {
616        matches!(self, ProcedureState::Poisoned { .. })
617    }
618
619    /// Returns true if the procedure state failed.
620    pub fn is_failed(&self) -> bool {
621        matches!(self, ProcedureState::Failed { .. })
622    }
623
624    /// Returns true if the procedure state is retrying.
625    pub fn is_retrying(&self) -> bool {
626        matches!(self, ProcedureState::Retrying { .. })
627    }
628
629    /// Returns true if the procedure state is rolling back.
630    pub fn is_rolling_back(&self) -> bool {
631        matches!(self, ProcedureState::RollingBack { .. })
632    }
633
634    /// Returns true if the procedure state is prepare rollback.
635    pub fn is_prepare_rollback(&self) -> bool {
636        matches!(self, ProcedureState::PrepareRollback { .. })
637    }
638
639    /// Returns the error.
640    pub fn error(&self) -> Option<&Arc<Error>> {
641        match self {
642            ProcedureState::Failed { error } => Some(error),
643            ProcedureState::Retrying { error } => Some(error),
644            ProcedureState::RollingBack { error } => Some(error),
645            ProcedureState::Poisoned { error, .. } => Some(error),
646            _ => None,
647        }
648    }
649
650    /// Return the string values of the enum field names.
651    pub fn as_str_name(&self) -> &str {
652        match self {
653            ProcedureState::Running => "Running",
654            ProcedureState::Done { .. } => "Done",
655            ProcedureState::Retrying { .. } => "Retrying",
656            ProcedureState::Failed { .. } => "Failed",
657            ProcedureState::PrepareRollback { .. } => "PrepareRollback",
658            ProcedureState::RollingBack { .. } => "RollingBack",
659            ProcedureState::Poisoned { .. } => "Poisoned",
660        }
661    }
662}
663
664/// The initial procedure state.
665#[derive(Debug, Clone)]
666pub enum InitProcedureState {
667    Running,
668    RollingBack,
669}
670
671// TODO(yingwen): Shutdown
672/// `ProcedureManager` executes [Procedure] submitted to it.
673#[async_trait]
674pub trait ProcedureManager: Send + Sync + 'static {
675    /// Registers loader for specific procedure type `name`.
676    fn register_loader(&self, name: &str, loader: BoxedProcedureLoader) -> Result<()>;
677
678    /// Starts the background GC task.
679    ///
680    /// Recovers unfinished procedures and reruns them.
681    ///
682    /// Callers should ensure all loaders are registered.
683    async fn start(&self) -> Result<()>;
684
685    /// Stops the background GC task.
686    async fn stop(&self) -> Result<()>;
687
688    /// Submits a procedure to execute.
689    ///
690    /// Returns a [Watcher] to watch the created procedure.
691    async fn submit(&self, procedure: ProcedureWithId) -> Result<Watcher>;
692
693    /// Query the procedure state.
694    ///
695    /// Returns `Ok(None)` if the procedure doesn't exist.
696    async fn procedure_state(&self, procedure_id: ProcedureId) -> Result<Option<ProcedureState>>;
697
698    /// Returns a [Watcher] to watch [ProcedureState] of specific procedure.
699    fn procedure_watcher(&self, procedure_id: ProcedureId) -> Option<Watcher>;
700
701    /// Returns the details of the procedure.
702    async fn list_procedures(&self) -> Result<Vec<ProcedureInfo>>;
703
704    /// Returns whether persisted unfinished procedures contain any requested type.
705    ///
706    /// This inspects durable state without submitting procedures for execution.
707    async fn has_unfinished_procedure(&self, type_names: &[&str]) -> Result<bool>;
708}
709
710/// Ref-counted pointer to the [ProcedureManager].
711pub type ProcedureManagerRef = Arc<dyn ProcedureManager>;
712
713#[derive(Debug, Clone)]
714pub struct ProcedureInfo {
715    /// Id of this procedure.
716    pub id: ProcedureId,
717    /// Type name of this procedure.
718    pub type_name: String,
719    /// Start execution time of this procedure.
720    pub start_time_ms: i64,
721    /// End execution time of this procedure.
722    pub end_time_ms: i64,
723    /// status of this procedure.
724    pub state: ProcedureState,
725    /// Lock keys of this procedure.
726    pub lock_keys: Vec<String>,
727}
728
729#[cfg(test)]
730mod tests {
731    use async_trait::async_trait;
732    use common_error::mock::MockError;
733    use common_error::status_code::StatusCode;
734
735    use super::*;
736
737    struct DefaultEventProcedure;
738
739    #[async_trait]
740    impl Procedure for DefaultEventProcedure {
741        fn type_name(&self) -> &str {
742            "default_event"
743        }
744
745        async fn execute(&mut self, _: &Context) -> Result<Status> {
746            Ok(Status::done())
747        }
748
749        fn dump(&self) -> Result<String> {
750            Ok(String::new())
751        }
752
753        fn lock_key(&self) -> LockKey {
754            LockKey::default()
755        }
756    }
757
758    #[test]
759    fn test_default_procedure_event_hook() {
760        let state = ProcedureState::Running;
761        let context = EventContext {
762            procedure_id: ProcedureId::random(),
763            lifecycle_state: &state,
764            trigger: EventTrigger::Succeeded,
765            event_type_filter: Arc::new(common_event_recorder::EventTypeFilter::All),
766            event_context: None,
767        };
768
769        assert!(DefaultEventProcedure.event(&context).is_none());
770        assert!(Box::new(DefaultEventProcedure).event(&context).is_none());
771    }
772
773    #[test]
774    fn test_status() {
775        let status = Status::executing(false);
776        assert!(!status.need_persist());
777
778        let status = Status::executing(true);
779        assert!(status.need_persist());
780
781        let status = Status::executing_with_clean_poisons(false);
782        assert!(status.need_clean_poisons());
783
784        let status = Status::executing_with_clean_poisons(true);
785        assert!(status.need_clean_poisons());
786
787        let status = Status::Suspended {
788            subprocedures: Vec::new(),
789            persist: false,
790        };
791        assert!(!status.need_persist());
792
793        let status = Status::Suspended {
794            subprocedures: Vec::new(),
795            persist: true,
796        };
797        assert!(status.need_persist());
798
799        let status = Status::done();
800        assert!(!status.need_persist());
801        assert!(status.need_clean_poisons());
802    }
803
804    #[test]
805    fn test_lock_key() {
806        let entity = "catalog.schema.my_table";
807        let key = LockKey::single_exclusive(entity);
808        assert_eq!(
809            vec![&StringKey::Exclusive(entity.to_string())],
810            key.keys_to_lock().collect::<Vec<_>>()
811        );
812
813        let key = LockKey::new_exclusive([
814            "b".to_string(),
815            "c".to_string(),
816            "a".to_string(),
817            "c".to_string(),
818        ]);
819        assert_eq!(
820            vec![
821                &StringKey::Exclusive("a".to_string()),
822                &StringKey::Exclusive("b".to_string()),
823                &StringKey::Exclusive("c".to_string())
824            ],
825            key.keys_to_lock().collect::<Vec<_>>()
826        );
827    }
828
829    #[test]
830    fn test_procedure_id() {
831        let id = ProcedureId::random();
832        let uuid_str = id.to_string();
833        assert_eq!(id.0.to_string(), uuid_str);
834
835        let parsed = ProcedureId::parse_str(&uuid_str).unwrap();
836        assert_eq!(id, parsed);
837        let parsed = uuid_str.parse().unwrap();
838        assert_eq!(id, parsed);
839    }
840
841    #[test]
842    fn test_procedure_id_serialization() {
843        let id = ProcedureId::random();
844        let json = serde_json::to_string(&id).unwrap();
845        assert_eq!(format!("\"{id}\""), json);
846
847        let parsed = serde_json::from_str(&json).unwrap();
848        assert_eq!(id, parsed);
849    }
850
851    #[test]
852    fn test_procedure_state() {
853        assert!(ProcedureState::Running.is_running());
854        assert!(ProcedureState::Running.error().is_none());
855        assert!(ProcedureState::Done { output: None }.is_done());
856
857        let state = ProcedureState::failed(Arc::new(Error::external(MockError::new(
858            StatusCode::Unexpected,
859        ))));
860        assert!(state.is_failed());
861        let _ = state.error().unwrap();
862    }
863}
common_procedure/procedure.rs

common_procedure/
procedure.rs
