Run tunnels as singleton process (for a --cli-data-dir) (#177002)

* wip on singleton

* wip

* windows support

* wip

* wip

* fix clippy

cli/src/util/errors.rs

@@ -2,11 +2,11 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
-use std::fmt::Display;
-
 use crate::constants::{
 	APPLICATION_NAME, CONTROL_PORT, DOCUMENTATION_URL, QUALITYLESS_PRODUCT_NAME,
 };
+use std::fmt::Display;
+use thiserror::Error;
 
 // Wraps another error with additional info.
 #[derive(Debug, Clone)]
@@ -475,6 +475,22 @@ macro_rules! makeAnyError {
     };
 }
 
+/// Internal errors in the VS Code CLI.
+/// Note: other error should be migrated to this type gradually
+#[derive(Error, Debug)]
+pub enum CodeError {
+	#[error("could not connect to socket/pipe")]
+	AsyncPipeFailed(std::io::Error),
+	#[error("could not listen on socket/pipe")]
+	AsyncPipeListenerFailed(std::io::Error),
+	#[error("could not create singleton lock file")]
+	SingletonLockfileOpenFailed(std::io::Error),
+	#[error("could not read singleton lock file")]
+	SingletonLockfileReadFailed(rmp_serde::decode::Error),
+	#[error("the process holding the singleton lock file exited")]
+	SingletonLockedProcessExited(u32),
+}
+
 makeAnyError!(
 	MissingLegalConsent,
 	MismatchConnectionToken,
@@ -505,7 +521,8 @@ makeAnyError!(
 	MissingHomeDirectory,
 	CommandFailed,
 	OAuthError,
-	InvalidRpcDataError
+	InvalidRpcDataError,
+	CodeError
 );
 
 impl From<reqwest::Error> for AnyError {

cli/src/util/file_lock.rs

@@ -0,0 +1,125 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+use crate::util::errors::CodeError;
+use std::{fs::File, io};
+
+pub struct FileLock {
+	file: File,
+	#[cfg(windows)]
+	overlapped: winapi::um::minwinbase::OVERLAPPED,
+}
+
+#[cfg(windows)] // overlapped is thread-safe, mark it so with this
+unsafe impl Send for FileLock {}
+
+pub enum Lock {
+	Acquired(FileLock),
+	AlreadyLocked(File),
+}
+
+/// Number of locked bytes in the file. On Windows, locking prevents reads,
+/// but consumers of the lock may still want to read what the locking file
+/// as written. Thus, only PREFIX_LOCKED_BYTES are locked, and any globally-
+/// readable content should be written after the prefix.
+#[cfg(windows)]
+pub const PREFIX_LOCKED_BYTES: usize = 1;
+
+#[cfg(unix)]
+pub const PREFIX_LOCKED_BYTES: usize = 0;
+
+impl FileLock {
+	#[cfg(windows)]
+	pub fn acquire(file: File) -> Result<Lock, CodeError> {
+		use std::os::windows::prelude::AsRawHandle;
+		use winapi::{
+			shared::winerror::{ERROR_IO_PENDING, ERROR_LOCK_VIOLATION},
+			um::{
+				fileapi::LockFileEx,
+				minwinbase::{LOCKFILE_EXCLUSIVE_LOCK, LOCKFILE_FAIL_IMMEDIATELY},
+			},
+		};
+
+		let handle = file.as_raw_handle();
+		let (overlapped, ok) = unsafe {
+			let mut overlapped = std::mem::zeroed();
+			let ok = LockFileEx(
+				handle,
+				LOCKFILE_EXCLUSIVE_LOCK | LOCKFILE_FAIL_IMMEDIATELY,
+				0,
+				PREFIX_LOCKED_BYTES as u32,
+				0,
+				&mut overlapped,
+			);
+
+			(overlapped, ok)
+		};
+
+		if ok != 0 {
+			return Ok(Lock::Acquired(Self { file, overlapped }));
+		}
+
+		let err = io::Error::last_os_error();
+		let raw = err.raw_os_error();
+		// docs report it should return ERROR_IO_PENDING, but in my testing it actually
+		// returns ERROR_LOCK_VIOLATION. Or maybe winapi is wrong?
+		if raw == Some(ERROR_IO_PENDING as i32) || raw == Some(ERROR_LOCK_VIOLATION as i32) {
+			return Ok(Lock::AlreadyLocked(file));
+		}
+
+		Err(CodeError::SingletonLockfileOpenFailed(err))
+	}
+
+	#[cfg(unix)]
+	pub fn acquire(file: File) -> Result<Lock, CodeError> {
+		use std::os::unix::io::AsRawFd;
+
+		let fd = file.as_raw_fd();
+		let res = unsafe { libc::flock(fd, libc::LOCK_EX | libc::LOCK_NB) };
+		if res == 0 {
+			return Ok(Lock::Acquired(Self { file }));
+		}
+
+		let err = io::Error::last_os_error();
+		if err.kind() == io::ErrorKind::WouldBlock {
+			return Ok(Lock::AlreadyLocked(file));
+		}
+
+		Err(CodeError::SingletonLockfileOpenFailed(err))
+	}
+
+	pub fn file(&self) -> &File {
+		&self.file
+	}
+
+	pub fn file_mut(&mut self) -> &mut File {
+		&mut self.file
+	}
+}
+
+impl Drop for FileLock {
+	#[cfg(windows)]
+	fn drop(&mut self) {
+		use std::os::windows::prelude::AsRawHandle;
+		use winapi::um::fileapi::UnlockFileEx;
+
+		unsafe {
+			UnlockFileEx(
+				self.file.as_raw_handle(),
+				0,
+				u32::MAX,
+				u32::MAX,
+				&mut self.overlapped,
+			)
+		};
+	}
+
+	#[cfg(unix)]
+	fn drop(&mut self) {
+		use std::os::unix::io::AsRawFd;
+
+		unsafe { libc::flock(self.file.as_raw_fd(), libc::LOCK_UN) };
+	}
+}

cli/src/util/io.rs

@@ -15,6 +15,8 @@ use tokio::{
 	time::sleep,
 };
 
+use super::ring_buffer::RingBuffer;
+
 pub trait ReportCopyProgress {
 	fn report_progress(&mut self, bytes_so_far: u64, total_bytes: u64);
 }
@@ -132,8 +134,7 @@ pub fn tailf(file: File, n: usize) -> mpsc::UnboundedReceiver<TailEvent> {
 
 	// Read the initial "n" lines back from the request. initial_lines
 	// is a small ring buffer.
-	let mut initial_lines = Vec::with_capacity(n);
-	let mut initial_lines_i = 0;
+	let mut initial_lines = RingBuffer::new(n);
 	loop {
 		let mut line = String::new();
 		let bytes_read = match reader.read_line(&mut line) {
@@ -151,26 +152,11 @@ pub fn tailf(file: File, n: usize) -> mpsc::UnboundedReceiver<TailEvent> {
 		}
 
 		pos += bytes_read as u64;
-		if initial_lines.len() < initial_lines.capacity() {
-			initial_lines.push(line)
-		} else {
-			initial_lines[initial_lines_i] = line;
-		}
-
-		initial_lines_i = (initial_lines_i + 1) % n;
+		initial_lines.push(line);
 	}
 
-	// remove tail lines...
-	if initial_lines_i < initial_lines.len() {
-		for line in initial_lines.drain((initial_lines_i)..) {
-			tx.send(TailEvent::Line(line)).ok();
-		}
-	}
-	// then the remaining lines
-	if !initial_lines.is_empty() {
-		for line in initial_lines.drain(0..) {
-			tx.send(TailEvent::Line(line)).ok();
-		}
+	for line in initial_lines.into_iter() {
+		tx.send(TailEvent::Line(line)).ok();
 	}
 
 	// now spawn the poll process to keep reading new lines

cli/src/util/machine.rs

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-use std::path::Path;
+use std::{path::Path, time::Duration};
 use sysinfo::{Pid, PidExt, ProcessExt, System, SystemExt};
 
 pub fn process_at_path_exists(pid: u32, name: &Path) -> bool {
@@ -29,6 +29,14 @@ pub fn process_exists(pid: u32) -> bool {
 	sys.refresh_process(Pid::from_u32(pid))
 }
 
+pub async fn wait_until_process_exits(pid: Pid, poll_ms: u64) {
+	let mut s = System::new();
+	let duration = Duration::from_millis(poll_ms);
+	while s.refresh_process(pid) {
+		tokio::time::sleep(duration).await;
+	}
+}
+
 pub fn find_running_process(name: &Path) -> Option<u32> {
 	let mut sys = System::new();
 	sys.refresh_processes();

cli/src/util/ring_buffer.rs

@@ -0,0 +1,142 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+pub struct RingBuffer<T> {
+	data: Vec<T>,
+	i: usize,
+}
+
+impl<T> RingBuffer<T> {
+	pub fn new(capacity: usize) -> Self {
+		Self {
+			data: Vec::with_capacity(capacity),
+			i: 0,
+		}
+	}
+
+	pub fn capacity(&self) -> usize {
+		self.data.capacity()
+	}
+
+	pub fn len(&self) -> usize {
+		self.data.len()
+	}
+
+	pub fn is_full(&self) -> bool {
+		self.data.len() == self.data.capacity()
+	}
+
+	pub fn is_empty(&self) -> bool {
+		self.data.len() == 0
+	}
+
+	pub fn push(&mut self, value: T) {
+		if self.data.len() == self.data.capacity() {
+			self.data[self.i] = value;
+		} else {
+			self.data.push(value);
+		}
+
+		self.i = (self.i + 1) % self.data.capacity();
+	}
+
+	pub fn iter(&self) -> RingBufferIter<'_, T> {
+		RingBufferIter {
+			index: 0,
+			buffer: self,
+		}
+	}
+}
+
+impl<T: Default> IntoIterator for RingBuffer<T> {
+	type Item = T;
+	type IntoIter = OwnedRingBufferIter<T>;
+
+	fn into_iter(self) -> OwnedRingBufferIter<T>
+	where
+		T: Default,
+	{
+		OwnedRingBufferIter {
+			index: 0,
+			buffer: self,
+		}
+	}
+}
+
+pub struct OwnedRingBufferIter<T: Default> {
+	buffer: RingBuffer<T>,
+	index: usize,
+}
+
+impl<T: Default> Iterator for OwnedRingBufferIter<T> {
+	type Item = T;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		if self.index == self.buffer.len() {
+			return None;
+		}
+
+		let ii = (self.index + self.buffer.i) % self.buffer.len();
+		let item = std::mem::take(&mut self.buffer.data[ii]);
+		self.index += 1;
+		Some(item)
+	}
+}
+
+pub struct RingBufferIter<'a, T> {
+	buffer: &'a RingBuffer<T>,
+	index: usize,
+}
+
+impl<'a, T> Iterator for RingBufferIter<'a, T> {
+	type Item = &'a T;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		if self.index == self.buffer.len() {
+			return None;
+		}
+
+		let ii = (self.index + self.buffer.i) % self.buffer.len();
+		let item = &self.buffer.data[ii];
+		self.index += 1;
+		Some(item)
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn test_inserts() {
+		let mut rb = RingBuffer::new(3);
+		assert_eq!(rb.capacity(), 3);
+		assert!(!rb.is_full());
+		assert_eq!(rb.len(), 0);
+		assert_eq!(rb.iter().copied().count(), 0);
+
+		rb.push(1);
+		assert!(!rb.is_full());
+		assert_eq!(rb.len(), 1);
+		assert_eq!(rb.iter().copied().collect::<Vec<i32>>(), vec![1]);
+
+		rb.push(2);
+		assert!(!rb.is_full());
+		assert_eq!(rb.len(), 2);
+		assert_eq!(rb.iter().copied().collect::<Vec<i32>>(), vec![1, 2]);
+
+		rb.push(3);
+		assert!(rb.is_full());
+		assert_eq!(rb.len(), 3);
+		assert_eq!(rb.iter().copied().collect::<Vec<i32>>(), vec![1, 2, 3]);
+
+		rb.push(4);
+		assert!(rb.is_full());
+		assert_eq!(rb.len(), 3);
+		assert_eq!(rb.iter().copied().collect::<Vec<i32>>(), vec![2, 3, 4]);
+
+		assert_eq!(rb.into_iter().collect::<Vec<i32>>(), vec![2, 3, 4]);
+	}
+}

cli/src/util/sync.rs

@@ -2,38 +2,53 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
-use tokio::sync::watch::{
-	self,
-	error::{RecvError, SendError},
+use async_trait::async_trait;
+use std::{marker::PhantomData, sync::Arc};
+use tokio::sync::{
+	broadcast, mpsc,
+	watch::{self, error::RecvError},
 };
 
 #[derive(Clone)]
 pub struct Barrier<T>(watch::Receiver<Option<T>>)
 where
-	T: Copy;
+	T: Clone;
 
 impl<T> Barrier<T>
 where
-	T: Copy,
+	T: Clone,
 {
 	/// Waits for the barrier to be closed, returning a value if one was sent.
 	pub async fn wait(&mut self) -> Result<T, RecvError> {
 		loop {
 			self.0.changed().await?;
 
-			if let Some(v) = *(self.0.borrow()) {
+			if let Some(v) = self.0.borrow().clone() {
 				return Ok(v);
 			}
 		}
 	}
+
+	/// Gets whether the barrier is currently open
+	pub fn is_open(&self) -> bool {
+		self.0.borrow().is_some()
+	}
 }
 
-pub struct BarrierOpener<T>(watch::Sender<Option<T>>);
+#[derive(Clone)]
+pub struct BarrierOpener<T: Clone>(Arc<watch::Sender<Option<T>>>);
 
-impl<T> BarrierOpener<T> {
-	/// Closes the barrier.
-	pub fn open(self, value: T) -> Result<(), SendError<Option<T>>> {
-		self.0.send(Some(value))
+impl<T: Clone> BarrierOpener<T> {
+	/// Opens the barrier.
+	pub fn open(&self, value: T) {
+		self.0.send_if_modified(|v| {
+			if v.is_none() {
+				*v = Some(value);
+				true
+			} else {
+				false
+			}
+		});
 	}
 }
 
@@ -44,7 +59,119 @@ where
 	T: Copy,
 {
 	let (closed_tx, closed_rx) = watch::channel(None);
-	(Barrier(closed_rx), BarrierOpener(closed_tx))
+	(Barrier(closed_rx), BarrierOpener(Arc::new(closed_tx)))
+}
+
+/// Type that can receive messages in an async way.
+#[async_trait]
+pub trait Receivable<T> {
+	async fn recv_msg(&mut self) -> Option<T>;
+}
+
+// todo: ideally we would use an Arc in the broadcast::Receiver to avoid having
+// to clone bytes everywhere, requires updating rpc consumers as well.
+#[async_trait]
+impl<T: Clone + Send> Receivable<T> for broadcast::Receiver<T> {
+	async fn recv_msg(&mut self) -> Option<T> {
+		loop {
+			match self.recv().await {
+				Ok(v) => return Some(v),
+				Err(broadcast::error::RecvError::Lagged(_)) => continue,
+				Err(broadcast::error::RecvError::Closed) => return None,
+			}
+		}
+	}
+}
+
+#[async_trait]
+impl<T: Send> Receivable<T> for mpsc::UnboundedReceiver<T> {
+	async fn recv_msg(&mut self) -> Option<T> {
+		self.recv().await
+	}
+}
+
+#[async_trait]
+impl<T: Send> Receivable<T> for () {
+	async fn recv_msg(&mut self) -> Option<T> {
+		futures::future::pending().await
+	}
+}
+
+pub struct ConcatReceivable<T: Send, A: Receivable<T>, B: Receivable<T>> {
+	left: Option<A>,
+	right: B,
+	_marker: PhantomData<T>,
+}
+
+impl<T: Send, A: Receivable<T>, B: Receivable<T>> ConcatReceivable<T, A, B> {
+	pub fn new(left: A, right: B) -> Self {
+		Self {
+			left: Some(left),
+			right,
+			_marker: PhantomData,
+		}
+	}
+}
+
+#[async_trait]
+impl<T: Send, A: Send + Receivable<T>, B: Send + Receivable<T>> Receivable<T>
+	for ConcatReceivable<T, A, B>
+{
+	async fn recv_msg(&mut self) -> Option<T> {
+		if let Some(left) = &mut self.left {
+			match left.recv_msg().await {
+				Some(v) => return Some(v),
+				None => {
+					self.left = None;
+				}
+			}
+		}
+
+		return self.right.recv_msg().await;
+	}
+}
+
+pub struct MergedReceivable<T: Send, A: Receivable<T>, B: Receivable<T>> {
+	left: Option<A>,
+	right: Option<B>,
+	_marker: PhantomData<T>,
+}
+
+impl<T: Send, A: Receivable<T>, B: Receivable<T>> MergedReceivable<T, A, B> {
+	pub fn new(left: A, right: B) -> Self {
+		Self {
+			left: Some(left),
+			right: Some(right),
+			_marker: PhantomData,
+		}
+	}
+}
+
+#[async_trait]
+impl<T: Send, A: Send + Receivable<T>, B: Send + Receivable<T>> Receivable<T>
+	for MergedReceivable<T, A, B>
+{
+	async fn recv_msg(&mut self) -> Option<T> {
+		loop {
+			match (&mut self.left, &mut self.right) {
+				(Some(left), Some(right)) => {
+					tokio::select! {
+						left = left.recv_msg() => match left {
+							Some(v) => return Some(v),
+							None => { self.left = None; continue; },
+						},
+						right = right.recv_msg() => match right {
+							Some(v) => return Some(v),
+							None => { self.right = None; continue; },
+						},
+					}
+				}
+				(Some(a), None) => break a.recv_msg().await,
+				(None, Some(b)) => break b.recv_msg().await,
+				(None, None) => break None,
+			}
+		}
+	}
 }
 
 #[cfg(test)]
@@ -60,7 +187,7 @@ mod tests {
 			tx.send(barrier.wait().await.unwrap()).unwrap();
 		});
 
-		opener.open(42).unwrap();
+		opener.open(42);
 
 		assert!(rx.await.unwrap() == 42);
 	}
@@ -71,7 +198,7 @@ mod tests {
 		let (tx1, rx1) = tokio::sync::oneshot::channel::<u32>();
 		let (tx2, rx2) = tokio::sync::oneshot::channel::<u32>();
 
-		opener.open(42).unwrap();
+		opener.open(42);
 		let mut b1 = barrier.clone();
 		tokio::spawn(async move {
 			tx1.send(b1.wait().await.unwrap()).unwrap();