Cleanup running containers on the Control-C signal #422

ammernico · 2024-09-25T16:33:30Z

Add the signal feature to tokio to interrupt and handle the Control-C signal in Butido.
Add Control-C signal handling into the Orchestrator.
Implement Drop on the JobHandle to ensure container cleanup.

This is a working draft pr for testing purposes and still missing some features.

christophprokop · 2024-09-26T14:10:55Z

Tested while multiple containers were running on all build hosts.
All containeres were successfully stopped after just a few seconds.
Nice! :)

- Add the `signal` feature to `tokio` to interrupt and handle the Control-C signal in Butido. - Add Control-C signal handling into the `Orchestrator`. - Implement `Drop` on the `JobHandle` to ensure container cleanup. Fixes science-computing#409 Signed-off-by: Nico Steinle <[email protected]>

primeos-work

I haven't done any testing yet but already found some potential issues in the code.

primeos-work · 2024-10-14T18:58:43Z

src/endpoint/scheduler.rs

-        drop(self.bar);
+        drop(self.bar.clone());


Huh, we're cloning and immediately dropping the clone? Shouldn't this be a no-op? I'm a bit surprised that Clippy doesn't catch this.

(By only looking at this context I'm surprised why there is a self.bar related change at all...)

primeos-work · 2024-10-14T19:02:30Z

src/endpoint/scheduler.rs

@@ -370,6 +374,36 @@ impl JobHandle {
    }
 }

+impl Drop for JobHandle {
+    fn drop(&mut self) {
+        debug!("Cleaning up JobHandle");


Nit: Might be nice/useful to include the job ID (but debugging is obviously optional).

primeos-work · 2024-10-14T19:04:50Z

src/endpoint/scheduler.rs

+        if self.container_id.is_some() {
+            debug!("Container was already started");
+            let docker = self.endpoint.docker().clone();
+            let container_id = self.container_id.take().unwrap();


This isn't pretty and not guaranteed to be safe - please use if let Some(container_id) = self.container_id or something similar instead.

primeos-work · 2024-10-14T19:09:54Z

src/endpoint/scheduler.rs

+                let container_info = container.inspect().await.unwrap();
+
+                if container_info.state.running {
+                    debug!("Container is still running, cleaning up...");


Nit: Including the container ID would be nice.

primeos-work · 2024-10-14T19:17:45Z

src/endpoint/scheduler.rs

+
+            tokio::spawn(async move {
+                let container = docker.containers().get(&container_id);
+                let container_info = container.inspect().await.unwrap();


We should avoid the unwrap here - I'd probably just log the error (in theory we might occasionally run into such errors when the containers terminate between the if and this inspect).

PS: We want to avoid unwrap() as much as possible in general (but there are of course exceptions where it's fine).

primeos-work · 2024-10-14T19:35:14Z

src/orchestrator/orchestrator.rs

@@ -265,12 +267,26 @@ impl Borrow<ArtifactPath> for ProducedArtifact {

 impl<'a> Orchestrator<'a> {
    pub async fn run(self, output: &mut Vec<ArtifactPath>) -> Result<HashMap<Uuid, Error>> {
-        let (results, errors) = self.run_tree().await?;
+        let token = CancellationToken::new();
+        let cloned_token = token.clone();


I'd just clone it where needed instead - why use a variable here?

primeos-work · 2024-10-14T19:36:20Z

src/orchestrator/orchestrator.rs

+        let cloned_token = token.clone();
+
+        tokio::spawn(async move {
+            info!("Received the ctl-c signal, stopping...");


Looking at the code I would assume that this always gets printed - shouldn't you move this below tokio::signal::ctrl_c().await?

primeos-work · 2024-10-14T19:43:21Z

src/orchestrator/orchestrator.rs

-        let running_jobs = jobs
-            .into_iter()
-            .map(|prep| {
-                trace!(parent: &run_span, job_uuid = %prep.1.jobdef.job.uuid(), "Creating JobTask");
-                // the sender is set or we need to use the root sender
-                let sender = prep
-                    .3
-                    .into_inner()
-                    .unwrap_or_else(|| vec![root_sender.clone()]);
-                JobTask::new(prep.0, prep.1, sender)
-            })
-            .inspect(
-                |task| trace!(parent: &run_span, job_uuid = %task.jobdef.job.uuid(), "Running job"),
-            )
-            .map(|task| {
-                task.run()
-                    .instrument(tracing::debug_span!(parent: &run_span, "JobTask::run"))
-            })
-            .collect::<futures::stream::FuturesUnordered<_>>();
-        debug!("Built {} jobs", running_jobs.len());
-
-        running_jobs
-            .collect::<Result<()>>()
-            .instrument(run_span.clone())
-            .await?;
-        trace!(parent: &run_span, "All jobs finished");
-        drop(run_span);
-
-        match root_receiver.recv().await {
-            None => Err(anyhow!("No result received...")),
-            Some(Ok(results)) => {
-                let results = results
-                    .into_iter()
-                    .flat_map(|tpl| tpl.1.into_iter())
-                    .map(ProducedArtifact::unpack)
-                    .collect();
-                Ok((results, HashMap::with_capacity(0)))
+
+        tokio::select! {
+            _ = token.cancelled() => {
+                anyhow::bail!("Received Control-C signal");
+            }
+            r = async {
+                let running_jobs = jobs
+                        .into_iter()
+                        .map(|prep| {
+                            trace!(parent: &run_span, job_uuid = %prep.1.jobdef.job.uuid(), "Creating JobTask");
+                            // the sender is set or we need to use the root sender
+                            let sender = prep
+                                .3
+                                .into_inner()
+                                .unwrap_or_else(|| vec![root_sender.clone()]);
+                            JobTask::new(prep.0, prep.1, sender)
+                        })
+                        .inspect(
+                            |task| trace!(parent: &run_span, job_uuid = %task.jobdef.job.uuid(), "Running job"),
+                        )
+                        .map(|task| {
+                            task.run()
+                                .instrument(tracing::debug_span!(parent: &run_span, "JobTask::run"))
+                        })
+                        .collect::<futures::stream::FuturesUnordered<_>>();
+                debug!("Built {} jobs", running_jobs.len());
+
+                running_jobs
+                    .collect::<Result<()>>()
+                    .instrument(run_span.clone())
+                    .await?;
+                trace!(parent: &run_span, "All jobs finished");
+                drop(run_span);
+
+                match root_receiver.recv().await {
+                    None => Err(anyhow!("No result received...")),
+                    Some(Ok(results)) => {
+                        let results = results
+                            .into_iter()
+                            .flat_map(|tpl| tpl.1.into_iter())
+                            .map(ProducedArtifact::unpack)
+                            .collect();
+                        Ok((results, HashMap::with_capacity(0)))
+                    }
+                    Some(Err(errors)) => Ok((vec![], errors)),
+                }
+            } => {
+                r
            }
-            Some(Err(errors)) => Ok((vec![], errors)),


Note to self: I still need to properly review this part.

primeos-work · 2024-10-14T19:46:25Z

src/endpoint/scheduler.rs

 pub struct JobHandle {
    log_dir: Option<PathBuf>,
    endpoint: EndpointHandle,
+    container_id: Option<String>,


This approach seems fine but we never set container_id back to None for finished jobs - it would be more elegant if we could do so (after ensuring that the container has indeed exited but that might already be implemented to check if the job has finished).

primeos-work · 2024-10-14T19:49:08Z

src/endpoint/scheduler.rs

@@ -202,12 +206,12 @@ impl JobHandle {
            package_name: &package.name,
            package_version: &package.version,
            log_dir: self.log_dir.as_ref(),
-            job: self.job,
+            job: self.job.clone(),


TODO/question: Why do we now need to clone here? (The whole cloning is irritating me a bit in general and can be quite dangerous if the data can diverge)

ammernico requested review from christophprokop and primeos-work September 25, 2024 16:33

christophprokop added the prerelease PRs which are merged to staging branch but not in main/master label Oct 1, 2024

primeos-work linked an issue Oct 7, 2024 that may be closed by this pull request

cleanup job for running containers #409

Open

ammernico force-pushed the container-cleanup branch from 54b622f to 770ffb1 Compare October 14, 2024 09:29

ammernico marked this pull request as ready for review October 14, 2024 09:42

ammernico force-pushed the container-cleanup branch from 770ffb1 to 284f396 Compare October 14, 2024 09:42

ammernico force-pushed the container-cleanup branch from 284f396 to 752d2fc Compare October 14, 2024 13:17

primeos-work requested changes Oct 14, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Cleanup running containers on the Control-C signal #422

Cleanup running containers on the Control-C signal #422

ammernico commented Sep 25, 2024

christophprokop commented Sep 26, 2024

primeos-work left a comment

primeos-work Oct 14, 2024

primeos-work Oct 14, 2024

primeos-work Oct 14, 2024

primeos-work Oct 14, 2024

primeos-work Oct 14, 2024

primeos-work Oct 14, 2024

primeos-work Oct 14, 2024

primeos-work Oct 14, 2024

primeos-work Oct 14, 2024

primeos-work Oct 14, 2024

Cleanup running containers on the Control-C signal #422

Are you sure you want to change the base?

Cleanup running containers on the Control-C signal #422

Conversation

ammernico commented Sep 25, 2024

christophprokop commented Sep 26, 2024

primeos-work left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment