Merge pull request #258 from zhaofengli/refactor/asyncbufread

server/upload_path: Get rid of necessary double-buffering

Cargo.lock

@@ -242,6 +242,7 @@ dependencies = [
  "hex",
  "lazy_static",
  "nix-base32",
+ "pin-project",
  "regex",
  "serde",
  "serde_json",

attic/Cargo.toml

@@ -16,6 +16,7 @@ futures = "0.3.31"
 hex = "0.4.3"
 lazy_static = "1.5.0"
 nix-base32 = "0.2.0"
+pin-project = "1.1.10"
 regex = "1.11.1"
 serde = { version = "1.0.219", features = ["derive"] }
 serde_with = "3.14.0"
@@ -53,13 +54,13 @@ nix-main = { version = "2.24", feature = "nix_store" }
 [features]
 default = [
 	"chunking",
+	"io",
 	"nix_store",
-	"stream",
 	"tokio",
 ]
 
 # Chunking.
-chunking = ["tokio", "stream", "dep:async-stream"]
+chunking = ["tokio", "io", "dep:async-stream"]
 
 # Native libnixstore bindings.
 #
@@ -73,8 +74,8 @@ nix_store = [
 	"dep:system-deps",
 ]
 
-# Stream utilities.
-stream = ["tokio", "dep:async-stream"]
+# IO utilities.
+io = ["tokio", "dep:async-stream"]
 
 # Tokio runtime.
 tokio = ["dep:tokio", "tokio/rt", "tokio/time"]

attic/src/chunking/mod.rs

@@ -9,7 +9,7 @@ use fastcdc::ronomon::FastCDC;
 use futures::stream::Stream;
 use tokio::io::AsyncRead;
 
-use crate::stream::read_chunk_async;
+use crate::io::read_chunk_async;
 
 /// Splits a streams into content-defined chunks.
 ///

attic/src/io/hash_reader.rs

@@ -0,0 +1,234 @@
+use std::marker::Unpin;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+use digest::{Digest, Output as DigestOutput};
+use pin_project::pin_project;
+use tokio::io::{self, AsyncBufRead, AsyncRead, ReadBuf};
+use tokio::sync::OnceCell;
+
+/// AsyncRead filter that hashes the bytes that have been read.
+///
+/// The hash is finalized when EOF is reached.
+#[pin_project(project = HashReaderProj)]
+pub struct HashReader<R, D>
+where
+    R: AsyncRead + Unpin,
+    D: Digest + Unpin,
+{
+    #[pin]
+    inner: R,
+    state: State<D>,
+}
+
+struct State<D>
+where
+    D: Digest + Unpin,
+{
+    digest: Option<D>,
+    bytes_hashed: usize,
+    bytes_consumed: usize,
+    finalized: Arc<OnceCell<(DigestOutput<D>, usize)>>,
+}
+
+impl<D> State<D>
+where
+    D: Digest + Unpin,
+{
+    fn hash_unconsumed(&mut self, unconsumed: &[u8]) {
+        let unhashed_offset = self.bytes_hashed - self.bytes_consumed;
+
+        // It's technically possible for the `poll_read`/`poll_fill_buf` implementation
+        // to return less data than the unconsumed portion returned by a previous
+        // call to `AsyncBufRead::poll_fill_buf`.
+        if unhashed_offset < unconsumed.len() {
+            let unhashed = &unconsumed[unhashed_offset..];
+            self.bytes_hashed += unhashed.len();
+
+            let digest = self.digest.as_mut().expect("Stream has data after EOF");
+            digest.update(unhashed);
+        }
+    }
+
+    fn eof(&mut self) {
+        if let Some(digest) = self.digest.take() {
+            assert!(self.bytes_hashed == self.bytes_consumed, "bytes_hashed != bytes_consumed but EOF - Unconsumed bytes disappeared from buffer??");
+            self.finalized
+                .set((digest.finalize(), self.bytes_hashed))
+                .expect("Hash has already been finalized");
+        }
+    }
+}
+
+impl<R, D> HashReader<R, D>
+where
+    R: AsyncRead + Unpin,
+    D: Digest + Unpin,
+{
+    pub fn new(inner: R, digest: D) -> (Self, Arc<OnceCell<(DigestOutput<D>, usize)>>) {
+        let finalized = Arc::new(OnceCell::new());
+
+        (
+            Self {
+                inner,
+                state: State {
+                    digest: Some(digest),
+                    bytes_hashed: 0,
+                    bytes_consumed: 0,
+                    finalized: finalized.clone(),
+                },
+            },
+            finalized,
+        )
+    }
+}
+
+impl<R, D> AsyncRead for HashReader<R, D>
+where
+    R: AsyncRead + Unpin,
+    D: Digest + Unpin,
+{
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let this = self.project();
+
+        let old_filled = buf.filled().len();
+        ready!(this.inner.poll_read(cx, buf))?;
+
+        let filled = buf.filled();
+        let unconsumed = &filled[old_filled..];
+        if unconsumed.len() == 0 {
+            this.state.eof();
+        } else {
+            this.state.hash_unconsumed(unconsumed);
+            this.state.bytes_consumed += unconsumed.len();
+        }
+
+        debug_assert!(this.state.bytes_consumed <= this.state.bytes_hashed);
+        Poll::Ready(Ok(()))
+    }
+}
+
+impl<R, D> AsyncBufRead for HashReader<R, D>
+where
+    R: AsyncBufRead + Unpin,
+    D: Digest + Unpin,
+{
+    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
+        let this = self.project();
+        let unconsumed = ready!(this.inner.poll_fill_buf(cx))?;
+
+        if unconsumed.len() == 0 {
+            this.state.eof();
+        } else {
+            this.state.hash_unconsumed(unconsumed);
+        }
+
+        debug_assert!(this.state.bytes_consumed <= this.state.bytes_hashed);
+        Poll::Ready(Ok(unconsumed))
+    }
+
+    fn consume(self: Pin<&mut Self>, amt: usize) {
+        let this = self.project();
+        this.inner.consume(amt);
+        this.state.bytes_consumed += amt;
+
+        debug_assert!(this.state.bytes_consumed <= this.state.bytes_hashed);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use tokio::io::{AsyncBufReadExt, AsyncReadExt};
+
+    #[tokio::test]
+    async fn test_hash_reader() {
+        let expected = b"hello world";
+        let expected_sha256 =
+            hex::decode("b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9")
+                .unwrap();
+
+        let (mut read, finalized) = HashReader::new(expected.as_slice(), sha2::Sha256::new());
+        assert!(finalized.get().is_none());
+
+        // force multiple reads
+        let mut buf = vec![0u8; 100];
+        let mut bytes_read = 0;
+        bytes_read += read
+            .read(&mut buf[bytes_read..bytes_read + 5])
+            .await
+            .unwrap();
+        bytes_read += read
+            .read(&mut buf[bytes_read..bytes_read + 5])
+            .await
+            .unwrap();
+        bytes_read += read
+            .read(&mut buf[bytes_read..bytes_read + 5])
+            .await
+            .unwrap();
+        bytes_read += read
+            .read(&mut buf[bytes_read..bytes_read + 5])
+            .await
+            .unwrap();
+
+        assert_eq!(expected.len(), bytes_read);
+        assert_eq!(expected, &buf[..bytes_read]);
+
+        let (hash, count) = finalized.get().expect("Hash wasn't finalized");
+
+        assert_eq!(expected_sha256.as_slice(), hash.as_slice());
+        assert_eq!(expected.len(), *count);
+        eprintln!("finalized = {:x?}", finalized);
+    }
+
+    #[tokio::test]
+    async fn test_hash_reader_buf() {
+        let expected = b"hello world";
+        let expected_sha256 =
+            hex::decode("b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9")
+                .unwrap();
+
+        let (mut read, finalized) = HashReader::new(expected.as_slice(), sha2::Sha256::new());
+        assert!(finalized.get().is_none());
+
+        let mut buf = vec![0u8; 100];
+        let mut bytes_read = 0;
+
+        // Mix AsyncRead::read() and AsyncBufRead::fill_buf()
+
+        bytes_read += read
+            .read(&mut buf[bytes_read..bytes_read + 1])
+            .await
+            .unwrap();
+
+        loop {
+            // Perform multiple AsyncBufRead::fill_buf()s _without_ consuming
+            let _ = read.fill_buf().await.unwrap();
+            let _ = read.fill_buf().await.unwrap();
+            let read_buf = read.fill_buf().await.unwrap();
+
+            if read_buf.is_empty() {
+                break;
+            }
+
+            buf[bytes_read] = read_buf[0];
+            read.consume(1);
+            bytes_read += 1;
+        }
+
+        assert_eq!(expected.len(), bytes_read);
+        assert_eq!(expected, &buf[..bytes_read]);
+
+        let (hash, count) = finalized.get().expect("Hash wasn't finalized");
+
+        assert_eq!(expected_sha256.as_slice(), hash.as_slice());
+        assert_eq!(expected.len(), *count);
+        eprintln!("finalized = {:x?}", finalized);
+    }
+}

attic/src/io/mod.rs

@@ -1,29 +1,19 @@
 //! Stream utilities.
 
+mod hash_reader;
+
 use std::collections::VecDeque;
 use std::future::Future;
 use std::marker::Unpin;
 use std::pin::Pin;
-use std::sync::Arc;
-use std::task::{Context, Poll};
 
 use async_stream::try_stream;
 use bytes::{Bytes, BytesMut};
-use digest::{Digest, Output as DigestOutput};
 use futures::stream::{BoxStream, Stream, StreamExt};
-use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};
-use tokio::sync::OnceCell;
+use tokio::io::{AsyncRead, AsyncReadExt};
 use tokio::task::spawn;
 
-/// Stream filter that hashes the bytes that have been read.
-///
-/// The hash is finalized when EOF is reached.
-pub struct StreamHasher<R: AsyncRead + Unpin, D: Digest + Unpin> {
-    inner: R,
-    digest: Option<D>,
-    bytes_read: usize,
-    finalized: Arc<OnceCell<(DigestOutput<D>, usize)>>,
-}
+pub use hash_reader::HashReader;
 
 /// Merge chunks lazily into a continuous stream.
 ///
@@ -98,60 +88,6 @@ where
     Box::pin(s)
 }
 
-impl<R: AsyncRead + Unpin, D: Digest + Unpin> StreamHasher<R, D> {
-    pub fn new(inner: R, digest: D) -> (Self, Arc<OnceCell<(DigestOutput<D>, usize)>>) {
-        let finalized = Arc::new(OnceCell::new());
-
-        (
-            Self {
-                inner,
-                digest: Some(digest),
-                bytes_read: 0,
-                finalized: finalized.clone(),
-            },
-            finalized,
-        )
-    }
-}
-
-impl<R: AsyncRead + Unpin, D: Digest + Unpin> AsyncRead for StreamHasher<R, D> {
-    fn poll_read(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        buf: &mut ReadBuf<'_>,
-    ) -> Poll<tokio::io::Result<()>> {
-        let old_filled = buf.filled().len();
-        let r = Pin::new(&mut self.inner).poll_read(cx, buf);
-        let read_len = buf.filled().len() - old_filled;
-
-        match r {
-            Poll::Ready(Ok(())) => {
-                if read_len == 0 {
-                    // EOF
-                    if let Some(digest) = self.digest.take() {
-                        self.finalized
-                            .set((digest.finalize(), self.bytes_read))
-                            .expect("Hash has already been finalized");
-                    }
-                } else {
-                    // Read something
-                    let digest = self.digest.as_mut().expect("Stream has data after EOF");
-
-                    let filled = buf.filled();
-                    digest.update(&filled[filled.len() - read_len..]);
-                    self.bytes_read += read_len;
-                }
-            }
-            Poll::Ready(Err(_)) => {
-                assert!(read_len == 0);
-            }
-            Poll::Pending => {}
-        }
-
-        r
-    }
-}
-
 /// Greedily reads from a stream to fill a buffer.
 pub async fn read_chunk_async<S: AsyncRead + Unpin + Send>(
     stream: &mut S,
@@ -175,47 +111,6 @@ mod tests {
     use async_stream::stream;
     use bytes::{BufMut, BytesMut};
     use futures::future;
-    use tokio::io::AsyncReadExt;
-
-    #[tokio::test]
-    async fn test_stream_hasher() {
-        let expected = b"hello world";
-        let expected_sha256 =
-            hex::decode("b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9")
-                .unwrap();
-
-        let (mut read, finalized) = StreamHasher::new(expected.as_slice(), sha2::Sha256::new());
-        assert!(finalized.get().is_none());
-
-        // force multiple reads
-        let mut buf = vec![0u8; 100];
-        let mut bytes_read = 0;
-        bytes_read += read
-            .read(&mut buf[bytes_read..bytes_read + 5])
-            .await
-            .unwrap();
-        bytes_read += read
-            .read(&mut buf[bytes_read..bytes_read + 5])
-            .await
-            .unwrap();
-        bytes_read += read
-            .read(&mut buf[bytes_read..bytes_read + 5])
-            .await
-            .unwrap();
-        bytes_read += read
-            .read(&mut buf[bytes_read..bytes_read + 5])
-            .await
-            .unwrap();
-
-        assert_eq!(expected.len(), bytes_read);
-        assert_eq!(expected, &buf[..bytes_read]);
-
-        let (hash, count) = finalized.get().expect("Hash wasn't finalized");
-
-        assert_eq!(expected_sha256.as_slice(), hash.as_slice());
-        assert_eq!(expected.len(), *count);
-        eprintln!("finalized = {:x?}", finalized);
-    }
 
     #[tokio::test]
     async fn test_merge_chunks() {

attic/src/lib.rs

@@ -21,11 +21,11 @@ pub mod cache;
 pub mod chunking;
 pub mod error;
 pub mod hash;
+#[cfg(feature = "io")]
+pub mod io;
 pub mod mime;
 pub mod nix_store;
 pub mod signing;
-#[cfg(feature = "stream")]
-pub mod stream;
 #[cfg(target_family = "unix")]
 pub mod testing;
 #[cfg(feature = "tokio")]

client/src/api/mod.rs

@@ -3,7 +3,7 @@ use std::fmt;
 
 use anyhow::Result;
 use bytes::Bytes;
-use const_format::concatcp;
+use const_format::formatcp;
 use displaydoc::Display;
 use futures::{
     future,
@@ -27,7 +27,7 @@ use attic::nix_store::StorePathHash;
 
 /// The User-Agent string of Attic.
 const ATTIC_USER_AGENT: &str =
-    concatcp!("Attic/{} ({})", env!("CARGO_PKG_NAME"), ATTIC_DISTRIBUTOR);
+    formatcp!("Attic/{} ({})", env!("CARGO_PKG_NAME"), ATTIC_DISTRIBUTOR);
 
 /// The size threshold to send the upload info as part of the PUT body.
 const NAR_INFO_PREAMBLE_THRESHOLD: usize = 4 * 1024; // 4 KiB

flake/devshells.nix

@@ -40,9 +40,11 @@ in
         ];
 
         rust = [
+          cargo-audit
           cargo-expand
           cargo-outdated
           cargo-edit
+          cargo-udeps
           tokio-console
         ];
 
@@ -62,7 +64,6 @@ in
           postgresql
           sqlite-interactive
 
-          flyctl
           skopeo
           manifest-tool
         ];

justfile

@@ -27,7 +27,7 @@ ci-build-wasm:
 	export RUST_MIN_STACK=16777216
 
 	pushd attic
-	cargo build --target wasm32-unknown-unknown --no-default-features -F chunking -F stream
+	cargo build --target wasm32-unknown-unknown --no-default-features -F chunking -F io
 	popd
 	pushd token
 	cargo build --target wasm32-unknown-unknown

server/Cargo.toml

@@ -19,7 +19,7 @@ path = "src/adm/main.rs"
 doc = false
 
 [dependencies]
-attic = { path = "../attic", default-features = false, features = ["chunking", "stream", "tokio"] }
+attic = { path = "../attic", default-features = false, features = ["chunking", "io", "tokio"] }
 attic-token = { path = "../token" }
 
 anyhow = "1.0.98"

server/src/api/binary_cache.rs

@@ -32,9 +32,9 @@ use crate::nix_manifest;
 use crate::storage::{Download, StorageBackend};
 use crate::{RequestState, State};
 use attic::cache::CacheName;
+use attic::io::merge_chunks;
 use attic::mime;
 use attic::nix_store::StorePathHash;
-use attic::stream::merge_chunks;
 
 /// Nix cache information.
 ///

server/src/api/v1/upload_path.rs

@@ -14,7 +14,6 @@ use axum::{
 };
 use bytes::{Bytes, BytesMut};
 use chrono::Utc;
-use digest::Output as DigestOutput;
 use futures::future::join_all;
 use futures::StreamExt;
 use sea_orm::entity::prelude::*;
@@ -22,13 +21,14 @@ use sea_orm::sea_query::Expr;
 use sea_orm::ActiveValue::Set;
 use sea_orm::{QuerySelect, TransactionTrait};
 use sha2::{Digest, Sha256};
-use tokio::io::{AsyncBufRead, AsyncRead, AsyncReadExt, BufReader};
-use tokio::sync::{OnceCell, Semaphore};
+use tokio::io::{AsyncBufRead, AsyncReadExt};
+use tokio::sync::Semaphore;
 use tokio::task::spawn;
 use tokio_util::io::StreamReader;
 use tracing::instrument;
 use uuid::Uuid;
 
+use crate::compression::{CompressionStream, CompressorFn};
 use crate::config::CompressionType;
 use crate::error::{ErrorKind, ServerError, ServerResult};
 use crate::narinfo::Compression;
@@ -39,7 +39,7 @@ use attic::api::v1::upload_path::{
 };
 use attic::chunking::chunk_stream;
 use attic::hash::Hash;
-use attic::stream::{read_chunk_async, StreamHasher};
+use attic::io::{read_chunk_async, HashReader};
 use attic::util::Finally;
 
 use crate::database::entity::cache;
@@ -55,15 +55,13 @@ use crate::database::{AtticDatabase, ChunkGuard, NarGuard};
 /// TODO: Make this configurable
 const CONCURRENT_CHUNK_UPLOADS: usize = 10;
 
-type CompressorFn<C> = Box<dyn FnOnce(C) -> Box<dyn AsyncRead + Unpin + Send> + Send>;
-
 /// Data of a chunk.
 enum ChunkData {
     /// Some bytes in memory.
     Bytes(Bytes),
 
     /// A stream with a user-claimed hash and size that are potentially incorrect.
-    Stream(Box<dyn AsyncRead + Send + Unpin + 'static>, Hash, usize),
+    Stream(Box<dyn AsyncBufRead + Send + Unpin + 'static>, Hash, usize),
 }
 
 /// Result of a chunk upload.
@@ -72,33 +70,6 @@ struct UploadChunkResult {
     deduplicated: bool,
 }
 
-/// Applies compression to a stream, computing hashes along the way.
-///
-/// Our strategy is to stream directly onto a UUID-keyed file on the
-/// storage backend, performing compression and computing the hashes
-/// along the way. We delete the file if the hashes do not match.
-///
-/// ```text
-///                    ┌───────────────────────────────────►NAR Hash
-///                    │
-///                    │
-///                    ├───────────────────────────────────►NAR Size
-///                    │
-///              ┌─────┴────┐  ┌──────────┐  ┌───────────┐
-/// NAR Stream──►│NAR Hasher├─►│Compressor├─►│File Hasher├─►File Stream
-///              └──────────┘  └──────────┘  └─────┬─────┘
-///                                                │
-///                                                ├───────►File Hash
-///                                                │
-///                                                │
-///                                                └───────►File Size
-/// ```
-struct CompressionStream {
-    stream: Box<dyn AsyncRead + Unpin + Send>,
-    nar_compute: Arc<OnceCell<(DigestOutput<Sha256>, usize)>>,
-    file_compute: Arc<OnceCell<(DigestOutput<Sha256>, usize)>>,
-}
-
 trait UploadPathNarInfoExt {
     fn to_active_model(&self) -> object::ActiveModel;
 }
@@ -180,40 +151,33 @@ pub(crate) async fn upload_path(
     let username = req_state.auth.username().map(str::to_string);
 
     // Try to acquire a lock on an existing NAR
-    let existing_nar = database.find_and_lock_nar(&upload_info.nar_hash).await?;
-    match existing_nar {
-        Some(existing_nar) => {
-            // Deduplicate?
-            let missing_chunk = ChunkRef::find()
-                .filter(chunkref::Column::NarId.eq(existing_nar.id))
-                .filter(chunkref::Column::ChunkId.is_null())
-                .limit(1)
-                .one(database)
-                .await
-                .map_err(ServerError::database_error)?;
+    if let Some(existing_nar) = database.find_and_lock_nar(&upload_info.nar_hash).await? {
+        // Deduplicate?
+        let missing_chunk = ChunkRef::find()
+            .filter(chunkref::Column::NarId.eq(existing_nar.id))
+            .filter(chunkref::Column::ChunkId.is_null())
+            .limit(1)
+            .one(database)
+            .await
+            .map_err(ServerError::database_error)?;
 
-            if missing_chunk.is_some() {
-                // Need to repair
-                upload_path_new(username, cache, upload_info, stream, database, &state).await
-            } else {
-                // Can actually be deduplicated
-                upload_path_dedup(
-                    username,
-                    cache,
-                    upload_info,
-                    stream,
-                    database,
-                    &state,
-                    existing_nar,
-                )
-                .await
-            }
-        }
-        None => {
-            // New NAR
-            upload_path_new(username, cache, upload_info, stream, database, &state).await
+        if missing_chunk.is_none() {
+            // Can actually be deduplicated
+            return upload_path_dedup(
+                username,
+                cache,
+                upload_info,
+                stream,
+                database,
+                &state,
+                existing_nar,
+            )
+            .await;
         }
     }
+
+    // New NAR or need to repair
+    upload_path_new(username, cache, upload_info, stream, database, &state).await
 }
 
 /// Uploads a path when there is already a matching NAR in the global cache.
@@ -221,13 +185,13 @@ async fn upload_path_dedup(
     username: Option<String>,
     cache: cache::Model,
     upload_info: UploadPathNarInfo,
-    stream: impl AsyncRead + Unpin,
+    stream: impl AsyncBufRead + Unpin,
     database: &DatabaseConnection,
     state: &State,
     existing_nar: NarGuard,
 ) -> ServerResult<Json<UploadPathResult>> {
     if state.config.require_proof_of_possession {
-        let (mut stream, nar_compute) = StreamHasher::new(stream, Sha256::new());
+        let (mut stream, nar_compute) = HashReader::new(stream, Sha256::new());
         tokio::io::copy(&mut stream, &mut tokio::io::sink())
             .await
             .map_err(ServerError::request_error)?;
@@ -301,7 +265,7 @@ async fn upload_path_new(
     username: Option<String>,
     cache: cache::Model,
     upload_info: UploadPathNarInfo,
-    stream: impl AsyncRead + Send + Unpin + 'static,
+    stream: impl AsyncBufRead + Send + Unpin + 'static,
     database: &DatabaseConnection,
     state: &State,
 ) -> ServerResult<Json<UploadPathResult>> {
@@ -319,7 +283,7 @@ async fn upload_path_new_chunked(
     username: Option<String>,
     cache: cache::Model,
     upload_info: UploadPathNarInfo,
-    stream: impl AsyncRead + Send + Unpin + 'static,
+    stream: impl AsyncBufRead + Send + Unpin + 'static,
     database: &DatabaseConnection,
     state: &State,
 ) -> ServerResult<Json<UploadPathResult>> {
@@ -371,7 +335,7 @@ async fn upload_path_new_chunked(
     });
 
     let stream = stream.take(upload_info.nar_size as u64);
-    let (stream, nar_compute) = StreamHasher::new(stream, Sha256::new());
+    let (stream, nar_compute) = HashReader::new(stream, Sha256::new());
     let mut chunks = chunk_stream(
         stream,
         chunking_config.min_size,
@@ -510,7 +474,7 @@ async fn upload_path_new_unchunked(
     username: Option<String>,
     cache: cache::Model,
     upload_info: UploadPathNarInfo,
-    stream: impl AsyncRead + Send + Unpin + 'static,
+    stream: impl AsyncBufRead + Send + Unpin + 'static,
     database: &DatabaseConnection,
     state: &State,
 ) -> ServerResult<Json<UploadPathResult>> {
@@ -623,9 +587,9 @@ async fn upload_chunk(
     {
         // There's an existing chunk matching the hash
         if require_proof_of_possession && !data.is_hash_trusted() {
-            let stream = data.into_async_read();
+            let stream = data.into_async_buf_read();
 
-            let (mut stream, nar_compute) = StreamHasher::new(stream, Sha256::new());
+            let (mut stream, nar_compute) = HashReader::new(stream, Sha256::new());
             tokio::io::copy(&mut stream, &mut tokio::io::sink())
                 .await
                 .map_err(ServerError::request_error)?;
@@ -705,7 +669,7 @@ async fn upload_chunk(
 
     // Compress and stream to the storage backend
     let compressor = get_compressor_fn(compression_type, compression_level);
-    let mut stream = CompressionStream::new(data.into_async_read(), compressor);
+    let mut stream = CompressionStream::new(data.into_async_buf_read(), compressor);
 
     backend
         .upload_file(key, stream.stream())
@@ -809,8 +773,8 @@ impl ChunkData {
         matches!(self, ChunkData::Bytes(_))
     }
 
-    /// Turns the data into a stream.
-    fn into_async_read(self) -> Box<dyn AsyncRead + Unpin + Send> {
+    /// Turns the data into an AsyncBufRead.
+    fn into_async_buf_read(self) -> Box<dyn AsyncBufRead + Unpin + Send> {
         match self {
             Self::Bytes(bytes) => Box::new(Cursor::new(bytes)),
             Self::Stream(stream, _, _) => stream,
@@ -818,73 +782,6 @@ impl ChunkData {
     }
 }
 
-impl CompressionStream {
-    /// Creates a new compression stream.
-    fn new<R>(stream: R, compressor: CompressorFn<BufReader<StreamHasher<R, Sha256>>>) -> Self
-    where
-        R: AsyncRead + Unpin + Send + 'static,
-    {
-        // compute NAR hash and size
-        let (stream, nar_compute) = StreamHasher::new(stream, Sha256::new());
-
-        // compress NAR
-        let stream = compressor(BufReader::new(stream));
-
-        // compute file hash and size
-        let (stream, file_compute) = StreamHasher::new(stream, Sha256::new());
-
-        Self {
-            stream: Box::new(stream),
-            nar_compute,
-            file_compute,
-        }
-    }
-
-    /*
-    /// Creates a compression stream without compute the uncompressed hash/size.
-    ///
-    /// This is useful if you already know the hash. `nar_hash_and_size` will
-    /// always return `None`.
-    fn new_without_nar_hash<R>(stream: R, compressor: CompressorFn<BufReader<R>>) -> Self
-    where
-        R: AsyncRead + Unpin + Send + 'static,
-    {
-        // compress NAR
-        let stream = compressor(BufReader::new(stream));
-
-        // compute file hash and size
-        let (stream, file_compute) = StreamHasher::new(stream, Sha256::new());
-
-        Self {
-            stream: Box::new(stream),
-            nar_compute: Arc::new(OnceCell::new()),
-            file_compute,
-        }
-    }
-    */
-
-    /// Returns the stream of the compressed object.
-    fn stream(&mut self) -> &mut (impl AsyncRead + Unpin) {
-        &mut self.stream
-    }
-
-    /// Returns the NAR hash and size.
-    ///
-    /// The hash is only finalized when the stream is fully read.
-    /// Otherwise, returns `None`.
-    fn nar_hash_and_size(&self) -> Option<&(DigestOutput<Sha256>, usize)> {
-        self.nar_compute.get()
-    }
-
-    /// Returns the file hash and size.
-    ///
-    /// The hash is only finalized when the stream is fully read.
-    /// Otherwise, returns `None`.
-    fn file_hash_and_size(&self) -> Option<&(DigestOutput<Sha256>, usize)> {
-        self.file_compute.get()
-    }
-}
-
 impl UploadPathNarInfoExt for UploadPathNarInfo {
     fn to_active_model(&self) -> object::ActiveModel {
         object::ActiveModel {

server/src/compression.rs

@@ -0,0 +1,81 @@
+use std::sync::Arc;
+
+use digest::Output as DigestOutput;
+use sha2::{Digest, Sha256};
+use tokio::io::{AsyncBufRead, AsyncRead};
+use tokio::sync::OnceCell;
+
+use attic::io::HashReader;
+
+pub type CompressorFn<C> = Box<dyn FnOnce(C) -> Box<dyn AsyncRead + Unpin + Send> + Send>;
+
+/// Applies compression to a stream, computing hashes along the way.
+///
+/// Our strategy is to stream directly onto a UUID-keyed file on the
+/// storage backend, performing compression and computing the hashes
+/// along the way. We delete the file if the hashes do not match.
+///
+/// ```text
+///                    ┌───────────────────────────────────►NAR Hash
+///                    │
+///                    │
+///                    ├───────────────────────────────────►NAR Size
+///                    │
+///              ┌─────┴────┐  ┌──────────┐  ┌───────────┐
+/// NAR Stream──►│NAR Hasher├─►│Compressor├─►│File Hasher├─►File Stream
+///              └──────────┘  └──────────┘  └─────┬─────┘
+///                                                │
+///                                                ├───────►File Hash
+///                                                │
+///                                                │
+///                                                └───────►File Size
+/// ```
+pub struct CompressionStream {
+    stream: Box<dyn AsyncRead + Unpin + Send>,
+    nar_compute: Arc<OnceCell<(DigestOutput<Sha256>, usize)>>,
+    file_compute: Arc<OnceCell<(DigestOutput<Sha256>, usize)>>,
+}
+
+impl CompressionStream {
+    /// Creates a new compression stream.
+    pub fn new<R>(stream: R, compressor: CompressorFn<HashReader<R, Sha256>>) -> Self
+    where
+        R: AsyncBufRead + Unpin + Send + 'static,
+    {
+        // compute NAR hash and size
+        let (stream, nar_compute) = HashReader::new(stream, Sha256::new());
+
+        // compress NAR
+        let stream = compressor(stream);
+
+        // compute file hash and size
+        let (stream, file_compute) = HashReader::new(stream, Sha256::new());
+
+        Self {
+            stream: Box::new(stream),
+            nar_compute,
+            file_compute,
+        }
+    }
+
+    /// Returns the stream of the compressed object.
+    pub fn stream(&mut self) -> &mut (impl AsyncRead + Unpin) {
+        &mut self.stream
+    }
+
+    /// Returns the NAR hash and size.
+    ///
+    /// The hash is only finalized when the stream is fully read.
+    /// Otherwise, returns `None`.
+    pub fn nar_hash_and_size(&self) -> Option<&(DigestOutput<Sha256>, usize)> {
+        self.nar_compute.get()
+    }
+
+    /// Returns the file hash and size.
+    ///
+    /// The hash is only finalized when the stream is fully read.
+    /// Otherwise, returns `None`.
+    pub fn file_hash_and_size(&self) -> Option<&(DigestOutput<Sha256>, usize)> {
+        self.file_compute.get()
+    }
+}

server/src/lib.rs

@@ -15,6 +15,7 @@
 
 pub mod access;
 mod api;
+mod compression;
 pub mod config;
 pub mod database;
 pub mod error;

server/src/storage/s3.rs

@@ -19,7 +19,7 @@ use tokio::io::AsyncRead;
 
 use super::{Download, RemoteFile, StorageBackend};
 use crate::error::{ErrorKind, ServerError, ServerResult};
-use attic::stream::read_chunk_async;
+use attic::io::read_chunk_async;
 use attic::util::Finally;
 
 /// The chunk size for each part in a multipart upload.