From ea5f307e81101ab5666953d9ddf2f788624b41cc Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Mon, 22 Apr 2024 14:59:04 +0000 Subject: [PATCH] refactor sandwiching code Goals: 1. don't junk up the owned_buffer_io interface with the write_all_borrowed, eliminating one `panic!("should not happen")` 1. make the sandwiching clearer --- pageserver/src/tenant/ephemeral_file.rs | 152 ++--------------- .../ephemeral_file/page_caching_sandwich.rs | 104 ++++++++++++ .../src/tenant/ephemeral_file/sandwich.rs | 87 ++++++++++ .../ephemeral_file/size_tracker_borrowed.rs | 48 ++++++ .../ephemeral_file/zero_padded_buffer.rs} | 54 +++--- .../tenant/remote_timeline_client/download.rs | 8 +- .../tenant/storage_layer/inmemory_layer.rs | 2 +- pageserver/src/virtual_file.rs | 6 - .../util/size_tracking_writer.rs | 7 - .../virtual_file/owned_buffers_io/write.rs | 160 ++++++++++-------- 10 files changed, 379 insertions(+), 249 deletions(-) create mode 100644 pageserver/src/tenant/ephemeral_file/page_caching_sandwich.rs create mode 100644 pageserver/src/tenant/ephemeral_file/sandwich.rs create mode 100644 pageserver/src/tenant/ephemeral_file/size_tracker_borrowed.rs rename pageserver/src/{virtual_file/owned_buffers_io/write/zero_initialized_buffer.rs => tenant/ephemeral_file/zero_padded_buffer.rs} (85%) diff --git a/pageserver/src/tenant/ephemeral_file.rs b/pageserver/src/tenant/ephemeral_file.rs index 5769212150..b558a05fae 100644 --- a/pageserver/src/tenant/ephemeral_file.rs +++ b/pageserver/src/tenant/ephemeral_file.rs @@ -3,42 +3,29 @@ use crate::config::PageServerConf; use crate::context::RequestContext; -use crate::page_cache::{self, PAGE_SZ}; +use crate::page_cache; use crate::tenant::block_io::{BlockCursor, BlockLease, BlockReader}; -use crate::virtual_file::owned_buffers_io::write::OwnedAsyncWriter; -use crate::virtual_file::{self, owned_buffers_io, VirtualFile}; +use crate::virtual_file::{self, VirtualFile}; use camino::Utf8PathBuf; use pageserver_api::shard::TenantShardId; use std::io; use std::sync::atomic::AtomicU64; -use tracing::*; use utils::id::TimelineId; pub struct EphemeralFile { - page_cache_file_id: page_cache::FileId, - _tenant_shard_id: TenantShardId, _timeline_id: TimelineId, - /// We sandwich the buffered writer between two size-tracking writers. - /// This allows us to "elegantly" track in-memory bytes vs flushed bytes, - /// enabling [`Self::read_blk`] to determine whether to read from the - /// buffered writer's buffer, versus going to the VirtualFile. - /// - /// TODO: longer-term, we probably wand to get rid of this in favor - /// of a double-buffering scheme. See this commit's commit message - /// and git history for what we had before this sandwich, it might be useful. - file: owned_buffers_io::util::size_tracking_writer::Writer< - owned_buffers_io::write::BufferedWriter< - { Self::TAIL_SZ }, - owned_buffers_io::util::size_tracking_writer::Writer, - >, - >, + + sandwich: page_caching_sandwich::PageCachingSandwich, } -impl EphemeralFile { - const TAIL_SZ: usize = 64 * 1024; +mod page_caching_sandwich; +mod sandwich; +mod size_tracker_borrowed; +mod zero_padded_buffer; +impl EphemeralFile { pub async fn create( conf: &PageServerConf, tenant_shard_id: TenantShardId, @@ -62,24 +49,20 @@ impl EphemeralFile { .create(true), ) .await?; - let file = owned_buffers_io::util::size_tracking_writer::Writer::new(file); - let file = owned_buffers_io::write::BufferedWriter::new(file); - let file = owned_buffers_io::util::size_tracking_writer::Writer::new(file); Ok(EphemeralFile { - page_cache_file_id: page_cache::next_file_id(), _tenant_shard_id: tenant_shard_id, _timeline_id: timeline_id, - file, + sandwich: page_caching_sandwich::PageCachingSandwich::new(file), }) } pub(crate) fn len(&self) -> u64 { - self.file.bytes_written() + self.sandwich.bytes_written() } - pub(crate) fn id(&self) -> page_cache::FileId { - self.page_cache_file_id + pub(crate) fn page_cache_file_id(&self) -> page_cache::FileId { + self.sandwich.page_cache_file_id() } pub(crate) async fn read_blk( @@ -87,81 +70,7 @@ impl EphemeralFile { blknum: u32, ctx: &RequestContext, ) -> Result { - let buffered_offset = self.file.bytes_written(); - let flushed_offset = self.file.as_inner().as_inner().bytes_written(); - assert!(buffered_offset >= flushed_offset); - let read_offset = (blknum as u64) * (PAGE_SZ as u64); - - assert_eq!( - flushed_offset % (PAGE_SZ as u64), - 0, - "we need this in the logic below, because it assumes the page isn't spread across flushed part and in-memory buffer" - ); - - if read_offset < flushed_offset { - assert!( - read_offset + (PAGE_SZ as u64) <= flushed_offset, - "this impl can't deal with pages spread across flushed & buffered part" - ); - let cache = page_cache::get(); - match cache - .read_immutable_buf(self.page_cache_file_id, blknum, ctx) - .await - .map_err(|e| { - std::io::Error::new( - std::io::ErrorKind::Other, - // order path before error because error is anyhow::Error => might have many contexts - format!( - "ephemeral file: read immutable page #{}: {}: {:#}", - blknum, - self.file.as_inner().as_inner().as_inner().path, - e, - ), - ) - })? { - page_cache::ReadBufResult::Found(guard) => { - return Ok(BlockLease::PageReadGuard(guard)) - } - page_cache::ReadBufResult::NotFound(write_guard) => { - let write_guard = self - .file - .as_inner() - .as_inner() - .as_inner() - .read_exact_at_page(write_guard, blknum as u64 * PAGE_SZ as u64) - .await?; - let read_guard = write_guard.mark_valid(); - return Ok(BlockLease::PageReadGuard(read_guard)); - } - }; - } else { - let read_until_offset = read_offset + (PAGE_SZ as u64); - if !(0..buffered_offset).contains(&read_until_offset) { - // The blob_io code relies on the reader allowing reads past - // the end of what was written, up to end of the current PAGE_SZ chunk. - // This is a relict of the past where we would get a pre-zeroed page from the page cache. - // - // DeltaLayer probably has the same issue, not sure why it needs no special treatment. - let nbytes_past_end = read_until_offset.checked_sub(buffered_offset).unwrap(); - if nbytes_past_end >= (PAGE_SZ as u64) { - // TODO: treat this as error. Pre-existing issue before this patch. - panic!( - "return IO error: read past end of file: read=0x{read_offset:x} buffered=0x{buffered_offset:x} flushed=0x{flushed_offset}" - ) - } - } - let buffer: &[u8; Self::TAIL_SZ] = self.file.as_inner().inspect_buffer(); - let read_offset_in_buffer = read_offset - .checked_sub(flushed_offset) - .expect("would have taken `if` branch instead of this one"); - - let read_offset_in_buffer = usize::try_from(read_offset_in_buffer).unwrap(); - let page = &buffer[read_offset_in_buffer..(read_offset_in_buffer + PAGE_SZ)]; - Ok(BlockLease::EphemeralFileMutableTail( - page.try_into() - .expect("the slice above got it as page-size slice"), - )) - } + self.sandwich.read_blk(blknum, ctx).await } pub(crate) async fn write_blob( @@ -169,24 +78,22 @@ impl EphemeralFile { srcbuf: &[u8], _ctx: &RequestContext, ) -> Result { - let pos = self.file.bytes_written(); + let pos = self.sandwich.bytes_written(); // Write the length field if srcbuf.len() < 0x80 { // short one-byte length header let len_buf = [srcbuf.len() as u8]; - self.file.write_all_borrowed(&len_buf).await?; + self.sandwich.write_all_borrowed(&len_buf).await?; } else { let mut len_buf = u32::to_be_bytes(srcbuf.len() as u32); len_buf[0] |= 0x80; - self.file.write_all_borrowed(&len_buf).await?; + self.sandwich.write_all_borrowed(&len_buf).await?; } // Write the payload - self.file.write_all_borrowed(srcbuf).await?; - - // TODO: bring back pre-warming of page cache, using another sandwich layer + self.sandwich.write_all_borrowed(srcbuf).await?; Ok(pos) } @@ -201,29 +108,6 @@ pub fn is_ephemeral_file(filename: &str) -> bool { } } -impl Drop for EphemeralFile { - fn drop(&mut self) { - // There might still be pages in the [`crate::page_cache`] for this file. - // We leave them there, [`crate::page_cache::PageCache::find_victim`] will evict them when needed. - - // unlink the file - let res = std::fs::remove_file(&self.file.as_inner().as_inner().as_inner().path); - if let Err(e) = res { - if e.kind() != std::io::ErrorKind::NotFound { - // just never log the not found errors, we cannot do anything for them; on detach - // the tenant directory is already gone. - // - // not found files might also be related to https://github.com/neondatabase/neon/issues/2442 - error!( - "could not remove ephemeral file '{}': {}", - self.file.as_inner().as_inner().as_inner().path, - e - ); - } - } - } -} - impl BlockReader for EphemeralFile { fn block_cursor(&self) -> super::block_io::BlockCursor<'_> { BlockCursor::new(super::block_io::BlockReaderRef::EphemeralFile(self)) diff --git a/pageserver/src/tenant/ephemeral_file/page_caching_sandwich.rs b/pageserver/src/tenant/ephemeral_file/page_caching_sandwich.rs new file mode 100644 index 0000000000..0029e58385 --- /dev/null +++ b/pageserver/src/tenant/ephemeral_file/page_caching_sandwich.rs @@ -0,0 +1,104 @@ +//! Implementation of append-only file data structure +//! used to keep in-memory layers spilled on disk. + +use crate::context::RequestContext; +use crate::page_cache::{self, PAGE_SZ}; +use crate::tenant::block_io::BlockLease; +use crate::virtual_file::VirtualFile; + +use std::io; +use tracing::*; + +use super::sandwich; + +pub struct PageCachingSandwich { + page_cache_file_id: page_cache::FileId, + sandwich: super::sandwich::Sandwich, +} + +impl PageCachingSandwich { + pub fn new(file: VirtualFile) -> Self { + Self { + page_cache_file_id: page_cache::next_file_id(), + sandwich: super::sandwich::Sandwich::new(file), + } + } + + pub fn page_cache_file_id(&self) -> page_cache::FileId { + self.page_cache_file_id + } + + pub(crate) async fn write_all_borrowed(&mut self, srcbuf: &[u8]) -> Result { + // It doesn't make sense to proactively fill the page cache on the Pageserver write path + // because Compute is unlikely to access recently written data. + self.sandwich.write_all_borrowed(srcbuf).await + } + + pub(crate) fn bytes_written(&self) -> u64 { + self.sandwich.bytes_written() + } + + pub(crate) async fn read_blk( + &self, + blknum: u32, + ctx: &RequestContext, + ) -> Result { + match self.sandwich.read_blk(blknum).await? { + sandwich::ReadResult::NeedsReadFromVirtualFile { virtual_file } => { + let cache = page_cache::get(); + match cache + .read_immutable_buf(self.page_cache_file_id, blknum, ctx) + .await + .map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::Other, + // order path before error because error is anyhow::Error => might have many contexts + format!( + "ephemeral file: read immutable page #{}: {}: {:#}", + blknum, + self.sandwich.as_inner_virtual_file().path, + e, + ), + ) + })? { + page_cache::ReadBufResult::Found(guard) => { + return Ok(BlockLease::PageReadGuard(guard)) + } + page_cache::ReadBufResult::NotFound(write_guard) => { + let write_guard = virtual_file + .read_exact_at_page(write_guard, blknum as u64 * PAGE_SZ as u64) + .await?; + let read_guard = write_guard.mark_valid(); + return Ok(BlockLease::PageReadGuard(read_guard)); + } + } + } + sandwich::ReadResult::ServedFromZeroPaddedMutableTail { buffer } => { + Ok(BlockLease::EphemeralFileMutableTail(buffer)) + } + } + } +} + +impl Drop for PageCachingSandwich { + fn drop(&mut self) { + // There might still be pages in the [`crate::page_cache`] for this file. + // We leave them there, [`crate::page_cache::PageCache::find_victim`] will evict them when needed. + + // unlink the file + let res = std::fs::remove_file(&self.sandwich.as_inner_virtual_file().path); + if let Err(e) = res { + if e.kind() != std::io::ErrorKind::NotFound { + // just never log the not found errors, we cannot do anything for them; on detach + // the tenant directory is already gone. + // + // not found files might also be related to https://github.com/neondatabase/neon/issues/2442 + error!( + "could not remove ephemeral file '{}': {}", + self.sandwich.as_inner_virtual_file().path, + e + ); + } + } + } +} diff --git a/pageserver/src/tenant/ephemeral_file/sandwich.rs b/pageserver/src/tenant/ephemeral_file/sandwich.rs new file mode 100644 index 0000000000..ba2706bd6a --- /dev/null +++ b/pageserver/src/tenant/ephemeral_file/sandwich.rs @@ -0,0 +1,87 @@ +use crate::{page_cache::PAGE_SZ, virtual_file::VirtualFile}; + +use super::size_tracker_borrowed; + +pub struct Sandwich { + sandwich: size_tracker_borrowed::Writer<{ Self::TAIL_SZ }>, +} + +pub enum ReadResult<'a> { + NeedsReadFromVirtualFile { virtual_file: &'a VirtualFile }, + ServedFromZeroPaddedMutableTail { buffer: &'a [u8; PAGE_SZ] }, +} + +impl Sandwich { + const TAIL_SZ: usize = 64 * 1024; + + pub fn new(file: VirtualFile) -> Self { + let size_borrowed_tracker = size_tracker_borrowed::Writer::new(file); + Self { + sandwich: size_borrowed_tracker, + } + } + + pub(crate) fn as_inner_virtual_file(&self) -> &VirtualFile { + self.sandwich.as_inner_virtual_file() + } + + pub async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { + self.sandwich.write_all_borrowed(buf).await + } + + pub fn bytes_written(&self) -> u64 { + self.sandwich.buffered_offset() + } + + pub(crate) async fn read_blk(&self, blknum: u32) -> Result { + let buffered_offset = self.sandwich.buffered_offset(); + let flushed_offset = self.sandwich.flushed_offset(); + assert!(buffered_offset >= flushed_offset); + let read_offset = (blknum as u64) * (PAGE_SZ as u64); + + assert_eq!( + flushed_offset % (PAGE_SZ as u64), + 0, + "we need this in the logic below, because it assumes the page isn't spread across flushed part and in-memory buffer" + ); + + if read_offset < flushed_offset { + assert!( + read_offset + (PAGE_SZ as u64) <= flushed_offset, + "this impl can't deal with pages spread across flushed & buffered part" + ); + Ok(ReadResult::NeedsReadFromVirtualFile { + virtual_file: self.as_inner_virtual_file(), + }) + } else { + let read_until_offset = read_offset + (PAGE_SZ as u64); + if !(0..buffered_offset).contains(&read_until_offset) { + // The blob_io code relies on the reader allowing reads past + // the end of what was written, up to end of the current PAGE_SZ chunk. + // This is a relict of the past where we would get a pre-zeroed page from the page cache. + // + // DeltaLayer probably has the same issue, not sure why it needs no special treatment. + let nbytes_past_end = read_until_offset.checked_sub(buffered_offset).unwrap(); + if nbytes_past_end >= (PAGE_SZ as u64) { + // TODO: treat this as error. Pre-existing issue before this patch. + panic!( + "return IO error: read past end of file: read=0x{read_offset:x} buffered=0x{buffered_offset:x} flushed=0x{flushed_offset}" + ) + } + } + let buffer: &[u8; Self::TAIL_SZ] = + self.sandwich.inspect_buffer().as_zero_padded_slice(); + let read_offset_in_buffer = read_offset + .checked_sub(flushed_offset) + .expect("would have taken `if` branch instead of this one"); + + let read_offset_in_buffer = usize::try_from(read_offset_in_buffer).unwrap(); + let page = &buffer[read_offset_in_buffer..(read_offset_in_buffer + PAGE_SZ)]; + Ok(ReadResult::ServedFromZeroPaddedMutableTail { + buffer: page + .try_into() + .expect("the slice above got it as page-size slice"), + }) + } + } +} diff --git a/pageserver/src/tenant/ephemeral_file/size_tracker_borrowed.rs b/pageserver/src/tenant/ephemeral_file/size_tracker_borrowed.rs new file mode 100644 index 0000000000..ba7527d253 --- /dev/null +++ b/pageserver/src/tenant/ephemeral_file/size_tracker_borrowed.rs @@ -0,0 +1,48 @@ +use crate::virtual_file::{owned_buffers_io, VirtualFile}; + +use super::zero_padded_buffer; + +pub struct Writer { + buffered_writer: owned_buffers_io::write::BufferedWriter< + zero_padded_buffer::Buf, + owned_buffers_io::util::size_tracking_writer::Writer, + >, + bytes_amount: u64, +} + +impl Writer { + pub fn new(dst: VirtualFile) -> Self { + let bytes_flushed_tracker = owned_buffers_io::util::size_tracking_writer::Writer::new(dst); + let buffered = owned_buffers_io::write::BufferedWriter::new( + bytes_flushed_tracker, + zero_padded_buffer::Buf::default(), + ); + Self { + buffered_writer: { buffered }, + bytes_amount: 0, + } + } + + pub fn buffered_offset(&self) -> u64 { + self.bytes_amount + } + + pub fn flushed_offset(&self) -> u64 { + self.buffered_writer.as_inner().bytes_written() + } + + pub fn inspect_buffer(&self) -> &zero_padded_buffer::Buf { + self.buffered_writer.inspect_buffer() + } + + pub(crate) fn as_inner_virtual_file(&self) -> &VirtualFile { + self.buffered_writer.as_inner().as_inner() + } + + #[inline(always)] + pub async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { + let nwritten = self.buffered_writer.write_buffered_borrowed(buf).await?; + self.bytes_amount += u64::try_from(nwritten).unwrap(); + Ok(nwritten) + } +} diff --git a/pageserver/src/virtual_file/owned_buffers_io/write/zero_initialized_buffer.rs b/pageserver/src/tenant/ephemeral_file/zero_padded_buffer.rs similarity index 85% rename from pageserver/src/virtual_file/owned_buffers_io/write/zero_initialized_buffer.rs rename to pageserver/src/tenant/ephemeral_file/zero_padded_buffer.rs index a48bb9767a..fc83377e92 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/write/zero_initialized_buffer.rs +++ b/pageserver/src/tenant/ephemeral_file/zero_padded_buffer.rs @@ -1,5 +1,7 @@ use std::mem::MaybeUninit; +use crate::virtual_file::owned_buffers_io; + pub struct Buf { allocation: Box<[u8; N]>, written: usize, @@ -26,32 +28,6 @@ impl Buf { pub fn as_zero_padded_slice(&self) -> &[u8; N] { &self.allocation } - - /// panics if there's not enough capacity left - pub fn extend_from_slice(&mut self, buf: &[u8]) { - self.invariants(); - let can = N - self.written; - let want = buf.len(); - assert!(want <= can, "{:x} {:x}", want, can); - self.allocation[self.written..(self.written + want)].copy_from_slice(buf); - self.written += want; - self.invariants(); - } - - pub fn len(&self) -> usize { - self.written - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - pub fn clear(&mut self) { - self.invariants(); - self.written = 0; - self.allocation[..].fill(0); - self.invariants(); - } } /// SAFETY: @@ -71,3 +47,29 @@ unsafe impl tokio_epoll_uring::IoBuf for Buf { self.written // ? } } + +impl owned_buffers_io::write::Buffer for Buf { + const BUFFER_SIZE: usize = N; + + /// panics if there's not enough capacity left + fn extend_from_slice(&mut self, buf: &[u8]) { + self.invariants(); + let can = N - self.written; + let want = buf.len(); + assert!(want <= can, "{:x} {:x}", want, can); + self.allocation[self.written..(self.written + want)].copy_from_slice(buf); + self.written += want; + self.invariants(); + } + + fn len(&self) -> usize { + self.written + } + + fn clear(&mut self) { + self.invariants(); + self.written = 0; + self.allocation[..].fill(0); + self.invariants(); + } +} diff --git a/pageserver/src/tenant/remote_timeline_client/download.rs b/pageserver/src/tenant/remote_timeline_client/download.rs index 6ee8ad7155..56ddb9aa9a 100644 --- a/pageserver/src/tenant/remote_timeline_client/download.rs +++ b/pageserver/src/tenant/remote_timeline_client/download.rs @@ -194,10 +194,10 @@ async fn download_object<'a>( // There's chunks_vectored() on the stream. let (bytes_amount, destination_file) = async { let size_tracking = size_tracking_writer::Writer::new(destination_file); - let mut buffered = owned_buffers_io::write::BufferedWriter::< - { super::BUFFER_SIZE }, - _, - >::new(size_tracking); + let mut buffered = owned_buffers_io::write::BufferedWriter::new( + size_tracking, + owned_buffers_io::write::BytesMutBuffer::<{ super::BUFFER_SIZE }>::new(), + ); while let Some(res) = futures::StreamExt::next(&mut download.download_stream).await { diff --git a/pageserver/src/tenant/storage_layer/inmemory_layer.rs b/pageserver/src/tenant/storage_layer/inmemory_layer.rs index 29751641b4..3a77f972ee 100644 --- a/pageserver/src/tenant/storage_layer/inmemory_layer.rs +++ b/pageserver/src/tenant/storage_layer/inmemory_layer.rs @@ -454,7 +454,7 @@ impl InMemoryLayer { trace!("initializing new empty InMemoryLayer for writing on timeline {timeline_id} at {start_lsn}"); let file = EphemeralFile::create(conf, tenant_shard_id, timeline_id).await?; - let key = InMemoryLayerFileId(file.id()); + let key = InMemoryLayerFileId(file.page_cache_file_id()); Ok(InMemoryLayer { file_id: key, diff --git a/pageserver/src/virtual_file.rs b/pageserver/src/virtual_file.rs index 779a2ecdeb..911227c9c6 100644 --- a/pageserver/src/virtual_file.rs +++ b/pageserver/src/virtual_file.rs @@ -1093,12 +1093,6 @@ impl OwnedAsyncWriter for VirtualFile { let (buf, res) = VirtualFile::write_all(self, buf).await; res.map(move |v| (v, buf)) } - - #[inline(always)] - async fn write_all_borrowed(&mut self, _buf: &[u8]) -> std::io::Result { - // TODO: ensure this through the type system - panic!("this should not happen"); - } } impl OpenFiles { diff --git a/pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs b/pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs index 6f123bc195..8fe6400885 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs +++ b/pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs @@ -39,11 +39,4 @@ where self.bytes_amount += u64::try_from(nwritten).unwrap(); Ok((nwritten, buf)) } - - #[inline(always)] - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { - let nwritten = self.dst.write_all_borrowed(buf).await?; - self.bytes_amount += u64::try_from(nwritten).unwrap(); - Ok(nwritten) - } } diff --git a/pageserver/src/virtual_file/owned_buffers_io/write.rs b/pageserver/src/virtual_file/owned_buffers_io/write.rs index 6852fd55c1..1a711aeb22 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/write.rs +++ b/pageserver/src/virtual_file/owned_buffers_io/write.rs @@ -1,3 +1,4 @@ +use bytes::BytesMut; use tokio_epoll_uring::{BoundedBuf, IoBuf, Slice}; /// A trait for doing owned-buffer write IO. @@ -7,7 +8,57 @@ pub trait OwnedAsyncWriter { &mut self, buf: B, ) -> std::io::Result<(usize, B::Buf)>; - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result; +} + +pub trait Buffer { + const BUFFER_SIZE: usize; + + fn len(&self) -> usize; + fn clear(&mut self); + fn extend_from_slice(&mut self, other: &[u8]); +} + +pub struct BytesMutBuffer { + buf: BytesMut, +} + +/// SAFETY: just forwards to the pre-existing impl for BytesMut +unsafe impl IoBuf for BytesMutBuffer { + fn stable_ptr(&self) -> *const u8 { + IoBuf::stable_ptr(&self.buf) + } + + fn bytes_init(&self) -> usize { + IoBuf::bytes_init(&self.buf) + } + + fn bytes_total(&self) -> usize { + IoBuf::bytes_total(&self.buf) + } +} + +impl Buffer for BytesMutBuffer { + const BUFFER_SIZE: usize = BUFFER_SIZE; + + fn len(&self) -> usize { + self.buf.len() + } + + fn clear(&mut self) { + self.buf.clear() + } + + fn extend_from_slice(&mut self, other: &[u8]) { + self.buf.extend_from_slice(other) + } +} + +impl BytesMutBuffer { + pub fn new() -> Self { + BytesMutBuffer { + buf: BytesMut::with_capacity(BUFFER_SIZE), + } + } } /// A wrapper aorund an [`OwnedAsyncWriter`] that batches smaller writers @@ -25,26 +76,25 @@ pub trait OwnedAsyncWriter { /// /// In such cases, a different implementation that always buffers in memory /// may be preferable. -pub struct BufferedWriter { +pub struct BufferedWriter { writer: W, // invariant: always remains Some(buf) // with buf.capacity() == BUFFER_SIZE except // - while IO is ongoing => goes back to Some() once the IO completed successfully // - after an IO error => stays `None` forever // In these exceptional cases, it's `None`. - buf: Option>, + buf: Option, } -mod zero_initialized_buffer; - -impl BufferedWriter +impl BufferedWriter where + B: Buffer + IoBuf + Send, W: OwnedAsyncWriter, { - pub fn new(writer: W) -> Self { + pub fn new(writer: W, buffer: B) -> Self { Self { writer, - buf: Some(zero_initialized_buffer::Buf::default()), + buf: Some(buffer), } } @@ -52,13 +102,9 @@ where &self.writer } - /// panics if used after an error - pub fn inspect_buffer(&self) -> &[u8; BUFFER_SIZE] { - self.buf - .as_ref() - // TODO: can this happen on the EphemeralFile read path? - .expect("must not use after an error") - .as_zero_padded_slice() + /// Panics if used after any of the write paths returned an error + pub fn inspect_buffer(&self) -> &B { + self.buf.as_ref().expect("must not use after an error") } pub async fn flush_and_into_inner(mut self) -> std::io::Result { @@ -68,13 +114,13 @@ where Ok(writer) } - pub async fn write_buffered(&mut self, chunk: Slice) -> std::io::Result<(usize, B)> + pub async fn write_buffered(&mut self, chunk: Slice) -> std::io::Result<(usize, S)> where - B: IoBuf + Send, + S: IoBuf + Send, { let chunk_len = chunk.len(); // avoid memcpy for the middle of the chunk - if chunk.len() >= BUFFER_SIZE { + if chunk.len() >= B::BUFFER_SIZE { self.flush().await?; // do a big write, bypassing `buf` assert_eq!( @@ -89,17 +135,17 @@ where return Ok((nwritten, chunk)); } // in-memory copy the < BUFFER_SIZED tail of the chunk - assert!(chunk.len() < BUFFER_SIZE); + assert!(chunk.len() < B::BUFFER_SIZE); let mut slice = &chunk[..]; while !slice.is_empty() { let buf = self.buf.as_mut().expect("must not use after an error"); - let need = BUFFER_SIZE - buf.len(); + let need = B::BUFFER_SIZE - buf.len(); let have = slice.len(); let n = std::cmp::min(need, have); buf.extend_from_slice(&slice[..n]); slice = &slice[n..]; - if buf.len() >= BUFFER_SIZE { - assert_eq!(buf.len(), BUFFER_SIZE); + if buf.len() >= B::BUFFER_SIZE { + assert_eq!(buf.len(), B::BUFFER_SIZE); self.flush().await?; } } @@ -107,19 +153,22 @@ where Ok((chunk_len, chunk.into_inner())) } - /// Always goes through the internal buffer. - /// Guaranteed to never invoke [`OwnedAsyncWriter::write_all_borrowed`] on the underlying. - pub async fn write_all_borrowed(&mut self, mut chunk: &[u8]) -> std::io::Result { + /// Strictly less performant variant of [`Self::write_buffered`] that allows writing borrowed data. + /// + /// It is less performant because we always have to copy the borrowed data into the internal buffer + /// before we can do the IO. The [`Self::write_buffered`] can avoid this, which is more performant + /// for large writes. + pub async fn write_buffered_borrowed(&mut self, mut chunk: &[u8]) -> std::io::Result { let chunk_len = chunk.len(); while !chunk.is_empty() { let buf = self.buf.as_mut().expect("must not use after an error"); - let need = BUFFER_SIZE - buf.len(); + let need = B::BUFFER_SIZE - buf.len(); let have = chunk.len(); let n = std::cmp::min(need, have); buf.extend_from_slice(&chunk[..n]); chunk = &chunk[n..]; - if buf.len() >= BUFFER_SIZE { - assert_eq!(buf.len(), BUFFER_SIZE); + if buf.len() >= B::BUFFER_SIZE { + assert_eq!(buf.len(), B::BUFFER_SIZE); self.flush().await?; } } @@ -128,7 +177,7 @@ where async fn flush(&mut self) -> std::io::Result<()> { let buf = self.buf.take().expect("must not use after an error"); - if buf.is_empty() { + if buf.bytes_init() == 0 { self.buf = Some(buf); return std::io::Result::Ok(()); } @@ -141,27 +190,6 @@ where } } -impl OwnedAsyncWriter - for BufferedWriter -{ - #[inline(always)] - async fn write_all, Buf: IoBuf + Send>( - &mut self, - buf: B, - ) -> std::io::Result<(usize, B::Buf)> { - let nbytes = buf.bytes_init(); - if nbytes == 0 { - return Ok((0, Slice::into_inner(buf.slice_full()))); - } - let slice = buf.slice(0..nbytes); - BufferedWriter::write_buffered(self, slice).await - } - #[inline(always)] - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { - BufferedWriter::write_all_borrowed(self, buf).await - } -} - impl OwnedAsyncWriter for Vec { async fn write_all, Buf: IoBuf + Send>( &mut self, @@ -175,11 +203,6 @@ impl OwnedAsyncWriter for Vec { self.extend_from_slice(&buf[..]); Ok((buf.len(), Slice::into_inner(buf))) } - - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { - self.extend_from_slice(buf); - Ok(buf.len()) - } } #[cfg(test)] @@ -204,11 +227,6 @@ mod tests { self.writes.push(Vec::from(&buf[..])); Ok((buf.len(), Slice::into_inner(buf))) } - - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { - self.writes.push(Vec::from(buf)); - Ok(buf.len()) - } } macro_rules! write { @@ -222,7 +240,7 @@ mod tests { #[tokio::test] async fn test_buffered_writes_only() -> std::io::Result<()> { let recorder = RecorderWriter::default(); - let mut writer = BufferedWriter::<2, _>::new(recorder); + let mut writer = BufferedWriter::new(recorder, BytesMutBuffer::<2>::new()); write!(writer, b"a"); write!(writer, b"b"); write!(writer, b"c"); @@ -239,7 +257,7 @@ mod tests { #[tokio::test] async fn test_passthrough_writes_only() -> std::io::Result<()> { let recorder = RecorderWriter::default(); - let mut writer = BufferedWriter::<2, _>::new(recorder); + let mut writer = BufferedWriter::new(recorder, BytesMutBuffer::<2>::new()); write!(writer, b"abc"); write!(writer, b"de"); write!(writer, b""); @@ -255,7 +273,7 @@ mod tests { #[tokio::test] async fn test_passthrough_write_with_nonempty_buffer() -> std::io::Result<()> { let recorder = RecorderWriter::default(); - let mut writer = BufferedWriter::<2, _>::new(recorder); + let mut writer = BufferedWriter::new(recorder, BytesMutBuffer::<2>::new()); write!(writer, b"a"); write!(writer, b"bc"); write!(writer, b"d"); @@ -271,15 +289,15 @@ mod tests { #[tokio::test] async fn test_write_all_borrowed_always_goes_through_buffer() -> std::io::Result<()> { let recorder = RecorderWriter::default(); - let mut writer = BufferedWriter::<2, _>::new(recorder); + let mut writer = BufferedWriter::new(recorder, BytesMutBuffer::<2>::new()); - writer.write_all_borrowed(b"abc").await?; - writer.write_all_borrowed(b"d").await?; - writer.write_all_borrowed(b"e").await?; - writer.write_all_borrowed(b"fg").await?; - writer.write_all_borrowed(b"hi").await?; - writer.write_all_borrowed(b"j").await?; - writer.write_all_borrowed(b"klmno").await?; + writer.write_buffered_borrowed(b"abc").await?; + writer.write_buffered_borrowed(b"d").await?; + writer.write_buffered_borrowed(b"e").await?; + writer.write_buffered_borrowed(b"fg").await?; + writer.write_buffered_borrowed(b"hi").await?; + writer.write_buffered_borrowed(b"j").await?; + writer.write_buffered_borrowed(b"klmno").await?; let recorder = writer.flush_and_into_inner().await?; assert_eq!(