more comments

src/fs.rs

@@ -11,31 +11,61 @@ use futures::{Future, Stream, StreamExt};
 use once_cell::sync::Lazy;
 use parking_lot::RwLock;
 use reqwest::Error;
-use tokio::fs::{remove_file, File, OpenOptions};
+use tokio::fs::{remove_file, File};
 use tokio::io::{AsyncRead, AsyncWriteExt, ReadBuf};
 use tokio::time::Sleep;
 
+/// Keeps track of files that are currently being written to.
+///
+/// Why is this necessary? Consider the following situation:
+///
+/// Client A requests file `foo.png`. We construct a transparent file stream,
+/// and now the file is being streamed into and from.
+///
+/// Client B requests the same file `foo.png`. A naive implementation would
+/// attempt to either read directly the file as it sees the file existing. This
+/// is problematic as the file could still be written to. If Client B catches
+/// up to Client A's request, then Client B could receive a broken image, as it
+/// thinks it's done reading the file.
+///
+/// We effectively use WRITING_STATUS as a status relay to ensure concurrent
+/// reads to the file while it's being written to will wait for writing to be
+/// completed.
 static WRITING_STATUS: Lazy<RwLock<HashMap<PathBuf, Arc<CacheStatus>>>> =
     Lazy::new(|| RwLock::new(HashMap::new()));
 
+/// Tries to read from the file, returning a byte stream if it exists
+pub async fn read_file(
+    path: &Path,
+) -> Option<Result<impl Stream<Item = Result<Bytes, UpstreamError>>, std::io::Error>> {
+    if path.exists() {
+        if let Some(status) = WRITING_STATUS.read().get(path) {
+            Some(FromFsStream::new(path, Arc::clone(status)).await)
+        } else {
+            Some(FromFsStream::new(path, Arc::new(CacheStatus::done())).await)
+        }
+    } else {
+        None
+    }
+}
+
+/// Maps the input byte stream into one that writes to disk instead, returning
+/// a stream that reads from disk instead.
 pub async fn transparent_file_stream(
     path: &Path,
     mut byte_stream: impl Stream<Item = Result<Bytes, Error>> + Unpin + Send + 'static,
 ) -> Result<impl Stream<Item = Result<Bytes, UpstreamError>>, std::io::Error> {
-    if let Some(arc) = WRITING_STATUS.read().get(path) {
-        FromFsStream::new(path, Arc::clone(arc)).await
-    } else {
     let done_writing_flag = Arc::new(CacheStatus::new());
 
-        {
+    let mut file = {
         let mut write_lock = WRITING_STATUS.write();
-            File::create(path).await?; // we need to make sure the file exists and is truncated.
+        let file = File::create(path).await?; // we need to make sure the file exists and is truncated.
         write_lock.insert(path.to_path_buf(), Arc::clone(&done_writing_flag));
-        }
+        file
+    };
 
     let write_flag = Arc::clone(&done_writing_flag);
     // need owned variant because async lifetime
-        let mut file = OpenOptions::new().write(true).open(path).await?;
     let path_buf = path.to_path_buf();
     tokio::spawn(async move {
         let path_buf = path_buf; // moves path buf into async
@@ -53,7 +83,7 @@ pub async fn transparent_file_stream(
             let _ = remove_file(&path_buf).await;
         } else {
             file.flush().await?;
-                file.sync_all().await?;
+            file.sync_all().await?; // we need metadata
         }
 
         let mut write_lock = WRITING_STATUS.write();
@@ -70,8 +100,7 @@ pub async fn transparent_file_stream(
         Ok::<_, std::io::Error>(())
     });
 
-        FromFsStream::new(path, Arc::clone(&done_writing_flag)).await
+    Ok(FromFsStream::new(path, done_writing_flag).await?)
-    }
 }
 
 struct FromFsStream {
@@ -125,6 +154,11 @@ impl CacheStatus {
         Self(AtomicU8::new(WritingStatus::NotDone as u8))
     }
 
+    #[inline]
+    const fn done() -> Self {
+        Self(AtomicU8::new(WritingStatus::Done as u8))
+    }
+
     #[inline]
     fn store(&self, status: WritingStatus) {
         self.0.store(status as u8, Ordering::Release);