fix(remote_storage): buffering, by using streams for upload and download (#5446)

There is double buffering in remote_storage and in pageserver for 8KiB
in using `tokio::io::copy` to read `BufReader<ReaderStream<_>>`.

Switches downloads and uploads to use `Stream<Item =
std::io::Result<Bytes>>`. Caller and only caller now handles setting up
buffering. For reading, `Stream<Item = ...>` is also a `AsyncBufRead`,
so when writing to a file, we now have `tokio::io::copy_buf` reading
full buffers and writing them to `tokio::io::BufWriter` which handles
the buffering before dispatching over to `tokio::fs::File`.

Additionally implements streaming uploads for azure. With azure
downloads are a bit nicer than before, but not much; instead of one huge
vec they just hold on to N allocations we got over the wire.

This PR will also make it trivial to switch reading and writing to
io-uring based methods.

Cc: #5563.

pageserver/src/tenant/delete.rs

@@ -77,8 +77,10 @@ async fn create_remote_delete_mark(
     let data: &[u8] = &[];
     backoff::retry(
         || async {
+            let data = bytes::Bytes::from_static(data);
+            let stream = futures::stream::once(futures::future::ready(Ok(data)));
             remote_storage
-                .upload(data, 0, &remote_mark_path, None)
+                .upload(stream, 0, &remote_mark_path, None)
                 .await
         },
         |_e| false,

pageserver/src/tenant/remote_timeline_client/download.rs

@@ -75,12 +75,11 @@ pub async fn download_layer_file<'a>(
 
     let (mut destination_file, bytes_amount) = download_retry(
         || async {
-            // TODO: this doesn't use the cached fd for some reason?
-            let mut destination_file = fs::File::create(&temp_file_path)
+            let destination_file = tokio::fs::File::create(&temp_file_path)
                 .await
                 .with_context(|| format!("create a destination file for layer '{temp_file_path}'"))
                 .map_err(DownloadError::Other)?;
-            let mut download = storage
+            let download = storage
                 .download(&remote_path)
                 .await
                 .with_context(|| {
@@ -90,9 +89,14 @@ pub async fn download_layer_file<'a>(
                 })
                 .map_err(DownloadError::Other)?;
 
+            let mut destination_file =
+                tokio::io::BufWriter::with_capacity(8 * 1024, destination_file);
+
+            let mut reader = tokio_util::io::StreamReader::new(download.download_stream);
+
             let bytes_amount = tokio::time::timeout(
                 MAX_DOWNLOAD_DURATION,
-                tokio::io::copy(&mut download.download_stream, &mut destination_file),
+                tokio::io::copy_buf(&mut reader, &mut destination_file),
             )
             .await
             .map_err(|e| DownloadError::Other(anyhow::anyhow!("Timed out  {:?}", e)))?
@@ -103,6 +107,8 @@ pub async fn download_layer_file<'a>(
             })
             .map_err(DownloadError::Other)?;
 
+            let destination_file = destination_file.into_inner();
+
             Ok((destination_file, bytes_amount))
         },
         &format!("download {remote_path:?}"),
@@ -220,20 +226,22 @@ async fn do_download_index_part(
     index_generation: Generation,
     cancel: CancellationToken,
 ) -> Result<IndexPart, DownloadError> {
+    use futures::stream::StreamExt;
+
     let remote_path = remote_index_path(tenant_shard_id, timeline_id, index_generation);
 
     let index_part_bytes = download_retry_forever(
         || async {
-            let mut index_part_download = storage.download(&remote_path).await?;
+            let index_part_download = storage.download(&remote_path).await?;
 
             let mut index_part_bytes = Vec::new();
-            tokio::io::copy(
-                &mut index_part_download.download_stream,
-                &mut index_part_bytes,
-            )
-            .await
-            .with_context(|| format!("download index part at {remote_path:?}"))
-            .map_err(DownloadError::Other)?;
+            let mut stream = std::pin::pin!(index_part_download.download_stream);
+            while let Some(chunk) = stream.next().await {
+                let chunk = chunk
+                    .with_context(|| format!("download index part at {remote_path:?}"))
+                    .map_err(DownloadError::Other)?;
+                index_part_bytes.extend_from_slice(&chunk[..]);
+            }
             Ok(index_part_bytes)
         },
         &format!("download {remote_path:?}"),
@@ -398,7 +406,7 @@ pub(crate) async fn download_initdb_tar_zst(
 
     let file = download_retry(
         || async {
-            let mut file = OpenOptions::new()
+            let file = OpenOptions::new()
                 .create(true)
                 .truncate(true)
                 .read(true)
@@ -408,13 +416,17 @@ pub(crate) async fn download_initdb_tar_zst(
                 .with_context(|| format!("tempfile creation {temp_path}"))
                 .map_err(DownloadError::Other)?;
 
-            let mut download = storage.download(&remote_path).await?;
+            let download = storage.download(&remote_path).await?;
+            let mut download = tokio_util::io::StreamReader::new(download.download_stream);
+            let mut writer = tokio::io::BufWriter::with_capacity(8 * 1024, file);
 
-            tokio::io::copy(&mut download.download_stream, &mut file)
+            tokio::io::copy_buf(&mut download, &mut writer)
                 .await
                 .with_context(|| format!("download initdb.tar.zst at {remote_path:?}"))
                 .map_err(DownloadError::Other)?;
 
+            let mut file = writer.into_inner();
+
             file.seek(std::io::SeekFrom::Start(0))
                 .await
                 .with_context(|| format!("rewinding initdb.tar.zst at: {remote_path:?}"))

pageserver/src/tenant/remote_timeline_client/upload.rs

@@ -41,11 +41,15 @@ pub(super) async fn upload_index_part<'a>(
         .to_s3_bytes()
         .context("serialize index part file into bytes")?;
     let index_part_size = index_part_bytes.len();
-    let index_part_bytes = tokio::io::BufReader::new(std::io::Cursor::new(index_part_bytes));
+    let index_part_bytes = bytes::Bytes::from(index_part_bytes);
 
     let remote_path = remote_index_path(tenant_shard_id, timeline_id, generation);
     storage
-        .upload_storage_object(Box::new(index_part_bytes), index_part_size, &remote_path)
+        .upload_storage_object(
+            futures::stream::once(futures::future::ready(Ok(index_part_bytes))),
+            index_part_size,
+            &remote_path,
+        )
         .await
         .with_context(|| format!("upload index part for '{tenant_shard_id} / {timeline_id}'"))
 }
@@ -101,8 +105,10 @@ pub(super) async fn upload_timeline_layer<'a>(
     let fs_size = usize::try_from(fs_size)
         .with_context(|| format!("convert {source_path:?} size {fs_size} usize"))?;
 
+    let reader = tokio_util::io::ReaderStream::with_capacity(source_file, 8 * 1024);
+
     storage
-        .upload(source_file, fs_size, &storage_path, None)
+        .upload(reader, fs_size, &storage_path, None)
         .await
         .with_context(|| format!("upload layer from local path '{source_path}'"))?;
 
@@ -119,7 +125,8 @@ pub(crate) async fn upload_initdb_dir(
     tracing::trace!("uploading initdb dir");
 
     let size = initdb_dir.len();
-    let bytes = tokio::io::BufReader::new(std::io::Cursor::new(initdb_dir));
+
+    let bytes = futures::stream::once(futures::future::ready(Ok(initdb_dir)));
 
     let remote_path = remote_initdb_archive_path(tenant_id, timeline_id);
     storage