feat(procedure): Support multi-lock keys and querying procedure state from context (#1006)

* feat: Add ContextProvider to Context

So procedures can query states of other procedures via the
ContextProvider and they don't need to hold a ProcedureManagerRef

* feat: Procedure supports acquring multiple lock keys

* test: Use multi-locks in test

* feat: Add keys_to_lock/unlock

Author: evenyag

Cargo.lock

@@ -13,6 +13,7 @@ futures.workspace = true
 object-store = { path = "../../object-store" }
 serde.workspace = true
 serde_json = "1.0"
+smallvec = "1"
 snafu.workspace = true
 tokio.workspace = true
 uuid.workspace = true

src/common/procedure/Cargo.toml

@@ -24,6 +24,6 @@ mod store;
 
 pub use crate::error::{Error, Result};
 pub use crate::procedure::{
-    BoxedProcedure, Context, LockKey, Procedure, ProcedureId, ProcedureManager,
+    BoxedProcedure, Context, ContextProvider, LockKey, Procedure, ProcedureId, ProcedureManager,
     ProcedureManagerRef, ProcedureState, ProcedureWithId, Status,
 };

src/common/procedure/src/lib.rs

@@ -30,7 +30,8 @@ use crate::local::runner::Runner;
 use crate::procedure::BoxedProcedureLoader;
 use crate::store::{ObjectStateStore, ProcedureMessage, ProcedureStore, StateStoreRef};
 use crate::{
-    BoxedProcedure, LockKey, ProcedureId, ProcedureManager, ProcedureState, ProcedureWithId,
+    BoxedProcedure, ContextProvider, LockKey, ProcedureId, ProcedureManager, ProcedureState,
+    ProcedureWithId,
 };
 
 /// Mutable metadata of a procedure during execution.
@@ -70,7 +71,7 @@ pub(crate) struct ProcedureMeta {
     /// Notify to wait for subprocedures.
     child_notify: Notify,
     /// Lock required by this procedure.
-    lock_key: Option<LockKey>,
+    lock_key: LockKey,
     /// Mutable status during execution.
     exec_meta: Mutex<ExecMeta>,
 }
@@ -128,6 +129,13 @@ pub(crate) struct ManagerContext {
     messages: Mutex<HashMap<ProcedureId, ProcedureMessage>>,
 }
 
+#[async_trait]
+impl ContextProvider for ManagerContext {
+    async fn procedure_state(&self, procedure_id: ProcedureId) -> Result<Option<ProcedureState>> {
+        Ok(self.state(procedure_id))
+    }
+}
+
 impl ManagerContext {
     /// Returns a new [ManagerContext].
     fn new() -> ManagerContext {
@@ -409,7 +417,7 @@ mod test_util {
             lock_notify: Notify::new(),
             parent_id: None,
             child_notify: Notify::new(),
-            lock_key: None,
+            lock_key: LockKey::default(),
             exec_meta: Mutex::new(ExecMeta::default()),
         }
     }
@@ -508,8 +516,8 @@ mod tests {
             Ok(self.content.clone())
         }
 
-        fn lock_key(&self) -> Option<LockKey> {
-            None
+        fn lock_key(&self) -> LockKey {
+            LockKey::default()
         }
     }
 
@@ -617,8 +625,8 @@ mod tests {
                 unimplemented!()
             }
 
-            fn lock_key(&self) -> Option<LockKey> {
-                Some(LockKey::new("test.submit"))
+            fn lock_key(&self) -> LockKey {
+                LockKey::single("test.submit")
             }
         }