Merge pull request #94 from guzman-raphael/agent_store_pod_result

`Agent`: Add configurable storage of `PodResult`

Author: Synicix

cspell.json

@@ -18,6 +18,7 @@
         "cffi",
         "zenoh",
         "PodJob",
+        "stresser"
     ],
     "ignoreWords": [
         "relpath",

src/core/orchestrator/agent.rs

@@ -123,8 +123,9 @@ pub async fn start_service<
 ) -> Result<()>
 where
     EventClassifierF: Fn(&RequestI) -> EventPayload + Send + 'static,
-    RequestI: for<'serde> Deserialize<'serde>,
-    RequestF: Fn(Arc<Agent>, HashMap<String, PathBuf>, EventMetadata, RequestI) -> RequestR
+    RequestI: for<'serde> Deserialize<'serde> + Send + 'static,
+    RequestF: FnOnce(Arc<Agent>, HashMap<String, PathBuf>, EventMetadata, RequestI) -> RequestR
+        + Clone
         + Send
         + 'static,
     RequestR: Future<Output = Result<ResponseI>> + Send + 'static,
@@ -164,19 +165,25 @@ where
                         subgroup: metadata["pod_job_hash"].to_string(),
                     };
                     let _event_payload = event_classifier(&input);
-                    tasks.spawn(
-                        request_task(
-                            Arc::clone(&inner_agent),
-                            namespace_lookup.clone(),
-                            event_metadata,
-                            input,
-                        )
-                        .then(move |response| async move {
-                            let _: Result<(), SendError<Result<ResponseI>>> =
-                                inner_response_tx.send(response).await;
-                            Ok::<_, OrcaError>(())
-                        }),
-                    );
+                    tasks.spawn({
+                        let inner_request_task = request_task.clone();
+                        let inner_inner_agent = Arc::clone(&inner_agent);
+                        let inner_namespace_lookup = namespace_lookup.clone();
+                        async move {
+                            inner_request_task(
+                                inner_inner_agent,
+                                inner_namespace_lookup,
+                                event_metadata,
+                                input,
+                            )
+                            .then(move |response| async move {
+                                let _: Result<(), SendError<Result<ResponseI>>> =
+                                    inner_response_tx.send(response).await;
+                                Ok::<_, OrcaError>(())
+                            })
+                            .await
+                        }
+                    });
                 }
             }
             Ok(())

src/core/orchestrator/docker.rs

@@ -170,6 +170,7 @@ impl LocalDockerOrchestrator {
         clippy::cast_precision_loss,
         clippy::cast_possible_truncation,
         clippy::indexing_slicing,
+        clippy::too_many_lines,
         reason = r#"
         - Timestamp and memory should always have a value > 0
         - Container will always have a name with more than 1 character
@@ -210,13 +211,13 @@ impl LocalDockerOrchestrator {
             let terminated_timestamp =
                 DateTime::parse_from_rfc3339(container_spec.state.as_ref()?.finished_at.as_ref()?)
                     .ok()?
-                    .timestamp() as u64;
+                    .timestamp();
             Some((
                 container_name,
                 RunInfo {
                     image: container_spec.config.as_ref()?.image.as_ref()?.clone(),
                     created: container_summary.created? as u64,
-                    terminated: (terminated_timestamp > 0).then_some(terminated_timestamp),
+                    terminated: (terminated_timestamp > 0).then_some(terminated_timestamp as u64),
                     env_vars: container_spec
                         .config
                         .as_ref()?
@@ -243,9 +244,20 @@ impl LocalDockerOrchestrator {
                         container_spec.state.as_ref()?.exit_code? as i16,
                     ) {
                         (ContainerStateStatusEnum::RUNNING, _) => Status::Running,
-                        (ContainerStateStatusEnum::EXITED, 0) => Status::Completed,
-                        (ContainerStateStatusEnum::EXITED, code) => Status::Failed(code),
-                        _ => todo!(),
+                        (
+                            ContainerStateStatusEnum::EXITED | ContainerStateStatusEnum::REMOVING,
+                            0,
+                        ) => Status::Completed,
+                        (
+                            ContainerStateStatusEnum::EXITED | ContainerStateStatusEnum::REMOVING,
+                            code,
+                        ) => Status::Failed(code),
+                        (_, code) => {
+                            todo!(
+                                "Unhandled container state: {}, exit code: {code}.",
+                                container_spec.state.as_ref()?.status.as_ref()?
+                            )
+                        }
                     },
                     mounts: container_spec
                         .mounts

src/uniffi/orchestrator/agent.rs

@@ -2,8 +2,9 @@ use crate::{
     core::orchestrator::agent::{EventPayload, start_service},
     uniffi::{
         error::{OrcaError, Result, selector},
-        model::PodJob,
+        model::{PodJob, PodResult},
         orchestrator::{Orchestrator, Status, docker::LocalDockerOrchestrator},
+        store::{Store as _, filestore::LocalFileStore},
     },
 };
 use derive_more::Display;
@@ -17,6 +18,19 @@ use tokio::task::JoinSet;
 use uniffi;
 use zenoh;
 
+/// A response, similar to Rust's `Result` but casting error to `String`.
+///
+/// This is a workaround due to `UniFFI` limitations when trying to send a collection of `Result`
+/// over the CFFI boundary e.g. concurrent calls where it is necessary to know the status of each
+/// request to determine what to retry.
+#[derive(uniffi::Enum)]
+pub enum Response {
+    /// Success
+    Ok,
+    /// Error cast to `String`
+    Err(String),
+}
+
 /// Client to connect to an execution agent within a coordinated fleet. Connection optimized/rerouted by Zenoh.
 #[expect(
     clippy::field_scoped_visibility_modifiers,
@@ -58,14 +72,14 @@ impl AgentClient {
     }
     /// Submit many pod jobs to be processed in parallel.
     /// Return order will match inputs, casting outputs to `String` (since `uniffi` doesn't support sending unwrapped `Result`s).
-    pub async fn submit_pod_jobs(&self, pod_jobs: Vec<Arc<PodJob>>) -> Vec<String> {
+    pub async fn submit_pod_jobs(&self, pod_jobs: Vec<Arc<PodJob>>) -> Vec<Response> {
         join_all(pod_jobs.iter().map(|pod_job| async {
             match self
                 .publish(&format!("request/pod_job/{}", pod_job.hash), pod_job)
                 .await
             {
-                Ok(()) => "ok".into(),
-                Err(error) => error.to_string(),
+                Ok(()) => Response::Ok,
+                Err(error) => Response::Err(error.to_string()),
             }
         }))
         .await
@@ -126,13 +140,18 @@ impl Agent {
     /// # Errors
     ///
     /// Will stop and return an error if encounters an error while processing any pod job request.
-    pub async fn start(&self, namespace_lookup: &HashMap<String, PathBuf>) -> Result<()> {
+    #[expect(clippy::excessive_nesting, reason = "Nesting manageable.")]
+    pub async fn start(
+        &self,
+        namespace_lookup: &HashMap<String, PathBuf>,
+        available_store: Option<Arc<LocalFileStore>>,
+    ) -> Result<()> {
         let mut services = JoinSet::new();
         services.spawn(start_service(
             Arc::new(self.clone()),
             "request/pod_job/**".to_owned(),
             namespace_lookup.clone(),
-            |input: &PodJob| EventPayload::Request(input.clone()),
+            |pod_job: &PodJob| EventPayload::Request(pod_job.clone()),
             async |agent, inner_namespace_lookup, _, pod_job| {
                 let pod_run = agent
                     .orchestrator
@@ -152,6 +171,33 @@ impl Agent {
                 client.publish(response_topic, &pod_result).await
             },
         ));
+        if let Some(store) = available_store {
+            services.spawn(start_service(
+                Arc::new(self.clone()),
+                "success/pod_job/**".to_owned(),
+                namespace_lookup.clone(),
+                |pod_result: &PodResult| EventPayload::Success(pod_result.clone()),
+                {
+                    let inner_store = Arc::clone(&store);
+                    async move |_, _, _, pod_result| {
+                        inner_store.save_pod_result(&pod_result)?;
+                        Ok(())
+                    }
+                },
+                async |_, ()| Ok(()),
+            ));
+            services.spawn(start_service(
+                Arc::new(self.clone()),
+                "failure/pod_job/**".to_owned(),
+                namespace_lookup.clone(),
+                |pod_result: &PodResult| EventPayload::Failure(pod_result.clone()),
+                async move |_, _, _, pod_result| {
+                    store.save_pod_result(&pod_result)?;
+                    Ok(())
+                },
+                async |_, ()| Ok(()),
+            ));
+        }
         services
             .join_next()
             .await

src/uniffi/orchestrator/docker.rs

@@ -6,20 +6,21 @@ use crate::{
     uniffi::{
         error::{OrcaError, Result, selector},
         model::{PodJob, PodResult},
-        orchestrator::{ImageKind, Orchestrator, PodRun, RunInfo},
+        orchestrator::{ImageKind, Orchestrator, PodRun, RunInfo, Status},
     },
 };
 use async_trait;
 use bollard::{
     Docker,
     container::{RemoveContainerOptions, StartContainerOptions, WaitContainerOptions},
+    errors::Error::DockerContainerWaitError,
     image::{CreateImageOptions, ImportImageOptions},
 };
 use derive_more::Display;
 use futures_util::stream::{StreamExt as _, TryStreamExt as _};
 use snafu::{OptionExt as _, futures::TryFutureExt as _};
-use std::{collections::HashMap, path::PathBuf, sync::Arc};
-use tokio::fs::File;
+use std::{collections::HashMap, path::PathBuf, sync::Arc, time::Duration};
+use tokio::{fs::File, time::sleep as async_sleep};
 use tokio_util::{
     bytes::{Bytes, BytesMut},
     codec::{BytesCodec, FramedRead},
@@ -192,12 +193,32 @@ impl Orchestrator for LocalDockerOrchestrator {
             })?;
         Ok(run_info)
     }
+    #[expect(
+        clippy::wildcard_enum_match_arm,
+        reason = "Favor readability due to complexity in external dependency."
+    )]
     async fn get_result(&self, pod_run: &PodRun) -> Result<PodResult> {
-        self.api
+        match self
+            .api
             .wait_container(&pod_run.assigned_name, None::<WaitContainerOptions<String>>)
             .try_collect::<Vec<_>>()
-            .await?;
-        let result_info = self.get_info(pod_run).await?;
+            .await
+        {
+            Ok(_) => (),
+            Err(err) => match err {
+                DockerContainerWaitError { .. } => (),
+                _ => return Err(OrcaError::from(err)),
+            },
+        }
+
+        let mut result_info: RunInfo;
+        while {
+            result_info = self.get_info(pod_run).await?;
+            matches!(&result_info.status, Status::Running)
+        } {
+            async_sleep(Duration::from_millis(100)).await;
+        }
+
         PodResult::new(
             None,
             Arc::clone(&pod_run.pod_job),

src/uniffi/store/filestore.rs

@@ -8,7 +8,7 @@ use getset::CloneGetters;
 use std::{fs, path::PathBuf};
 use uniffi;
 /// Support for a storage backend on a local filesystem directory.
-#[derive(uniffi::Object, Debug, Display, CloneGetters)]
+#[derive(uniffi::Object, Debug, Display, CloneGetters, Clone)]
 #[getset(get_clone, impl_attrs = "#[uniffi::export]")]
 #[display("{self:#?}")]
 #[uniffi::export(Display)]