cli: reuse a fully-staged server download after an interrupted update

cli/src/download_cache.rs

@@ -18,6 +18,11 @@ use crate::{
 
 const KEEP_LRU: usize = 5;
 const STAGING_SUFFIX: &str = ".staging";
+/// Sentinel file written into a staging directory once `do_create` has
+/// finished populating it. Its presence is the only signal that staged
+/// content is complete and safe to reuse rather than a partial download
+/// left behind by an interrupted run; see `create()`.
+const STAGING_COMPLETE_MARKER: &str = ".complete";
 const RENAME_ATTEMPTS: u32 = 20;
 const RENAME_DELAY: std::time::Duration = std::time::Duration::from_millis(200);
 const PERSISTED_STATE_FILE_NAME: &str = "lru.json";
@@ -75,6 +80,14 @@ impl DownloadCache {
 	/// returning the path where the folder is. Note that the path passed to
 	/// the `do_create` method is a staging path and will not be the same as the
 	/// final returned path.
+	///
+	/// If a previous call was interrupted (e.g. the process was killed)
+	/// after `do_create` finished but before the staging directory could be
+	/// renamed into place, the next call reuses that already-downloaded
+	/// staging directory instead of deleting it and downloading again. A
+	/// staging directory without the completion marker is assumed to be
+	/// partial (interrupted mid-download/extract) and is discarded as
+	/// before.
 	pub async fn create<F, T>(
 		&self,
 		name: impl AsRef<str>,
@@ -91,10 +104,22 @@ impl DownloadCache {
 		}
 
 		let temp_dir = self.path.join(format!("{name}{STAGING_SUFFIX}"));
-		let _ = remove_dir_all(&temp_dir).await; // cleanup any existing
+		let marker = temp_dir.join(STAGING_COMPLETE_MARKER);
 
-		create_dir_all(&temp_dir).map_err(|e| wrap(e, "error creating server directory"))?;
-		do_create(temp_dir.clone()).await?;
+		if marker.exists() {
+			debug_assert!(temp_dir.exists());
+		} else {
+			let _ = remove_dir_all(&temp_dir).await; // cleanup any incomplete attempt
+
+			create_dir_all(&temp_dir).map_err(|e| wrap(e, "error creating server directory"))?;
+			do_create(temp_dir.clone()).await?;
+
+			// Record that staging finished successfully before attempting the
+			// rename below, so an interruption between here and a successful
+			// rename can be recovered on the next call without re-downloading.
+			std::fs::write(&marker, b"")
+				.map_err(|e| wrap(e, "error marking download as complete"))?;
+		}
 
 		let _ = self.touch(name.to_string());
 		// retry the rename, it seems on WoA sometimes it takes a second for the
@@ -113,6 +138,11 @@ impl DownloadCache {
 			}
 		}
 
+		// The marker is staging-only bookkeeping; strip it from the final
+		// location so callers that enumerate the returned directory (e.g.
+		// looking for the single extracted CLI binary) don't see it.
+		let _ = std::fs::remove_file(target_dir.join(STAGING_COMPLETE_MARKER));
+
 		Ok(target_dir)
 	}
 
@@ -138,3 +168,144 @@ impl DownloadCache {
 		Ok(())
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use std::sync::{
+		atomic::{AtomicBool, Ordering},
+		Arc,
+	};
+
+	/// An existing, fully-renamed target directory (the common case: the
+	/// commit was already installed on a previous run) must be returned
+	/// as-is, without ever invoking `do_create` -- i.e. without touching
+	/// the network.
+	#[tokio::test]
+	async fn test_existing_valid_target_is_reused_without_download() {
+		let dir = tempfile::tempdir().unwrap();
+		let cache = DownloadCache::new(dir.path().to_path_buf());
+		let name = "some-commit";
+		std::fs::create_dir_all(dir.path().join(name)).unwrap();
+		std::fs::write(dir.path().join(name).join("payload.bin"), b"already installed").unwrap();
+
+		let called = Arc::new(AtomicBool::new(false));
+		let called_inner = called.clone();
+		let result = cache
+			.create(name, move |_target_dir| {
+				called_inner.store(true, Ordering::SeqCst);
+				async move { Ok(()) }
+			})
+			.await
+			.unwrap();
+
+		assert!(
+			!called.load(Ordering::SeqCst),
+			"do_create must not run when the target directory already exists"
+		);
+		assert_eq!(result, dir.path().join(name));
+		assert!(result.join("payload.bin").exists());
+	}
+
+	/// When nothing is cached at all, `do_create` must run and its output
+	/// must end up at the final (non-staging) path.
+	#[tokio::test]
+	async fn test_missing_patch_triggers_download() {
+		let dir = tempfile::tempdir().unwrap();
+		let cache = DownloadCache::new(dir.path().to_path_buf());
+		let name = "some-commit";
+
+		let result = cache
+			.create(name, |target_dir| async move {
+				std::fs::write(target_dir.join("payload.bin"), b"downloaded").unwrap();
+				Ok(())
+			})
+			.await
+			.unwrap();
+
+		assert_eq!(result, dir.path().join(name));
+		assert!(result.join("payload.bin").exists());
+		assert!(!dir.path().join(format!("{name}{STAGING_SUFFIX}")).exists());
+	}
+
+	/// A staging directory left over without the completion marker
+	/// represents a download/extraction that was interrupted partway
+	/// through (e.g. connection lost mid-transfer). It must be discarded
+	/// and `do_create` must run again from scratch, so stale partial
+	/// content can never leak into the final installed directory.
+	#[tokio::test]
+	async fn test_incomplete_staging_download_triggers_fresh_download() {
+		let dir = tempfile::tempdir().unwrap();
+		let cache = DownloadCache::new(dir.path().to_path_buf());
+		let name = "some-commit";
+
+		let staging = dir.path().join(format!("{name}{STAGING_SUFFIX}"));
+		std::fs::create_dir_all(&staging).unwrap();
+		std::fs::write(staging.join("partial.tmp"), b"half-downloaded").unwrap();
+
+		let called = Arc::new(AtomicBool::new(false));
+		let called_inner = called.clone();
+		let result = cache
+			.create(name, move |target_dir| {
+				called_inner.store(true, Ordering::SeqCst);
+				async move {
+					std::fs::write(target_dir.join("payload.bin"), b"downloaded").unwrap();
+					Ok(())
+				}
+			})
+			.await
+			.unwrap();
+
+		assert!(
+			called.load(Ordering::SeqCst),
+			"do_create must run again for an incomplete staging directory"
+		);
+		assert!(result.join("payload.bin").exists());
+		assert!(
+			!result.join("partial.tmp").exists(),
+			"stale partial content from the interrupted attempt must not survive"
+		);
+	}
+
+	/// Regression test for https://github.com/microsoft/vscode/issues/331690:
+	/// if the process is killed after a download+extract fully finished in
+	/// the staging directory but before it could be renamed into place,
+	/// restarting must reuse that already-downloaded content instead of
+	/// re-downloading it from the update service.
+	#[tokio::test]
+	async fn test_interrupted_update_restart_reuses_completed_staging_without_redownload() {
+		let dir = tempfile::tempdir().unwrap();
+		let cache = DownloadCache::new(dir.path().to_path_buf());
+		let name = "some-commit";
+
+		// Simulate a prior run that finished populating the staging directory
+		// (do_create succeeded and the completion marker was written) but was
+		// killed before the final rename ran.
+		let staging = dir.path().join(format!("{name}{STAGING_SUFFIX}"));
+		std::fs::create_dir_all(&staging).unwrap();
+		std::fs::write(staging.join("payload.bin"), b"downloaded").unwrap();
+		std::fs::write(staging.join(STAGING_COMPLETE_MARKER), b"").unwrap();
+
+		let called = Arc::new(AtomicBool::new(false));
+		let called_inner = called.clone();
+		let result = cache
+			.create(name, move |_target_dir| {
+				called_inner.store(true, Ordering::SeqCst);
+				async move { Ok(()) }
+			})
+			.await
+			.unwrap();
+
+		assert!(
+			!called.load(Ordering::SeqCst),
+			"restarting after an interrupted update must not redownload an already-available patch"
+		);
+		assert_eq!(result, dir.path().join(name));
+		assert!(result.join("payload.bin").exists());
+		assert!(
+			!result.join(STAGING_COMPLETE_MARKER).exists(),
+			"the internal completion marker must not leak into the final installed directory"
+		);
+		assert!(!staging.exists(), "the staging directory must be renamed away, not left behind");
+	}
+}