From 6e7ddbd5582c5caef7805669fcfbc7e0f2e3e867 Mon Sep 17 00:00:00 2001 From: srikanthananthula63053 Date: Thu, 20 Aug 2026 11:22:49 +0530 Subject: [PATCH] cli: reuse a fully-staged server download after an interrupted update Fixes #331690 DownloadCache::create() unconditionally deleted its .staging directory before every download. If the process was killed after a download+ extract finished populating .staging but before the atomic rename into the final commit-keyed directory (e.g. VS Code closed right as a server update finished), the next launch discarded that already-valid download and re-fetched it from the update service, wasting bandwidth on a patch that was already available locally. Write a completion marker into the staging directory right after the do_create closure succeeds, before the rename is attempted. On the next call, a marker's presence means the staged content is already complete and can be renamed into place directly instead of being wiped and re-downloaded; a staging directory without the marker is still treated as partial and discarded, unchanged from before. The marker is stripped from the final directory after rename so callers that enumerate it don't see it. No call site of create() needed to change. --- cli/src/download_cache.rs | 177 +++++++++++++++++++++++++++++++++++++- 1 file changed, 174 insertions(+), 3 deletions(-) diff --git a/cli/src/download_cache.rs b/cli/src/download_cache.rs index 87ca1924a798cc..b857210ffc99d3 100644 --- a/cli/src/download_cache.rs +++ b/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( &self, name: impl AsRef, @@ -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"); + } +}