tunnel cli: add --parent-process-id option (microsoft#163935)

* tunnel cli: add --parent-process-id option

* Update cli/src/commands/args.rs

Co-authored-by: Connor Peet <[email protected]>

* Change parent process check interval to 2s

Co-authored-by: Connor Peet <[email protected]>

Author: aeschli

cli/src/commands/args.rs

@@ -559,6 +559,10 @@ pub struct TunnelServeArgs {
 	/// Randomly name machine for port forwarding service
 	#[clap(long)]
 	pub random_name: bool,
+
+	/// Optional parent process id. If provided, the server will be stopped when the process of the given pid no longer exists
+	#[clap(long, hide = true)]
+	pub parent_process_id: Option<i32>,
 }
 
 #[derive(Args, Debug, Clone)]

cli/src/commands/tunnels.rs

@@ -3,8 +3,11 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+use std::fmt;
 use async_trait::async_trait;
-use tokio::sync::oneshot;
+use sysinfo::{Pid, SystemExt};
+use tokio::sync::mpsc;
+use tokio::time::{sleep, Duration};
 
 use super::{
 	args::{
@@ -70,7 +73,7 @@ impl ServiceContainer for TunnelServiceContainer {
 		&mut self,
 		log: log::Logger,
 		launcher_paths: LauncherPaths,
-		shutdown_rx: oneshot::Receiver<()>,
+		shutdown_rx: mpsc::Receiver<ShutdownSignal>,
 	) -> Result<(), AnyError> {
 		let csa = (&self.args).into();
 		serve_with_csa(
@@ -87,6 +90,20 @@ impl ServiceContainer for TunnelServiceContainer {
 		Ok(())
 	}
 }
+/// Describes the signal to manully stop the server
+pub enum ShutdownSignal {
+	CtrlC,
+	ParentProcessKilled,
+}
+
+impl fmt::Display for ShutdownSignal {
+	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+		match self {
+			ShutdownSignal::CtrlC => write!(f, "Ctrl-C received"),
+			ShutdownSignal::ParentProcessKilled => write!(f, "Parent process no longer exists"),
+		}
+	}
+}
 
 pub async fn service(
 	ctx: CommandContext,
@@ -211,7 +228,7 @@ async fn serve_with_csa(
 	log: Logger,
 	gateway_args: TunnelServeArgs,
 	csa: CodeServerArgs,
-	shutdown_rx: Option<oneshot::Receiver<()>>,
+	shutdown_rx: Option<mpsc::Receiver<ShutdownSignal>>,
 ) -> Result<i32, AnyError> {
 	// Intentionally read before starting the server. If the server updated and
 	// respawn is requested, the old binary will get renamed, and then
@@ -230,10 +247,22 @@ async fn serve_with_csa(
 	let shutdown_tx = if let Some(tx) = shutdown_rx {
 		tx
 	} else {
-		let (tx, rx) = oneshot::channel();
+		let (tx, rx) = mpsc::channel::<ShutdownSignal>(2);
+		if let Some(process_id) = gateway_args.parent_process_id {
+			let tx = tx.clone();
+			info!(log, "checking for parent process {}", process_id);
+			tokio::spawn(async move {
+				let mut s = sysinfo::System::new();
+				let pid = Pid::from(process_id);
+				while s.refresh_process(pid) {
+					sleep(Duration::from_millis(2000)).await;
+				}
+				tx.send(ShutdownSignal::ParentProcessKilled).await.ok();
+			});
+		}
 		tokio::spawn(async move {
 			tokio::signal::ctrl_c().await.ok();
-			tx.send(()).ok();
+			tx.send(ShutdownSignal::CtrlC).await.ok();
 		});
 		rx
 	};

cli/src/tunnels/control_server.rs

@@ -2,6 +2,7 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
+use crate::commands::tunnels::ShutdownSignal;
 use crate::constants::{CONTROL_PORT, PROTOCOL_VERSION, VSCODE_CLI_VERSION};
 use crate::log;
 use crate::self_update::SelfUpdate;
@@ -23,7 +24,7 @@ use std::sync::Arc;
 use std::time::Instant;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
 use tokio::pin;
-use tokio::sync::{mpsc, oneshot, Mutex};
+use tokio::sync::{mpsc, Mutex};
 
 use super::code_server::{
 	AnyCodeServer, CodeServerArgs, ServerBuilder, ServerParamsRaw, SocketCodeServer,
@@ -181,7 +182,7 @@ pub async fn serve(
 	launcher_paths: &LauncherPaths,
 	code_server_args: &CodeServerArgs,
 	platform: Platform,
-	shutdown_rx: oneshot::Receiver<()>,
+	shutdown_rx: mpsc::Receiver<ShutdownSignal>,
 ) -> Result<ServerTermination, AnyError> {
 	let mut port = tunnel.add_port_direct(CONTROL_PORT).await?;
 	print_listening(log, &tunnel.name);
@@ -194,8 +195,8 @@ pub async fn serve(
 
 	loop {
 		tokio::select! {
-			_ = &mut shutdown_rx => {
-				info!(log, "Received interrupt, shutting down...");
+			Some(r) = shutdown_rx.recv() => {
+				info!(log, "Shutting down: {}", r );
 				drop(signal_exit);
 				return Ok(ServerTermination {
 					respawn: false,

cli/src/tunnels/service.rs

@@ -6,8 +6,9 @@
 use std::path::PathBuf;
 
 use async_trait::async_trait;
-use tokio::sync::oneshot;
+use tokio::sync::mpsc;
 
+use crate::commands::tunnels::ShutdownSignal;
 use crate::log;
 use crate::state::LauncherPaths;
 use crate::util::errors::AnyError;
@@ -20,7 +21,7 @@ pub trait ServiceContainer: Send {
 		&mut self,
 		log: log::Logger,
 		launcher_paths: LauncherPaths,
-		shutdown_rx: oneshot::Receiver<()>,
+		shutdown_rx: mpsc::Receiver<ShutdownSignal>,
 	) -> Result<(), AnyError>;
 }
 

src/vs/platform/remoteTunnel/electron-browser/remoteTunnelService.ts

@@ -145,7 +145,7 @@ export class RemoteTunnelService extends Disposable implements IRemoteTunnelServ
 				hostName = undefined;
 			}
 		}
-		const args = hostName ? [] : ['--random-name'];
+		const args = hostName ? [] : ['--random-name', '--parent-process-id', String(process.pid)];
 		const serveCommand = this.runCodeTunneCommand('tunnel', args, (message: string) => {
 		});
 		this._tunnelProcess = serveCommand;