support streaming content from S3, start using it in the webserver

Author: syphar

Cargo.lock

@@ -64,8 +64,10 @@ derive_more = { version = "2.0.0", features = ["display"] }
 
 # Async
 tokio = { version = "1.0", features = ["rt-multi-thread", "signal", "macros"] }
+tokio-util = { version = "0.7.15", default-features = false, features = ["io"] }
 futures-util = "0.3.5"
 async-stream = "0.3.5"
+async-compression = { version = "0.4.25", features = ["tokio", "bzip2", "zstd", "gzip"] }
 aws-config = "1.0.0"
 aws-sdk-s3 = "1.3.0"
 aws-sdk-cloudfront = "1.3.0"

Cargo.toml

@@ -1,9 +1,9 @@
-use super::{Blob, FileRange};
+use super::{Blob, FileRange, StreamingBlob};
 use crate::{InstanceMetrics, db::Pool, error::Result};
 use chrono::{DateTime, Utc};
 use futures_util::stream::{Stream, TryStreamExt};
 use sqlx::Acquire;
-use std::sync::Arc;
+use std::{io, sync::Arc};
 
 pub(crate) struct DatabaseBackend {
     pool: Pool,
@@ -58,15 +58,14 @@ impl DatabaseBackend {
         }
     }
 
-    pub(super) async fn get(
+    pub(super) async fn get_stream(
         &self,
         path: &str,
-        max_size: usize,
         range: Option<FileRange>,
-    ) -> Result<Blob> {
-        // The maximum size for a BYTEA (the type used for `content`) is 1GB, so this cast is safe:
+    ) -> Result<StreamingBlob> {
+        // The maximum size for a BYTEA (the type used for `content`) is 1GB.
         // https://www.postgresql.org/message-id/162867790712200946i7ba8eb92v908ac595c0c35aee%40mail.gmail.com
-        let max_size = max_size.min(i32::MAX as usize) as i32;
+        let max_size = 2_i32.pow(30);
 
         struct Result {
             path: String,
@@ -126,14 +125,17 @@ impl DatabaseBackend {
             i.try_into()
                 .expect("invalid compression algorithm stored in database")
         });
-        Ok(Blob {
+        let content = result.content.unwrap_or_default();
+        let content_len = content.len();
+        Ok(StreamingBlob {
             path: result.path,
             mime: result
                 .mime
                 .parse()
                 .unwrap_or(mime::APPLICATION_OCTET_STREAM),
             date_updated: result.date_updated,
-            content: result.content.unwrap_or_default(),
+            content: Box::new(io::Cursor::new(content)),
+            content_length: content_len,
             compression,
         })
     }

src/storage/database.rs

@@ -25,14 +25,17 @@ use path_slash::PathExt;
 use std::{
     fmt,
     fs::{self, File},
-    io::{self, BufReader},
+    io::{self, BufReader, Write as _},
     num::ParseIntError,
     ops::RangeInclusive,
     path::{Path, PathBuf},
     sync::Arc,
 };
 use std::{iter, str::FromStr};
-use tokio::{io::AsyncWriteExt, runtime::Runtime};
+use tokio::{
+    io::{AsyncRead, AsyncReadExt as _, AsyncWriteExt},
+    runtime::Runtime,
+};
 use tracing::{error, info_span, instrument, trace};
 use walkdir::WalkDir;
 
@@ -57,6 +60,80 @@ impl Blob {
     }
 }
 
+pub(crate) struct StreamingBlob {
+    pub(crate) path: String,
+    pub(crate) mime: Mime,
+    pub(crate) date_updated: DateTime<Utc>,
+    pub(crate) compression: Option<CompressionAlgorithm>,
+    pub(crate) content_length: usize,
+    pub(crate) content: Box<dyn AsyncRead + Unpin + Send>,
+}
+
+impl std::fmt::Debug for StreamingBlob {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("StreamingBlob")
+            .field("path", &self.path)
+            .field("mime", &self.mime)
+            .field("date_updated", &self.date_updated)
+            .field("compression", &self.compression)
+            .finish()
+    }
+}
+
+impl StreamingBlob {
+    /// wrap the content stream in a streaming decompressor according to the
+    /// algorithm found in `compression` attribute.
+    pub(crate) fn decompress(mut self) -> Self {
+        let Some(alg) = self.compression else {
+            return self;
+        };
+
+        match alg {
+            CompressionAlgorithm::Zstd => {
+                self.content = Box::new(async_compression::tokio::bufread::ZstdDecoder::new(
+                    tokio::io::BufReader::new(self.content),
+                ))
+            }
+            CompressionAlgorithm::Bzip2 => {
+                self.content = Box::new(async_compression::tokio::bufread::BzDecoder::new(
+                    tokio::io::BufReader::new(self.content),
+                ))
+            }
+            CompressionAlgorithm::Gzip => {
+                self.content = Box::new(async_compression::tokio::bufread::GzipDecoder::new(
+                    tokio::io::BufReader::new(self.content),
+                ))
+            }
+        };
+        self.compression = None;
+        self
+    }
+
+    pub(crate) async fn materialize(mut self, max_size: usize) -> Result<Blob> {
+        self = self.decompress();
+
+        let mut content = crate::utils::sized_buffer::SizedBuffer::new(max_size);
+        content.reserve(self.content_length);
+
+        let mut buf = [0u8; 8 * 1024];
+        loop {
+            let n = self.content.read(&mut buf).await?;
+            if n == 0 {
+                break;
+            }
+            content.write_all(&buf[..n])?;
+        }
+
+        Ok(Blob {
+            path: self.path,
+            mime: self.mime,
+            date_updated: self.date_updated,
+            content: content.into_inner(),
+            compression: self.compression,
+        })
+    }
+}
+
 pub fn get_file_list<P: AsRef<Path>>(path: P) -> Box<dyn Iterator<Item = Result<PathBuf>>> {
     let path = path.as_ref().to_path_buf();
     if path.is_file() {
@@ -210,6 +287,35 @@ impl AsyncStorage {
         })
     }
 
+    /// Fetch a rustdoc file from our blob storage.
+    /// * `name` - the crate name
+    /// * `version` - the crate version
+    /// * `latest_build_id` - the id of the most recent build. used purely to invalidate the local archive
+    ///   index cache, when `archive_storage` is `true.` Without it we wouldn't know that we have
+    ///   to invalidate the locally cached file after a rebuild.
+    /// * `path` - the wanted path inside the documentation.
+    /// * `archive_storage` - if `true`, we will assume we have a remove ZIP archive and an index
+    ///    where we can fetch the requested path from inside the ZIP file.
+    #[instrument]
+    pub(crate) async fn stream_rustdoc_file(
+        &self,
+        name: &str,
+        version: &str,
+        latest_build_id: Option<BuildId>,
+        path: &str,
+        archive_storage: bool,
+    ) -> Result<StreamingBlob> {
+        trace!("fetch rustdoc file");
+        Ok(if archive_storage {
+            self.stream_from_archive(&rustdoc_archive_path(name, version), latest_build_id, path)
+                .await?
+        } else {
+            // Add rustdoc prefix, name and version to the path for accessing the file stored in the database
+            let remote_path = format!("rustdoc/{name}/{version}/{path}");
+            self.get_stream(&remote_path).await?
+        })
+    }
+
     #[context("fetching {path} from {name} {version} (archive: {archive_storage})")]
     pub(crate) async fn fetch_source_file(
         &self,
@@ -282,15 +388,16 @@ impl AsyncStorage {
 
     #[instrument]
     pub(crate) async fn get(&self, path: &str, max_size: usize) -> Result<Blob> {
-        let mut blob = match &self.backend {
-            StorageBackend::Database(db) => db.get(path, max_size, None).await,
-            StorageBackend::S3(s3) => s3.get(path, max_size, None).await,
+        self.get_stream(path).await?.materialize(max_size).await
+    }
+
+    #[instrument]
+    pub(crate) async fn get_stream(&self, path: &str) -> Result<StreamingBlob> {
+        let blob = match &self.backend {
+            StorageBackend::Database(db) => db.get_stream(path, None).await,
+            StorageBackend::S3(s3) => s3.get_stream(path, None).await,
         }?;
-        if let Some(alg) = blob.compression {
-            blob.content = decompress(blob.content.as_slice(), alg, max_size)?;
-            blob.compression = None;
-        }
-        Ok(blob)
+        Ok(blob.decompress())
     }
 
     #[instrument]
@@ -301,18 +408,28 @@ impl AsyncStorage {
         range: FileRange,
         compression: Option<CompressionAlgorithm>,
     ) -> Result<Blob> {
+        self.get_range_stream(path, range, compression)
+            .await?
+            .materialize(max_size)
+            .await
+    }
+
+    #[instrument]
+    pub(super) async fn get_range_stream(
+        &self,
+        path: &str,
+        range: FileRange,
+        compression: Option<CompressionAlgorithm>,
+    ) -> Result<StreamingBlob> {
         let mut blob = match &self.backend {
-            StorageBackend::Database(db) => db.get(path, max_size, Some(range)).await,
-            StorageBackend::S3(s3) => s3.get(path, max_size, Some(range)).await,
+            StorageBackend::Database(db) => db.get_stream(path, Some(range)).await,
+            StorageBackend::S3(s3) => s3.get_stream(path, Some(range)).await,
         }?;
         // `compression` represents the compression of the file-stream inside the archive.
         // We don't compress the whole archive, so the encoding of the archive's blob is irrelevant
         // here.
-        if let Some(alg) = compression {
-            blob.content = decompress(blob.content.as_slice(), alg, max_size)?;
-            blob.compression = None;
-        }
-        Ok(blob)
+        blob.compression = compression;
+        Ok(blob.decompress())
     }
 
     #[instrument]
@@ -389,6 +506,38 @@ impl AsyncStorage {
         })
     }
 
+    #[instrument]
+    pub(crate) async fn stream_from_archive(
+        &self,
+        archive_path: &str,
+        latest_build_id: Option<BuildId>,
+        path: &str,
+    ) -> Result<StreamingBlob> {
+        let index_filename = self
+            .download_archive_index(archive_path, latest_build_id)
+            .await?;
+
+        let info = {
+            let path = path.to_owned();
+            spawn_blocking(move || archive_index::find_in_file(index_filename, &path)).await
+        }?
+        .ok_or(PathNotFoundError)?;
+
+        let blob = self
+            .get_range_stream(archive_path, info.range(), Some(info.compression()))
+            .await?;
+        assert_eq!(blob.compression, None);
+
+        Ok(StreamingBlob {
+            path: format!("{archive_path}/{path}"),
+            mime: detect_mime(path),
+            date_updated: blob.date_updated,
+            content: blob.content,
+            content_length: blob.content_length,
+            compression: None,
+        })
+    }
+
     #[instrument(skip(self))]
     pub(crate) async fn store_all_in_archive(
         &self,