rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-07-08 22:50:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

done

Browse Source
This commit is contained in:
Christian Schwarz
2024-04-23 11:45:27 +00:00
parent 6138941e6c
commit 842474f2d1
6 changed files with 120 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
pageserver/src/tenant/ephemeral_file.rs
View File

@@ -22,7 +22,6 @@ pub struct EphemeralFile {
mod page_caching_sandwich;
mod sandwich;
mod size_tracker_borrowed;
mod zero_padded_buffer;
impl EphemeralFile {

45
pageserver/src/tenant/ephemeral_file/sandwich.rs
View File

@@ -1,11 +1,17 @@
use crate::{
page_cache::PAGE_SZ, tenant::ephemeral_file::zero_padded_buffer, virtual_file::VirtualFile,
page_cache::PAGE_SZ,
tenant::ephemeral_file::zero_padded_buffer,
virtual_file::{
owned_buffers_io::{self, write::Buffer},
VirtualFile,
},
};
use super::size_tracker_borrowed;
pub struct Sandwich {
sandwich: size_tracker_borrowed::Writer<{ Self::TAIL_SZ }>,
buffered_writer: owned_buffers_io::write::BufferedWriter<
zero_padded_buffer::Buf<{ Self::TAIL_SZ }>,
owned_buffers_io::util::size_tracking_writer::Writer<VirtualFile>,
>,
}
pub enum ReadResult<'a> {
@@ -17,28 +23,35 @@ 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,
}
let bytes_flushed_tracker = owned_buffers_io::util::size_tracking_writer::Writer::new(file);
let buffered_writer = owned_buffers_io::write::BufferedWriter::new(
bytes_flushed_tracker,
zero_padded_buffer::Buf::default(),
);
Self { buffered_writer }
}
pub(crate) fn as_inner_virtual_file(&self) -> &VirtualFile {
self.sandwich.as_inner_virtual_file()
self.buffered_writer.as_inner().as_inner()
}
pub async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.sandwich.write_all_borrowed(buf).await
self.buffered_writer.write_buffered_borrowed(buf).await
}
pub fn bytes_written(&self) -> u64 {
self.sandwich.buffered_offset()
let flushed_offset = self.buffered_writer.as_inner().bytes_written();
let buffer: &zero_padded_buffer::Buf<{ Self::TAIL_SZ }> =
self.buffered_writer.inspect_buffer();
let buffered_offset = flushed_offset + u64::try_from(buffer.pending()).unwrap();
buffered_offset
}
pub(crate) async fn read_blk(&self, blknum: u32) -> Result<ReadResult, std::io::Error> {
let flushed_offset = self.sandwich.flushed_offset();
let buffer = self.sandwich.buffered_writer.inspect_buffer();
let buffered_offset = flushed_offset + buffer.pending();
let flushed_offset = self.buffered_writer.as_inner().bytes_written();
let buffer: &zero_padded_buffer::Buf<{ Self::TAIL_SZ }> =
self.buffered_writer.inspect_buffer();
let buffered_offset = flushed_offset + u64::try_from(buffer.pending()).unwrap();
let read_offset = (blknum as u64) * (PAGE_SZ as u64);
assert_eq!(
@@ -75,8 +88,8 @@ impl Sandwich {
.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 zero_padded_buffer: &zero_padded_buffer::Buf<TAIL_SZ> = buffer.raw_iobuf();
let page = zero_padded_buffer[read_offset_in_buffer..(read_offset_in_buffer + PAGE_SZ)];
let zero_padded_slice = buffer.as_zero_padded_slice();
let page = &zero_padded_slice[read_offset_in_buffer..(read_offset_in_buffer + PAGE_SZ)];
Ok(ReadResult::ServedFromZeroPaddedMutableTail {
buffer: page
.try_into()

53
pageserver/src/tenant/ephemeral_file/size_tracker_borrowed.rs
View File

@@ -1,53 +0,0 @@
use crate::virtual_file::{owned_buffers_io, VirtualFile};
use super::zero_padded_buffer;
pub struct Writer<const TAIL_SZ: usize> {
pub buffered_writer: owned_buffers_io::write::BufferedWriter<
zero_padded_buffer::Buf<TAIL_SZ>,
owned_buffers_io::util::size_tracking_writer::Writer<VirtualFile>,
>,
pub bytes_amount: u64,
}
impl<const TAIL_SZ: usize> Writer<TAIL_SZ> {
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 buffered_zero_padded_page_at(
&self,
offset_in_buffer: usize,
) -> &[u8; crate::page_cache::PAGE_SZ] {
let bufio_buffer = self.buffered_writer.inspect_buffer();
let zero_padded_buffer: &zero_padded_buffer::Buf<TAIL_SZ> = bufio_buffer.as_inner();
zero_padded_buffer[offset_in_buffer..(offset_in_buffer + crate::page_cache::PAGE_SZ)]
}
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<usize> {
let nwritten = self.buffered_writer.write_buffered_borrowed(buf).await?;
self.bytes_amount += u64::try_from(nwritten).unwrap();
Ok(nwritten)
}
}

66
pageserver/src/tenant/ephemeral_file/zero_padded_buffer.rs
View File

@@ -29,9 +29,56 @@ impl<const N: usize> Buf<N> {
}
}
impl<const N: usize> crate::virtual_file::owned_buffers_io::write::Buffer for Buf<N> {
type IoBuf = Self;
fn cap(&self) -> usize {
self.allocation.len()
}
fn extend_from_slice(&mut self, other: &[u8]) {
self.invariants();
let remaining = self.cap() - other.len();
if other.len() > remaining {
panic!("calling extend_from_slice() with insufficient remaining capacity");
}
self.allocation[self.written..(self.written + other.len())].copy_from_slice(other);
self.written += other.len();
self.invariants();
}
fn pending(&self) -> usize {
self.written
}
fn flush(self) -> tokio_epoll_uring::Slice<Self> {
self.invariants();
let written = self.written;
tokio_epoll_uring::BoundedBuf::slice(self, 0..written)
}
fn reconstruct_after_flush(this: Self::IoBuf) -> Self {
let Self {
mut allocation,
written,
} = this;
allocation[0..written].fill(0);
let new = Self {
allocation,
written: 0,
};
new.invariants();
new
}
}
/// We have this implementation so that Buf<N> can be used with [`crate::virtual_file::owned_buffers_io::BufferedWriter`].
///
///
/// SAFETY:
///
/// The [`Self::allocation`] is stable becauses boxes are stable.
/// The memory is zero-initialized, so, bytes_init is always N.
///
unsafe impl<const N: usize> tokio_epoll_uring::IoBuf for Buf<N> {
fn stable_ptr(&self) -> *const u8 {
@@ -46,22 +93,3 @@ unsafe impl<const N: usize> tokio_epoll_uring::IoBuf for Buf<N> {
N
}
}
/// SAFETY:
///
/// The [`Self::allocation`] is stable becauses boxes are stable.
///
unsafe impl<const N: usize> tokio_epoll_uring::IoBufMut for Buf<N> {
fn stable_mut_ptr(&mut self) -> *mut u8 {
self.allocation.as_mut_ptr()
}
unsafe fn set_init(&mut self, pos: usize) {
self.invariants();
if pos < self.written {
self.allocation[pos..self.written].fill(0);
}
self.written = pos;
self.invariants();
}
}

2
pageserver/src/tenant/remote_timeline_client/download.rs
View File

@@ -195,7 +195,7 @@ 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::new(
let mut buffered = owned_buffers_io::write::BufferedWriter::<BytesMut, _>::new(
size_tracking,
BytesMut::with_capacity(super::BUFFER_SIZE),
);

97
pageserver/src/virtual_file/owned_buffers_io/write.rs
View File

@@ -1,3 +1,4 @@
use bytes::BytesMut;
use tokio_epoll_uring::{BoundedBuf, IoBuf, Slice};
/// A trait for doing owned-buffer write IO.
@@ -9,61 +10,44 @@ pub trait OwnedAsyncWriter {
) -> std::io::Result<(usize, B::Buf)>;
}
struct Buffer<B> {
pending: usize,
iobuf: B,
pub trait Buffer {
type IoBuf: IoBuf;
fn cap(&self) -> usize;
fn extend_from_slice(&mut self, other: &[u8]);
fn pending(&self) -> usize;
fn flush(self) -> Slice<Self::IoBuf>;
fn reconstruct_after_flush(iobuf: Self::IoBuf) -> Self;
}
impl<B> Buffer<B>
where
B: tokio_epoll_uring::IoBufMut,
{
pub fn new(iobuf: B) -> Self {
Self { pending: 0, iobuf }
impl Buffer for BytesMut {
type IoBuf = BytesMut;
#[inline(always)]
fn cap(&self) -> usize {
self.capacity()
}
fn extend_from_slice(&mut self, other: &[u8]) {
BytesMut::extend_from_slice(self, other)
}
#[inline(always)]
fn invariants(&self) {
debug_assert!(self.pending <= self.iobuf.bytes_total());
debug_assert!(self.pending <= self.iobuf.bytes_init());
fn pending(&self) -> usize {
self.len()
}
#[inline(always)]
pub fn cap(&self) -> usize {
self.iobuf.bytes_total()
}
#[inline(always)]
pub fn pending(&self) -> usize {
self.pending
}
pub fn flush(self) -> Slice<B> {
self.invariants();
self.iobuf.slice(0..self.pending)
}
pub fn extend_from_slice(&mut self, other: &[u8]) {
self.invariants();
let remaining = self.iobuf.bytes_total() - self.pending;
if other.len() > remaining {
panic!("extend_from_slice() would extend beyond buffer capacity; cap={} self.len={} other.len()={}", self.iobuf.bytes_total(), self.pending, other.len());
fn flush(self) -> Slice<BytesMut> {
if self.len() == 0 {
return self.slice_full();
}
// SAFETY: we did bounds-checking above; non-overlapping is guaranteed by Rust borrowing
// (self is borrowed mut, so `other` can't reference self or anything in it).
unsafe {
self.iobuf
.stable_mut_ptr()
.add(self.pending)
.copy_from_nonoverlapping(other.as_ptr(), other.len());
self.iobuf.set_init(self.pending + other.len());
}
self.pending += other.len();
self.invariants();
let len = self.len();
self.slice(0..len)
}
pub fn raw_iobuf(&self) -> &B {
&self.iobuf
fn reconstruct_after_flush(mut iobuf: BytesMut) -> Self {
iobuf.clear();
iobuf
}
}
@@ -89,18 +73,19 @@ pub struct BufferedWriter<B, W> {
// - 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<Buffer<B>>,
buf: Option<B>,
}
impl<B, W> BufferedWriter<B, W>
impl<B, Buf, W> BufferedWriter<B, W>
where
B: tokio_epoll_uring::IoBufMut + Send,
B: Buffer<IoBuf = Buf> + Send,
Buf: IoBuf + Send,
W: OwnedAsyncWriter,
{
pub fn new(writer: W, iobuf: B) -> Self {
pub fn new(writer: W, buf: B) -> Self {
Self {
writer,
buf: Some(Buffer::new(iobuf)),
buf: Some(buf),
}
}
@@ -109,7 +94,7 @@ where
}
/// Panics if used after any of the write paths returned an error
pub fn inspect_buffer(&self) -> &Buffer<B> {
pub fn inspect_buffer(&self) -> &B {
self.buf()
}
@@ -121,7 +106,7 @@ where
}
#[inline(always)]
fn buf(&self) -> &Buffer<B> {
fn buf(&self) -> &B {
self.buf
.as_ref()
.expect("must not use after we returned an error")
@@ -191,9 +176,13 @@ where
async fn flush(&mut self) -> std::io::Result<()> {
let buf = self.buf.take().expect("must not use after an error");
let buf_len = buf.pending();
let (nwritten, mut io_buf) = self.writer.write_all(buf.flush()).await?;
if buf_len == 0 {
self.buf = Some(buf);
return Ok(());
}
let (nwritten, io_buf) = self.writer.write_all(buf.flush()).await?;
assert_eq!(nwritten, buf_len);
self.buf = Some(Buffer::new(io_buf));
self.buf = Some(Buffer::reconstruct_after_flush(io_buf));
Ok(())
}
}