Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 51 additions & 5 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,5 +16,6 @@ sysinfo = "0.31.4"
log = "0.4.0"
env_logger = "0.11.8"
regex = "1.11.1"
ctrlc = "3.4"
core-foundation = "0.10.1"
objc2-io-kit = "0.3.2"
40 changes: 40 additions & 0 deletions src/cleanup.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
// SPDX-License-Identifier: Apache-2.0

use std::path::PathBuf;
use std::sync::Arc;

/// Remove each socket file, logging the outcome. Missing files are silently ignored.
fn remove_sockets(paths: &[PathBuf]) {
for path in paths {
match std::fs::remove_file(path) {
Ok(()) => log::info!("removed vsock socket: {}", path.display()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
Err(e) => log::warn!("failed to remove vsock socket {}: {}", path.display(), e),
}
}
}

/// Guard that cleans up vsock socket files when dropped (normal exit) and
/// also installs a Ctrl-C handler so the sockets are removed on SIGINT.
pub struct VsockCleanupGuard {
paths: Arc<Vec<PathBuf>>,
}

impl VsockCleanupGuard {
/// Create the guard **and** register the SIGINT handler.
pub fn new(paths: Vec<PathBuf>) -> Result<Self, anyhow::Error> {
let paths = Arc::new(paths);
let handler_paths = Arc::clone(&paths);
ctrlc::set_handler(move || {
remove_sockets(&handler_paths);
std::process::exit(0);
})?;
Ok(Self { paths })
}
}

impl Drop for VsockCleanupGuard {
fn drop(&mut self) {
remove_sockets(&self.paths);
}
}
19 changes: 18 additions & 1 deletion src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use std::{
};

use crate::timesync::timesync_listener;
use crate::virtio::{VsockAction, VsockConfig};
use crate::virtio::{VirtioDeviceConfig, VsockAction, VsockConfig};
use anyhow::{anyhow, Context};
use env_logger::{Builder, Env, Target};

Expand Down Expand Up @@ -217,6 +217,23 @@ impl TryFrom<Args> for KrunContext {
}

impl KrunContext {
/// Collect all vsock socket paths that should be cleaned up on exit.
pub fn vsock_socket_paths(&self) -> Vec<PathBuf> {
let mut paths = Vec::new();
for device in &self.args.devices {
if let VirtioDeviceConfig::Vsock(vsock) = device {
paths.push(vsock.socket_url.clone());
}
}
if self.args.timesync.is_some() {
paths.push(PathBuf::from(format!(
"/tmp/krunkit_timesync_{}.sock",
std::process::id()
)));
}
paths
}
Comment on lines +220 to +235
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic for generating the timesync socket path is duplicated here and in the TryFrom<Args> for KrunContext implementation (lines 205-208). Refactoring this into a helper method improves maintainability and ensures consistency across the codebase. Please also update the usage in the TryFrom implementation to use this new method.

    fn timesync_socket_path() -> PathBuf {
        PathBuf::from(format!(
            "/tmp/krunkit_timesync_{}.sock",
            std::process::id()
        ))
    }

    /// Collect all vsock socket paths that should be cleaned up on exit.
    pub fn vsock_socket_paths(&self) -> Vec<PathBuf> {
        let mut paths = Vec::new();
        for device in &self.args.devices {
            if let VirtioDeviceConfig::Vsock(vsock) = device {
                paths.push(vsock.socket_url.clone());
            }
        }
        if self.args.timesync.is_some() {
            paths.push(Self::timesync_socket_path());
        }
        paths
    }


/// Spawn a thread to listen for shutdown requests and run the workload. If behaving properly,
/// the main thread will never return from this function.
pub fn run(&self) -> Result<(), anyhow::Error> {
Expand Down
5 changes: 5 additions & 0 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,14 @@

#![allow(dead_code)]

mod cleanup;
mod cmdline;
mod context;
mod status;
mod timesync;
mod virtio;

use cleanup::VsockCleanupGuard;
use cmdline::Args;
use context::KrunContext;

Expand All @@ -18,6 +20,9 @@ fn main() -> Result<(), anyhow::Error> {
// accordingly.
let ctx = KrunContext::try_from(Args::parse())?;

// The guard removes vsock socket files on drop (normal exit) and on Ctrl-C.
let _cleanup = VsockCleanupGuard::new(ctx.vsock_socket_paths())?;

// Run the workload. If behaving properly, the main thread will not return from this
// function.
ctx.run()?;
Expand Down