diff --git a/native-link-store/src/filesystem_store.rs b/native-link-store/src/filesystem_store.rs index 63c26fd36..0027d263f 100644 --- a/native-link-store/src/filesystem_store.rs +++ b/native-link-store/src/filesystem_store.rs @@ -17,7 +17,7 @@ use std::fmt::{Debug, Formatter}; use std::pin::Pin; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; -use std::time::SystemTime; +use std::time::{Duration, SystemTime}; use async_lock::RwLock; use async_trait::async_trait; @@ -33,7 +33,7 @@ use native_link_util::metrics_utils::{Collector, CollectorState, MetricsComponen use native_link_util::store_trait::{Store, UploadSizeInfo}; use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt, SeekFrom}; use tokio::task::spawn_blocking; -use tokio::time::timeout; +use tokio::time::{sleep, timeout, Sleep}; use tokio_stream::wrappers::ReadDirStream; // Default size to allocate memory of the buffer when reading files. @@ -439,10 +439,18 @@ pub struct FilesystemStore { shared_context: Arc, evicting_map: EvictingMap, SystemTime>, read_buffer_size: usize, + sleep_fn: fn(Duration) -> Sleep, } impl FilesystemStore { pub async fn new(config: &native_link_config::stores::FilesystemStore) -> Result { + Self::new_with_timeout(config, sleep).await + } + + pub async fn new_with_timeout( + config: &native_link_config::stores::FilesystemStore, + sleep_fn: fn(Duration) -> Sleep, + ) -> Result { let now = SystemTime::now(); let empty_policy = native_link_config::stores::EvictionPolicy::default(); @@ -473,6 +481,7 @@ impl FilesystemStore { shared_context, evicting_map, read_buffer_size, + sleep_fn, }; Ok(store) } @@ -654,19 +663,26 @@ impl Store for FilesystemStore { // because it is waiting for a file descriptor to open before receiving data. // Using `ResumeableFileSlot` will re-open the file in the event it gets closed on the // next iteration. + let buf_content = buf.split().freeze(); loop { - match timeout(fs::idle_file_descriptor_timeout(), writer.send(buf.split().freeze())).await { - Ok(Ok(())) => break, - Ok(Err(err)) => { - return Err(err).err_tip(|| "Failed to send chunk in filesystem store get_part"); - } - Err(_) => { + let sleep_fn = (self.sleep_fn)(fs::idle_file_descriptor_timeout()); + tokio::pin!(sleep_fn); + tokio::select! { + _ = & mut (sleep_fn) => { resumeable_temp_file .close_file() .await .err_tip(|| "Could not close file due to timeout in FileSystemStore::get_part")?; continue; } + res = writer.send(buf_content.clone()) => { + match res { + Ok(()) => break, + Err(err) => { + return Err(err).err_tip(|| "Failed to send chunk in filesystem store get_part"); + } + } + } } } } diff --git a/native-link-store/tests/filesystem_store_test.rs b/native-link-store/tests/filesystem_store_test.rs index 31275185f..0ff41873d 100644 --- a/native-link-store/tests/filesystem_store_test.rs +++ b/native-link-store/tests/filesystem_store_test.rs @@ -23,7 +23,7 @@ use std::pin::Pin; use std::rc::Rc; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{Arc, Mutex}; -use std::time::SystemTime; +use std::time::{Duration, SystemTime}; use async_lock::RwLock; use async_trait::async_trait; @@ -36,13 +36,14 @@ use native_link_store::filesystem_store::{ digest_from_filename, EncodedFilePath, FileEntry, FileEntryImpl, FilesystemStore, }; use native_link_util::buf_channel::{make_buf_channel_pair, DropCloserReadHalf}; -use native_link_util::common::{fs, DigestInfo}; +use native_link_util::common::{fs, DigestInfo, JoinHandleDropGuard}; use native_link_util::evicting_map::LenEntry; use native_link_util::store_trait::{Store, UploadSizeInfo}; use once_cell::sync::Lazy; use rand::{thread_rng, Rng}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::sync::Barrier; +use tokio::time::sleep; use tokio_stream::wrappers::ReadDirStream; use tokio_stream::StreamExt; @@ -273,7 +274,6 @@ mod filesystem_store_tests { .await?, ); - // Insert data into store. store.as_ref().update_oneshot(digest1, VALUE1.into()).await?; let expected_file_name = OsString::from(format!( @@ -865,4 +865,48 @@ mod filesystem_store_tests { Ok(()) } + + #[tokio::test] + async fn get_part_timeout_test() -> Result<(), Error> { + let large_value = "x".repeat(1024); + let digest = DigestInfo::try_new(HASH1, large_value.len())?; + let content_path = make_temp_path("content_path"); + let temp_path = make_temp_path("temp_path"); + + let store = Arc::new( + FilesystemStore::::new_with_timeout( + &native_link_config::stores::FilesystemStore { + content_path: content_path.clone(), + temp_path: temp_path.clone(), + read_buffer_size: 1, + ..Default::default() + }, + |_| sleep(Duration::ZERO), + ) + .await?, + ); + + let store_pin = Pin::new(store.as_ref()); + + store_pin + .as_ref() + .update_oneshot(digest, large_value.clone().into()) + .await?; + + let (writer, mut reader) = make_buf_channel_pair(); + let store_clone = store.clone(); + let digest_clone = digest; + + let _drop_guard = JoinHandleDropGuard::new(tokio::spawn(async move { + Pin::new(store_clone.as_ref()).get(digest_clone, writer).await + })); + + let file_data = DropCloserReadHalf::take(&mut reader, 1024) + .await + .err_tip(|| "Error reading bytes")?; + + assert_eq!(&file_data, large_value.as_bytes(), "Expected file content to match"); + + Ok(()) + } }