Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

stop tunnel when executable gets deleted #181505

Merged
merged 1 commit into from May 9, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
24 changes: 13 additions & 11 deletions cli/src/commands/tunnels.rs
Expand Up @@ -340,21 +340,23 @@ async fn serve_with_csa(
log = log.tee(log_broadcast.clone());
log::install_global_logger(log.clone()); // re-install so that library logs are captured

let shutdown = match gateway_args
.parent_process_id
.and_then(|p| Pid::from_str(&p).ok())
{
Some(pid) => ShutdownRequest::create_rx([
ShutdownRequest::CtrlC,
ShutdownRequest::ParentProcessKilled(pid),
]),
None => ShutdownRequest::create_rx([ShutdownRequest::CtrlC]),
};

// Intentionally read before starting the server. If the server updated and
// respawn is requested, the old binary will get renamed, and then
// current_exe will point to the wrong path.
let current_exe = std::env::current_exe().unwrap();

let mut vec = vec![
ShutdownRequest::CtrlC,
ShutdownRequest::ExeUninstalled(current_exe.to_owned()),
];
if let Some(p) = gateway_args
.parent_process_id
.and_then(|p| Pid::from_str(&p).ok())
{
vec.push(ShutdownRequest::ParentProcessKilled(p));
}
let shutdown = ShutdownRequest::create_rx(vec);

let server = loop {
if shutdown.is_open() {
return Ok(0);
Expand Down
13 changes: 11 additions & 2 deletions cli/src/tunnels/shutdown_signal.rs
Expand Up @@ -4,11 +4,11 @@
*--------------------------------------------------------------------------------------------*/

use futures::{stream::FuturesUnordered, StreamExt};
use std::fmt;
use std::{fmt, path::PathBuf};
use sysinfo::Pid;

use crate::util::{
machine::wait_until_process_exits,
machine::{wait_until_exe_deleted, wait_until_process_exits},
sync::{new_barrier, Barrier, Receivable},
};

Expand All @@ -17,6 +17,7 @@ use crate::util::{
pub enum ShutdownSignal {
CtrlC,
ParentProcessKilled(Pid),
ExeUninstalled,
ServiceStopped,
RpcShutdownRequested,
RpcRestartRequested,
Expand All @@ -29,6 +30,9 @@ impl fmt::Display for ShutdownSignal {
ShutdownSignal::ParentProcessKilled(p) => {
write!(f, "Parent process {} no longer exists", p)
}
ShutdownSignal::ExeUninstalled => {
write!(f, "Executable no longer exists")
}
ShutdownSignal::ServiceStopped => write!(f, "Service stopped"),
ShutdownSignal::RpcShutdownRequested => write!(f, "RPC client requested shutdown"),
ShutdownSignal::RpcRestartRequested => {
Expand All @@ -41,6 +45,7 @@ impl fmt::Display for ShutdownSignal {
pub enum ShutdownRequest {
CtrlC,
ParentProcessKilled(Pid),
ExeUninstalled(PathBuf),
Derived(Box<dyn Receivable<ShutdownSignal> + Send>),
}

Expand All @@ -56,6 +61,10 @@ impl ShutdownRequest {
wait_until_process_exits(pid, 2000).await;
Some(ShutdownSignal::ParentProcessKilled(pid))
}
ShutdownRequest::ExeUninstalled(exe_path) => {
wait_until_exe_deleted(&exe_path, 2000).await;
Some(ShutdownSignal::ExeUninstalled)
}
ShutdownRequest::Derived(mut rx) => rx.recv_msg().await,
}
}
Expand Down
7 changes: 7 additions & 0 deletions cli/src/util/machine.rs
Expand Up @@ -52,3 +52,10 @@ pub fn find_running_process(name: &Path) -> Option<u32> {
}
None
}

pub async fn wait_until_exe_deleted(current_exe: &Path, poll_ms: u64) {
let duration = Duration::from_millis(poll_ms);
while current_exe.exists() {
tokio::time::sleep(duration).await;
}
}