refactor: spooled local file and cleanup (#85)

I refactored the code to extract `StreamingAndLocal` into its own file
and add a `SpooledTemporaryFile`. The `SpooledTemporaryFile` keeps data
written to it in memory until it exceeds a certain limit (5MB in this
case).

This PR is in preparation to refactor wheel caching.

Author: baszalmstra

.codespell-whitelist.txt

@@ -1 +1,2 @@
 crate
+ser

crates/rattler_installs_packages/src/http.rs

@@ -16,6 +16,7 @@ use std::str::FromStr;
 use std::sync::Arc;
 use std::time::SystemTime;
 use thiserror::Error;
+use tokio_util::compat::FuturesAsyncReadCompatExt;
 use url::Url;
 
 // Attached to HTTP responses, to make testing easier
@@ -359,6 +360,7 @@ fn body_to_streaming_or_local(
     StreamingOrLocal::Streaming(Box::new(
         stream
             .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
-            .into_async_read(),
+            .into_async_read()
+            .compat(),
     ))
 }

crates/rattler_installs_packages/src/package_database.rs

@@ -267,7 +267,7 @@ impl PackageDb {
                 )
                 .await?;
 
-            let mut bytes = response.into_body().force_local().await.into_diagnostic()?;
+            let mut bytes = response.into_body().into_local().await.into_diagnostic()?;
             let mut source = String::new();
             bytes.read_to_string(&mut source).into_diagnostic()?;
             html::parse_package_names_html(&source)
@@ -301,7 +301,7 @@ impl PackageDb {
         // Turn the response into a seekable response.
         let bytes = artifact_bytes
             .into_body()
-            .force_local()
+            .into_local()
             .await
             .into_diagnostic()?;
         A::new(name.clone(), bytes)

crates/rattler_installs_packages/src/utils.rs

@@ -0,0 +1,39 @@
+mod read_and_seek;
+mod streaming_or_local;
+
+use include_dir::{include_dir, Dir};
+use url::Url;
+
+pub use read_and_seek::ReadAndSeek;
+pub use streaming_or_local::StreamingOrLocal;
+
+/// Keep retrying a certain IO function until it either succeeds or until it doesn't return
+/// [`std::io::ErrorKind::Interrupted`].
+pub fn retry_interrupted<F, T>(mut f: F) -> std::io::Result<T>
+where
+    F: FnMut() -> std::io::Result<T>,
+{
+    loop {
+        match f() {
+            Ok(result) => return Ok(result),
+            Err(err) if err.kind() != std::io::ErrorKind::Interrupted => {
+                return Err(err);
+            }
+            _ => {
+                // Otherwise keep looping!
+            }
+        }
+    }
+}
+
+/// Normalize url according to pip standards
+pub fn normalize_index_url(mut url: Url) -> Url {
+    let path = url.path();
+    if !path.ends_with('/') {
+        url.set_path(&format!("{path}/"));
+    }
+    url
+}
+
+pub(crate) static VENDORED_PACKAGING_DIR: Dir<'_> =
+    include_dir!("$CARGO_MANIFEST_DIR/vendor/packaging/");

crates/rattler_installs_packages/src/utils/mod.rs

@@ -0,0 +1,7 @@
+use std::io::{Read, Seek};
+
+/// Defines that a type can be read and seeked. This trait has a blanket implementation for any type
+/// that implements both [`Read`] and [`Seek`].
+pub trait ReadAndSeek: Read + Seek {}
+
+impl<T> ReadAndSeek for T where T: Read + Seek {}

crates/rattler_installs_packages/src/utils/read_and_seek.rs

@@ -0,0 +1,79 @@
+use crate::utils::ReadAndSeek;
+use futures::TryFutureExt;
+use std::{
+    io,
+    io::{Read, Seek, Write},
+};
+use tempfile::SpooledTempFile;
+use tokio::io::{AsyncRead, AsyncReadExt};
+use tokio::task::JoinError;
+
+/// Represents a stream of data that is either coming in asynchronously from a remote source or from
+/// a synchronous location (like the filesystem).
+///
+/// It is often useful to make this distinction because reading from a remote source is often slower
+/// than reading synchronously (from disk or memory).
+pub enum StreamingOrLocal {
+    /// Represents an asynchronous stream of data.
+    Streaming(Box<dyn AsyncRead + Unpin + Send>),
+
+    /// Represents a synchronous stream of data.
+    Local(Box<dyn ReadAndSeek + Send>),
+}
+
+impl StreamingOrLocal {
+    /// Stream in the contents of the stream and make sure we have a fast locally accessible stream.
+    ///
+    /// If the stream is already local this will simply return that stream. If however the file is
+    /// remote it will first be read to a temporary spooled file.
+    pub async fn into_local(self) -> io::Result<Box<dyn ReadAndSeek + Send>> {
+        match self {
+            StreamingOrLocal::Streaming(mut stream) => {
+                // Create a [`SpooledTempFile`] which is a blob of memory that is kept in memory if
+                // it does not grow beyond 5MB, otherwise it is written to disk.
+                let mut local_file = SpooledTempFile::new(5 * 1024 * 1024);
+
+                // Stream in the bytes and copy them to the temporary file.
+                let mut buf = [0u8; 1024 * 8];
+                loop {
+                    let bytes_read = stream.read(&mut buf).await?;
+                    if bytes_read == 0 {
+                        break;
+                    }
+                    local_file.write_all(&buf)?;
+                }
+
+                // Restart the file from the start so we can start reading from it.
+                local_file.rewind()?;
+
+                Ok(Box::new(local_file))
+            }
+            StreamingOrLocal::Local(stream) => Ok(stream),
+        }
+    }
+
+    /// Asynchronously read the contents of the stream into a vector of bytes.
+    pub async fn read_to_end(self, bytes: &mut Vec<u8>) -> std::io::Result<usize> {
+        match self {
+            StreamingOrLocal::Streaming(mut streaming) => streaming.read_to_end(bytes).await,
+            StreamingOrLocal::Local(mut local) => {
+                match tokio::task::spawn_blocking(move || {
+                    let mut bytes = Vec::new();
+                    local.read_to_end(&mut bytes).map(|_| bytes)
+                })
+                .map_err(JoinError::try_into_panic)
+                .await
+                {
+                    Ok(Ok(result)) => {
+                        *bytes = result;
+                        Ok(bytes.len())
+                    }
+                    Ok(Err(err)) => Err(err),
+                    // Resume the panic on the main task
+                    Err(Ok(panic)) => std::panic::resume_unwind(panic),
+                    Err(Err(_)) => Err(std::io::ErrorKind::Interrupted.into()),
+                }
+            }
+        }
+    }
+}