Add network transfer regression test for #942

xet_client/src/cas_client/simulation/network_simulation/bandwidth_limit_router.rs

@@ -36,6 +36,7 @@ pub struct NetworkSimulationProxy {
     current_network_profile: Mutex<NetworkConfig>,
     total_upload_possible: Arc<AtomicU64>,
     total_download_possible: Arc<AtomicU64>,
+    total_download_transferred: Arc<AtomicU64>,
     shutdown_flag: AtomicBool,
     network_profile_provider: NetworkProfile,
     start_instant: Instant,
@@ -74,6 +75,7 @@ impl NetworkSimulationProxy {
             current_network_profile: Mutex::new(initial),
             total_upload_possible: Arc::new(AtomicU64::new(0)),
             total_download_possible: Arc::new(AtomicU64::new(0)),
+            total_download_transferred: Arc::new(AtomicU64::new(0)),
             shutdown_flag: AtomicBool::new(false),
             network_profile_provider: network_profile,
             start_instant: now,
@@ -190,10 +192,16 @@ impl NetworkSimulationProxy {
         let latency = (profile.latency, profile.jitter);
         let upload_limiter = Arc::clone(&self.upload_limiter);
         let download_limiter = Arc::clone(&self.download_limiter);
+        let download_transferred = Arc::clone(&self.total_download_transferred);
         let to_upstream =
-            tokio::spawn(copy_with_rate_and_latency(client_read, upstream_write, Some(upload_limiter), latency));
-        let from_upstream =
-            tokio::spawn(copy_with_rate_and_latency(upstream_read, client_write, Some(download_limiter), latency));
+            tokio::spawn(copy_with_rate_and_latency(client_read, upstream_write, Some(upload_limiter), latency, None));
+        let from_upstream = tokio::spawn(copy_with_rate_and_latency(
+            upstream_read,
+            client_write,
+            Some(download_limiter),
+            latency,
+            Some(download_transferred),
+        ));
         let (to_res, from_res) = tokio::join!(to_upstream, from_upstream);
         to_res.map_err(std::io::Error::other)??;
         from_res.map_err(std::io::Error::other)??;
@@ -214,6 +222,11 @@ impl NetworkSimulationProxy {
         self.total_download_possible.load(Ordering::Relaxed)
     }
 
+    /// Total bytes actually forwarded from the upstream server to clients.
+    pub fn total_download_bytes_transferred(&self) -> u64 {
+        self.total_download_transferred.load(Ordering::Relaxed)
+    }
+
     /// Current bandwidth from the current network profile, in bytes per second, or `None` if unlimited.
     pub async fn current_bandwidth(&self) -> Option<u64> {
         let profile = self.current_network_profile.lock().await;
@@ -240,6 +253,7 @@ async fn copy_with_rate_and_latency<R, W>(
     writer: W,
     limiter: Option<Arc<Semaphore>>,
     latency: (Duration, Duration),
+    transferred: Option<Arc<AtomicU64>>,
 ) -> std::io::Result<u64>
 where
     R: tokio::io::AsyncRead + Unpin + Send + 'static,
@@ -248,7 +262,7 @@ where
     let has_latency = latency.0 > Duration::ZERO || latency.1 > Duration::ZERO;
 
     if !has_latency {
-        return copy_bandwidth_only(reader, writer, limiter).await;
+        return copy_bandwidth_only(reader, writer, limiter, transferred).await;
     }
 
     // Pipeline: reader task enqueues (delivery_time, data), current task dequeues and delivers.
@@ -293,6 +307,9 @@ where
         }
         writer.write_all(&chunk).await?;
         total += chunk.len() as u64;
+        if let Some(counter) = &transferred {
+            counter.fetch_add(chunk.len() as u64, Ordering::Relaxed);
+        }
     }
 
     reader_handle.await.map_err(std::io::Error::other)??;
@@ -305,6 +322,7 @@ async fn copy_bandwidth_only<R, W>(
     mut reader: R,
     mut writer: W,
     limiter: Option<Arc<Semaphore>>,
+    transferred: Option<Arc<AtomicU64>>,
 ) -> std::io::Result<u64>
 where
     R: tokio::io::AsyncRead + Unpin,
@@ -327,6 +345,9 @@ where
         }
         writer.write_all(&buf[..n]).await?;
         total += n as u64;
+        if let Some(counter) = &transferred {
+            counter.fetch_add(n as u64, Ordering::Relaxed);
+        }
     }
     Ok(total)
 }

xet_client/src/cas_client/simulation/simulation_server.rs

@@ -415,6 +415,14 @@ impl LocalTestServer {
             .unwrap_or(0)
     }
 
+    /// Returns total bytes actually forwarded from the server to clients (0 if no proxy).
+    pub fn total_download_bytes_transferred(&self) -> u64 {
+        self.network_simulation_proxy
+            .as_ref()
+            .map(|p| p.total_download_bytes_transferred())
+            .unwrap_or(0)
+    }
+
     /// Current bandwidth from the proxy's network profile (bytes/sec), or `None` if no proxy or unlimited.
     pub async fn current_bandwidth(&self) -> Option<u64> {
         match &self.network_simulation_proxy {

xet_pkg/src/xet_session/file_download_group.rs

@@ -448,7 +448,11 @@ mod tests {
 
     use anyhow::Result;
     use tempfile::tempdir;
+    #[cfg(all(feature = "simulation", not(target_family = "wasm")))]
+    use xet_client::cas_client::simulation::{ClientTestingUtils, LocalTestServerBuilder, NetworkProfileOptions};
     use xet_data::processing::Sha256Policy;
+    #[cfg(all(feature = "simulation", not(target_family = "wasm")))]
+    use xet_runtime::config::XetConfig;
     use xet_runtime::core::RuntimeMode;
 
     use super::*;
@@ -649,6 +653,92 @@ mod tests {
         assert!(matches!(err, XetError::UserCancelled(_)));
     }
 
+    // ── Abort behavior ────────────────────────────────────────────────────────
+
+    #[cfg(all(feature = "simulation", not(target_family = "wasm")))]
+    fn assert_network_transfer_stops_after(abort: impl FnOnce(&XetSession, &XetFileDownloadGroup)) {
+        const CHUNK_SIZE: usize = 64 * 1024;
+        const DOWNLOAD_STARTED_BYTES: u64 = 128 * 1024;
+
+        // The server must outlive an owned Xet runtime so sigint_abort() can tear
+        // down the client without also stopping the measurement endpoint.
+        let server_runtime = tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap();
+        let server = server_runtime.block_on(async {
+            LocalTestServerBuilder::new()
+                .with_network_profile(NetworkProfileOptions::new().with_bandwidth_bytes_per_sec(512 * 1024).build())
+                .start()
+                .await
+        });
+
+        // Keep each range at 64 KiB. With download concurrency capped at one, a
+        // cancellation may let one range drain but must stop scheduling further ranges.
+        let term_spec: Vec<(u64, (u64, u64))> = (1..=64).map(|seed| (seed, (0, 1))).collect();
+        let file = server_runtime
+            .block_on(server.client().upload_random_file(&term_spec, CHUNK_SIZE))
+            .unwrap();
+
+        let config = XetConfig::default()
+            .with_config("client.ac_max_download_concurrency", 1)
+            .unwrap();
+        let session = XetSessionBuilder::new_with_config(config).build().unwrap();
+        assert_eq!(session.inner.ctx.runtime.mode(), RuntimeMode::Owned);
+        let group = session
+            .new_file_download_group()
+            .unwrap()
+            .with_endpoint(server.http_endpoint())
+            .build_blocking()
+            .unwrap();
+        let temp = tempdir().unwrap();
+        let download_bytes_before = server.total_download_bytes_transferred();
+
+        group
+            .download_file_to_path_blocking(
+                XetFileInfo {
+                    hash: file.file_hash.to_string(),
+                    file_size: Some(file.data.len() as u64),
+                    sha256: None,
+                },
+                temp.path().join("download.bin"),
+            )
+            .unwrap();
+
+        let start_deadline = std::time::Instant::now() + Duration::from_secs(5);
+        while server.total_download_bytes_transferred() < download_bytes_before + DOWNLOAD_STARTED_BYTES {
+            assert!(std::time::Instant::now() < start_deadline, "download never started");
+            std::thread::sleep(Duration::from_millis(10));
+        }
+
+        abort(&session, &group);
+
+        // One second is ample for one 64 KiB in-flight range to drain at <= 512 KiB/s.
+        std::thread::sleep(Duration::from_secs(1));
+        let download_bytes_after_grace = server.total_download_bytes_transferred();
+        std::thread::sleep(Duration::from_millis(500));
+        let download_bytes_after_observation = server.total_download_bytes_transferred();
+
+        assert_eq!(
+            download_bytes_after_observation,
+            download_bytes_after_grace,
+            "download kept transferring after cancellation: {} bytes",
+            download_bytes_after_observation - download_bytes_after_grace,
+        );
+    }
+
+    #[cfg(all(feature = "simulation", not(target_family = "wasm")))]
+    #[test]
+    #[ignore = "known failure: #942"]
+    // group.abort() must stop the group's network transfer, not only the caller-facing task.
+    fn test_group_abort_stops_network_transfer() {
+        assert_network_transfer_stops_after(|_, group| group.abort().unwrap());
+    }
+
+    #[cfg(all(feature = "simulation", not(target_family = "wasm")))]
+    #[test]
+    // session.sigint_abort() is a positive control that stops the same network transfer.
+    fn test_sigint_abort_stops_network_transfer() {
+        assert_network_transfer_stops_after(|session, _| session.sigint_abort().unwrap());
+    }
+
     // ── Independence ─────────────────────────────────────────────────────────
 
     #[tokio::test(flavor = "multi_thread")]