cli: shut down service on windows more reliably (#181584)

* cli: shut down service on windows more reliably

Use the singleton kill logic.

Fixes #175268

* fix lint

cli/src/commands/tunnels.rs

@@ -7,7 +7,6 @@ use async_trait::async_trait;
 use sha2::{Digest, Sha256};
 use std::{str::FromStr, time::Duration};
 use sysinfo::Pid;
-use tokio::sync::mpsc;
 
 use super::{
 	args::{
@@ -18,12 +17,9 @@ use super::{
 };
 
 use crate::{
-	async_pipe::socket_stream_split,
 	auth::Auth,
 	constants::{APPLICATION_NAME, TUNNEL_CLI_LOCK_NAME, TUNNEL_SERVICE_LOCK_NAME},
-	json_rpc::{new_json_rpc, start_json_rpc},
 	log,
-	singleton::connect_as_client,
 	state::LauncherPaths,
 	tunnels::{
 		code_server::CodeServerArgs,
@@ -31,6 +27,7 @@ use crate::{
 		paths::get_all_servers,
 		protocol,
 		shutdown_signal::ShutdownRequest,
+		singleton_client::do_single_rpc_call,
 		singleton_server::{
 			make_singleton_server, start_singleton_server, BroadcastLogSink, SingletonServerArgs,
 		},
@@ -38,7 +35,7 @@ use crate::{
 	},
 	util::{
 		app_lock::AppMutex,
-		errors::{wrap, AnyError, CodeError},
+		errors::{wrap, AnyError},
 		prereqs::PreReqChecker,
 	},
 };
@@ -206,69 +203,34 @@ pub async fn unregister(ctx: CommandContext) -> Result<i32, AnyError> {
 	Ok(0)
 }
 
-async fn do_single_rpc_call<
-	P: serde::Serialize,
-	R: serde::de::DeserializeOwned + Send + 'static,
->(
-	ctx: &CommandContext,
-	method: &'static str,
-	params: P,
-) -> Result<R, AnyError> {
-	let client = match connect_as_client(&ctx.paths.tunnel_lockfile()).await {
-		Ok(p) => p,
-		Err(CodeError::SingletonLockfileOpenFailed(_))
-		| Err(CodeError::SingletonLockedProcessExited(_)) => {
-			return Err(CodeError::NoRunningTunnel.into());
-		}
-		Err(e) => return Err(e.into()),
-	};
-
-	let (msg_tx, msg_rx) = mpsc::unbounded_channel();
-	let mut rpc = new_json_rpc();
-	let caller = rpc.get_caller(msg_tx);
-	let (read, write) = socket_stream_split(client);
-	let log = ctx.log.clone();
-
-	let rpc = tokio::spawn(async move {
-		start_json_rpc(
-			rpc.methods(()).build(log),
-			read,
-			write,
-			msg_rx,
-			ShutdownRequest::create_rx([ShutdownRequest::CtrlC]),
-		)
-		.await
-		.unwrap();
-	});
-
-	let r = caller.call(method, params).await.unwrap();
-	rpc.abort();
-	r.map_err(|err| CodeError::TunnelRpcCallFailed(err).into())
-}
-
 pub async fn restart(ctx: CommandContext) -> Result<i32, AnyError> {
 	do_single_rpc_call::<_, ()>(
-		&ctx,
+		&ctx.paths.tunnel_lockfile(),
+		ctx.log,
 		protocol::singleton::METHOD_RESTART,
 		protocol::EmptyObject {},
 	)
 	.await
 	.map(|_| 0)
+	.map_err(|e| e.into())
 }
 
 pub async fn kill(ctx: CommandContext) -> Result<i32, AnyError> {
 	do_single_rpc_call::<_, ()>(
-		&ctx,
+		&ctx.paths.tunnel_lockfile(),
+		ctx.log,
 		protocol::singleton::METHOD_SHUTDOWN,
 		protocol::EmptyObject {},
 	)
 	.await
 	.map(|_| 0)
+	.map_err(|e| e.into())
 }
 
 pub async fn status(ctx: CommandContext) -> Result<i32, AnyError> {
 	let status: protocol::singleton::Status = do_single_rpc_call(
-		&ctx,
+		&ctx.paths.tunnel_lockfile(),
+		ctx.log.clone(),
 		protocol::singleton::METHOD_STATUS,
 		protocol::EmptyObject {},
 	)

cli/src/tunnels/service_windows.rs

@@ -6,17 +6,16 @@
 use async_trait::async_trait;
 use shell_escape::windows::escape as shell_escape;
 use std::{
-	io,
 	path::PathBuf,
 	process::{Command, Stdio},
 };
-use sysinfo::{ProcessExt, System, SystemExt};
 use winreg::{enums::HKEY_CURRENT_USER, RegKey};
 
 use crate::{
 	constants::TUNNEL_ACTIVITY_NAME,
 	log,
 	state::LauncherPaths,
+	tunnels::{protocol, singleton_client::do_single_rpc_call},
 	util::errors::{wrap, wrapdbg, AnyError},
 };
 
@@ -24,6 +23,7 @@ use super::service::{tail_log_file, ServiceContainer, ServiceManager as CliServi
 
 pub struct WindowsService {
 	log: log::Logger,
+	tunnel_lock: PathBuf,
 	log_file: PathBuf,
 }
 
@@ -31,6 +31,7 @@ impl WindowsService {
 	pub fn new(log: log::Logger, paths: &LauncherPaths) -> Self {
 		Self {
 			log,
+			tunnel_lock: paths.tunnel_lockfile(),
 			log_file: paths.service_log_file(),
 		}
 	}
@@ -94,30 +95,24 @@ impl CliServiceManager for WindowsService {
 
 	async fn unregister(&self) -> Result<(), AnyError> {
 		let key = WindowsService::open_key()?;
-		let prev_command_line: String = match key.get_value(TUNNEL_ACTIVITY_NAME) {
-			Ok(l) => l,
-			Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(()),
-			Err(e) => return Err(wrap(e, "error getting registry key").into()),
-		};
-
 		key.delete_value(TUNNEL_ACTIVITY_NAME)
 			.map_err(|e| AnyError::from(wrap(e, "error deleting registry key")))?;
 		info!(self.log, "Tunnel service uninstalled");
 
-		let mut sys = System::new();
-		sys.refresh_processes();
+		let r = do_single_rpc_call::<_, ()>(
+			&self.tunnel_lock,
+			self.log.clone(),
+			protocol::singleton::METHOD_SHUTDOWN,
+			protocol::EmptyObject {},
+		)
+		.await;
 
-		for process in sys.processes().values() {
-			let joined = process.cmd().join(" "); // this feels a little sketch, but seems to work fine
-			if joined == prev_command_line {
-				process.kill();
-				info!(self.log, "Successfully shut down running tunnel");
-				return Ok(());
-			}
+		if r.is_err() {
+			warning!(self.log, "The tunnel service has been unregistered, but we couldn't find a running tunnel process. You may need to restart or log out and back in to fully stop the tunnel.");
+		} else {
+			info!(self.log, "Successfully shut down running tunnel.");
 		}
 
-		warning!(self.log, "The tunnel service has been unregistered, but we couldn't find a running tunnel process. You may need to restart or log out and back in to fully stop the tunnel.");
-
 		Ok(())
 	}
 }

cli/src/tunnels/singleton_client.rs

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 use std::{
+	path::Path,
 	sync::{
 		atomic::{AtomicBool, Ordering},
 		Arc,
@@ -20,11 +21,15 @@ use crate::{
 	json_rpc::{new_json_rpc, start_json_rpc, JsonRpcSerializer},
 	log,
 	rpc::RpcCaller,
+	singleton::connect_as_client,
 	tunnels::{code_server::print_listening, protocol::EmptyObject},
-	util::sync::Barrier,
+	util::{errors::CodeError, sync::Barrier},
 };
 
-use super::{protocol, shutdown_signal::ShutdownSignal};
+use super::{
+	protocol,
+	shutdown_signal::{ShutdownRequest, ShutdownSignal},
+};
 
 pub struct SingletonClientArgs {
 	pub log: log::Logger,
@@ -144,3 +149,43 @@ pub async fn start_singleton_client(args: SingletonClientArgs) -> bool {
 
 	exit_entirely.load(Ordering::SeqCst)
 }
+
+pub async fn do_single_rpc_call<
+	P: serde::Serialize + 'static,
+	R: serde::de::DeserializeOwned + Send + 'static,
+>(
+	lock_file: &Path,
+	log: log::Logger,
+	method: &'static str,
+	params: P,
+) -> Result<R, CodeError> {
+	let client = match connect_as_client(lock_file).await {
+		Ok(p) => p,
+		Err(CodeError::SingletonLockfileOpenFailed(_))
+		| Err(CodeError::SingletonLockedProcessExited(_)) => {
+			return Err(CodeError::NoRunningTunnel);
+		}
+		Err(e) => return Err(e),
+	};
+
+	let (msg_tx, msg_rx) = mpsc::unbounded_channel();
+	let mut rpc = new_json_rpc();
+	let caller = rpc.get_caller(msg_tx);
+	let (read, write) = socket_stream_split(client);
+
+	let rpc = tokio::spawn(async move {
+		start_json_rpc(
+			rpc.methods(()).build(log),
+			read,
+			write,
+			msg_rx,
+			ShutdownRequest::create_rx([ShutdownRequest::CtrlC]),
+		)
+		.await
+		.unwrap();
+	});
+
+	let r = caller.call(method, params).await.unwrap();
+	rpc.abort();
+	r.map_err(CodeError::TunnelRpcCallFailed)
+}