Refactor filesystem store for timeout function passing (#439)

Refactor `filesystem_store::FilesystemStore` to allow for controlling sleep function from outside of implementation. 
Hoist `BytesMut.split()` to a position before the loop to prevent subsequent calls that result in an empty buffer. Original patch: #426 / @Jirixek 

Fixes (#426)
Co-authored-by: Jiří Szkandera @Jirixek

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<Fe: FileEntry = FileEntryImpl> {
     shared_context: Arc<SharedContext>,
     evicting_map: EvictingMap<Arc<Fe>, SystemTime>,
     read_buffer_size: usize,
+    sleep_fn: fn(Duration) -> Sleep,
 }
 
 impl<Fe: FileEntry> FilesystemStore<Fe> {
     pub async fn new(config: &native_link_config::stores::FilesystemStore) -> Result<Self, Error> {
+        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<Self, Error> {
         let now = SystemTime::now();
 
         let empty_policy = native_link_config::stores::EvictionPolicy::default();
@@ -473,6 +481,7 @@ impl<Fe: FileEntry> FilesystemStore<Fe> {
             shared_context,
             evicting_map,
             read_buffer_size,
+            sleep_fn,
         };
         Ok(store)
     }
@@ -654,19 +663,26 @@ impl<Fe: FileEntry> Store for FilesystemStore<Fe> {
             // 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");
+                            }
+                        }
+                    }
                 }
             }
         }

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::<FileEntryImpl>::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(())
+    }
 }