switch stream-upload api to use native SDK file-upload, calculate checksum

Author: syphar

.gitignore

@@ -1,4 +1,4 @@
-/ignored
+*/ignored
 /.env
 /.docker.env
 /src/web/badge/Cargo.lock

Cargo.lock

@@ -23,8 +23,10 @@ async-stream = { workspace = true }
 aws-config = { version = "1.0.0", default-features = false, features = ["default-https-client", "rt-tokio"] }
 aws-sdk-s3 = { version = "1.3.0", default-features = false, features = ["default-https-client", "rt-tokio"] }
 aws-smithy-types-convert = { version = "0.60.0", features = ["convert-chrono"] }
+base64 = { workspace = true }
 bzip2 = "0.6.0"
 chrono = { workspace = true }
+crc32fast = "1.4.2"
 dashmap = { version = "6.0.0", optional = true }
 docs_rs_config = { path = "../docs_rs_config" }
 docs_rs_env_vars = { path = "../docs_rs_env_vars" }
@@ -37,8 +39,6 @@ docs_rs_utils = { path = "../docs_rs_utils" }
 flate2 = "1.1.1"
 futures-util = { workspace = true }
 http = { workspace = true }
-http-body = "1.0.0"
-http-body-util = "0.1.3"
 itertools = { workspace = true }
 mime = { workspace = true }
 moka = { version = "0.12.14", features = ["future"] }
@@ -73,5 +73,9 @@ name = "archive_index_cache"
 harness = false
 required-features = ["testing"]
 
+[[bench]]
+name = "crc32"
+harness = false
+
 [lints]
 workspace = true

crates/lib/docs_rs_storage/Cargo.toml

@@ -0,0 +1,20 @@
+use criterion::{Criterion, Throughput, criterion_group, criterion_main};
+use docs_rs_storage::crc32_for_path;
+use std::{fs, hint::black_box};
+
+pub fn crc32_file(c: &mut Criterion) {
+    let fixture_path = tempfile::NamedTempFile::new().unwrap().into_temp_path();
+
+    let fixture = vec![b'x'; 16 * 1024 * 1024];
+    fs::write(&fixture_path, &fixture).unwrap();
+
+    let mut group = c.benchmark_group("crc32");
+    group.throughput(Throughput::Bytes(fixture.len() as u64));
+    group.bench_function("file_16mib", |b| {
+        b.iter(|| crc32_for_path(black_box(&fixture_path)));
+    });
+    group.finish();
+}
+
+criterion_group!(crc32_benches, crc32_file);
+criterion_main!(crc32_benches);

crates/lib/docs_rs_storage/benches/crc32.rs

@@ -1,7 +1,7 @@
 use crate::{
     Blob,
     backends::StorageBackendMethods,
-    blob::{StreamUpload, StreamingBlob},
+    blob::{StreamUpload, StreamUploadSource, StreamingBlob},
     errors::PathNotFoundError,
     metrics::StorageMetrics,
     types::FileRange,
@@ -12,7 +12,7 @@ use dashmap::DashMap;
 use docs_rs_headers::compute_etag;
 use futures_util::stream::{self, BoxStream};
 use itertools::Itertools as _;
-use tokio::io;
+use tokio::fs;
 
 pub(crate) struct MemoryBackend {
     otel_metrics: StorageMetrics,
@@ -56,16 +56,17 @@ impl StorageBackendMethods for MemoryBackend {
             compression,
         } = upload;
 
-        let mut content = source.reader().await?;
-        let mut buffer = Vec::new();
-        io::copy(&mut content, &mut buffer).await?;
+        let content = match source {
+            StreamUploadSource::Bytes(content) => content.to_vec(),
+            StreamUploadSource::File(path) => fs::read(&path).await?,
+        };
 
         let blob = Blob {
             path,
             mime,
             date_updated: Utc::now(),
-            etag: Some(compute_etag(&buffer)),
-            content: buffer,
+            etag: Some(compute_etag(&content)),
+            content,
             compression,
         };
 

crates/lib/docs_rs_storage/src/backends/memory.rs

@@ -1,7 +1,8 @@
 use crate::{
     Config,
     backends::StorageBackendMethods,
-    blob::{StreamUpload, StreamingBlob},
+    blob::{StreamUpload, StreamUploadSource, StreamingBlob},
+    crc32_for_path,
     errors::PathNotFoundError,
     metrics::StorageMetrics,
     types::FileRange,
@@ -13,20 +14,17 @@ use aws_sdk_s3::{
     Client,
     config::{Region, retry::RetryConfig},
     error::{ProvideErrorMetadata, SdkError},
-    primitives::ByteStream,
-    types::{Delete, ObjectIdentifier},
+    primitives::{ByteStream, Length},
+    types::{ChecksumAlgorithm, Delete, ObjectIdentifier},
 };
 use aws_smithy_types_convert::date_time::DateTimeExt;
+use base64::{Engine as _, engine::general_purpose::STANDARD as b64};
 use chrono::Utc;
 use docs_rs_headers::{ETag, compute_etag};
-use futures_util::{
-    TryStreamExt,
-    stream::{BoxStream, StreamExt},
-};
-use http_body::Frame;
-use http_body_util::StreamBody;
+use docs_rs_utils::spawn_blocking;
+use futures_util::stream::{BoxStream, StreamExt};
 use opentelemetry::KeyValue;
-use tokio_util::io::ReaderStream;
+use tokio::fs;
 use tracing::{error, warn};
 
 // error codes to check for when trying to determine if an error is
@@ -46,6 +44,8 @@ static NOT_FOUND_ERROR_CODES: [&str; 5] = [
     "XMinioInvalidObjectName",
 ];
 
+const S3_UPLOAD_BUFFER_SIZE: usize = 1024 * 1024; // 1 MiB
+
 trait S3ResultExt<T> {
     fn convert_errors(self) -> anyhow::Result<T>;
 }
@@ -243,34 +243,74 @@ impl StorageBackendMethods for S3Backend {
             compression,
         } = upload;
 
-        let content_length = source.content_length().await?;
+        let (content_length, checksum_crc32) = match &source {
+            StreamUploadSource::Bytes(bytes) => (bytes.len() as u64, None),
+            StreamUploadSource::File(local_path) => {
+                let local_path = local_path.clone();
+
+                (
+                    fs::metadata(&local_path).await?.len(),
+                    Some(
+                        spawn_blocking(move || Ok(b64.encode(crc32_for_path(local_path)?))).await?,
+                    ),
+                )
+            }
+        };
 
         let mut last_err = None;
 
         for attempt in 1..=3 {
-            let reader = source.reader().await?;
-            let stream = ReaderStream::new(reader).map_ok(Frame::data);
+            let body = match &source {
+                StreamUploadSource::Bytes(bytes) => ByteStream::from(bytes.clone()),
+                StreamUploadSource::File(path) => {
+                    // NOTE:
+                    // reading the upload-data from a local path is
+                    // "retryable" in the AWS SDK sense.
+                    // ".file" (file pointer) is not retryable.
+                    ByteStream::read_from()
+                        .path(path)
+                        .buffer_size(S3_UPLOAD_BUFFER_SIZE)
+                        .length(Length::Exact(content_length))
+                        .build()
+                        .await?
+                }
+            };
 
-            match self
+            let mut request = self
                 .client
                 .put_object()
                 .bucket(&self.bucket)
                 .key(&path)
-                .body(ByteStream::from_body_1_x(StreamBody::new(stream)))
+                .body(body)
                 .content_length(content_length as i64)
                 .content_type(mime.to_string())
-                .set_content_encoding(compression.map(|alg| alg.to_string()))
-                .send()
-                .await
-            {
+                .set_content_encoding(compression.map(|alg| alg.to_string()));
+
+            // NOTE: when you try to stream-upload a local file, the AWS SDK by default
+            // uses a "middleware" to calculate the checksum for the content, to compare it after
+            // uploading.
+            // This piece is broken right now, but only when using S3 directly. On minio, all is
+            // fiine.
+            // I don't want to disable checksums so we're sure the files are uploaded correctly.
+            // So the only alternative (outside of trying to fix the SDK) is to calculate the
+            // checksum ourselves. This is a little annoying because this means we have to read the
+            // whole file before upload. But since I don't want to load all files into memory before
+            // upload, this is the only option.
+            if let Some(checksum_crc32) = &checksum_crc32 {
+                request = request
+                    .checksum_algorithm(ChecksumAlgorithm::Crc32)
+                    .checksum_crc32(checksum_crc32);
+            }
+
+            match request.send().await {
                 Ok(_) => {
                     self.otel_metrics
                         .uploaded_files
                         .add(1, &[KeyValue::new("attempt", attempt.to_string())]);
                     return Ok(());
                 }
                 Err(err) => {
-                    warn!(?err, attempt = attempt + 1, %path, "failed to upload blob to S3");
+                    warn!(?err, attempt, %path, "failed to upload blob to S3");
                     last_err = Some(err);
                 }
             }

crates/lib/docs_rs_storage/src/backends/s3.rs

@@ -4,39 +4,13 @@ use chrono::{DateTime, Utc};
 use docs_rs_headers::{ETag, compute_etag};
 use docs_rs_types::CompressionAlgorithm;
 use mime::Mime;
-use std::{
-    fmt,
-    io::{Cursor, SeekFrom},
-    sync::Arc,
-};
-use tokio::{
-    fs,
-    io::{self, AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncSeekExt},
-};
+use std::{fmt, io::Cursor, path::PathBuf};
+use tokio::io::{self, AsyncBufRead, AsyncBufReadExt};
+use tokio_util::bytes::Bytes;
 
 pub enum StreamUploadSource {
-    Bytes(Arc<[u8]>),
-    File(fs::File),
-}
-
-impl StreamUploadSource {
-    pub async fn reader(&self) -> io::Result<Box<dyn AsyncRead + Unpin + Send + Sync>> {
-        Ok(match self {
-            Self::Bytes(bytes) => Box::new(Cursor::new(bytes.clone())),
-            Self::File(file) => {
-                let mut cloned = file.try_clone().await?;
-                cloned.seek(SeekFrom::Start(0)).await?;
-                Box::new(cloned)
-            }
-        })
-    }
-
-    pub async fn content_length(&self) -> io::Result<u64> {
-        Ok(match self {
-            Self::Bytes(bytes) => bytes.len() as u64,
-            Self::File(file) => file.metadata().await?.len(),
-        })
-    }
+    Bytes(Bytes),
+    File(PathBuf),
 }
 
 /// Represents a stream blob to be uploaded to storage.
@@ -59,7 +33,7 @@ impl From<BlobUpload> for StreamUpload {
         Self {
             path: value.path,
             mime: value.mime,
-            source: StreamUploadSource::Bytes(Arc::from(value.content)),
+            source: StreamUploadSource::Bytes(value.content.into()),
             compression: value.compression,
         }
     }
@@ -200,10 +174,6 @@ mod test {
     use super::*;
     use crate::compress_async;
     use docs_rs_headers::compute_etag;
-    use tokio::{
-        fs,
-        io::{AsyncReadExt as _, AsyncWriteExt as _},
-    };
 
     const ZSTD_EOF_BYTES: [u8; 3] = [0x01, 0x00, 0x00];
 
@@ -246,52 +216,6 @@ mod test {
         Ok(())
     }
 
-    #[tokio::test]
-    async fn test_stream_upload_source_bytes_creates_fresh_readers() -> Result<()> {
-        const CONTENT: &[u8] = b"Hello, world!";
-
-        let source = StreamUploadSource::Bytes(Arc::from(CONTENT));
-        assert_eq!(source.content_length().await?, CONTENT.len() as u64);
-
-        let mut first = source.reader().await?;
-        let mut first_buf = Vec::new();
-        first.read_to_end(&mut first_buf).await?;
-        assert_eq!(first_buf, CONTENT);
-
-        let mut second = source.reader().await?;
-        let mut second_buf = Vec::new();
-        second.read_to_end(&mut second_buf).await?;
-        assert_eq!(second_buf, CONTENT);
-
-        Ok(())
-    }
-
-    #[tokio::test]
-    async fn test_stream_upload_source_file_creates_fresh_readers() -> Result<()> {
-        const CONTENT: &[u8] = b"Hello, world!";
-
-        let tempfile = tempfile::NamedTempFile::new()?;
-        let mut file = fs::File::from_std(tempfile.reopen()?);
-        file.write_all(CONTENT).await?;
-        file.seek(std::io::SeekFrom::Start(CONTENT.len() as u64))
-            .await?;
-
-        let source = StreamUploadSource::File(file);
-        assert_eq!(source.content_length().await?, CONTENT.len() as u64);
-
-        let mut first = source.reader().await?;
-        let mut first_buf = Vec::new();
-        first.read_to_end(&mut first_buf).await?;
-        assert_eq!(first_buf, CONTENT);
-
-        let mut second = source.reader().await?;
-        let mut second_buf = Vec::new();
-        second.read_to_end(&mut second_buf).await?;
-        assert_eq!(second_buf, CONTENT);
-
-        Ok(())
-    }
-
     #[tokio::test]
     async fn test_streaming_broken_zstd_blob() -> Result<()> {
         const NOT_ZSTD: &[u8] = b"Hello, world!";

crates/lib/docs_rs_storage/src/blob.rs

@@ -75,8 +75,6 @@ impl AppConfig for ArchiveIndexCacheConfig {
 
 #[derive(Debug)]
 pub struct Config {
-    pub temp_dir: PathBuf,
-
     // Storage params
     pub storage_backend: StorageKind,
 
@@ -117,7 +115,6 @@ impl AppConfig for Config {
         let cores = std::thread::available_parallelism()?.get();
 
         Ok(Self {
-            temp_dir: prefix.join("tmp"),
             storage_backend: env("DOCSRS_STORAGE_BACKEND", StorageKind::default())?,
             aws_sdk_max_retries: env("DOCSRS_AWS_SDK_MAX_RETRIES", 6u32)?,
             s3_bucket: env("DOCSRS_S3_BUCKET", "rust-docs-rs".to_string())?,

crates/lib/docs_rs_storage/src/config.rs

@@ -22,6 +22,7 @@ pub use storage::blocking::Storage;
 pub use storage::non_blocking::AsyncStorage;
 pub use types::StorageKind;
 pub use utils::{
+    crc32::crc32_for_path,
     file_list::get_file_list,
     storage_path::{rustdoc_archive_path, rustdoc_json_path, source_archive_path},
 };