LCOV - code coverage report
Current view: top level - pageserver/src - virtual_file.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 90.5 % 1147 1038
Test Date: 2025-02-20 13:11:02 Functions: 86.8 % 266 231

            Line data    Source code
       1              : //! VirtualFile is like a normal File, but it's not bound directly to
       2              : //! a file descriptor.
       3              : //!
       4              : //! Instead, the file is opened when it's read from,
       5              : //! and if too many files are open globally in the system, least-recently
       6              : //! used ones are closed.
       7              : //!
       8              : //! To track which files have been recently used, we use the clock algorithm
       9              : //! with a 'recently_used' flag on each slot.
      10              : //!
      11              : //! This is similar to PostgreSQL's virtual file descriptor facility in
      12              : //! src/backend/storage/file/fd.c
      13              : //!
      14              : use crate::context::RequestContext;
      15              : use crate::metrics::{StorageIoOperation, STORAGE_IO_SIZE, STORAGE_IO_TIME_METRIC};
      16              : 
      17              : use crate::page_cache::{PageWriteGuard, PAGE_SZ};
      18              : use crate::tenant::TENANTS_SEGMENT_NAME;
      19              : use camino::{Utf8Path, Utf8PathBuf};
      20              : use once_cell::sync::OnceCell;
      21              : use owned_buffers_io::aligned_buffer::buffer::AlignedBuffer;
      22              : use owned_buffers_io::aligned_buffer::{AlignedBufferMut, AlignedSlice, ConstAlign};
      23              : use owned_buffers_io::io_buf_aligned::{IoBufAligned, IoBufAlignedMut};
      24              : use owned_buffers_io::io_buf_ext::FullSlice;
      25              : use pageserver_api::config::defaults::DEFAULT_IO_BUFFER_ALIGNMENT;
      26              : use pageserver_api::shard::TenantShardId;
      27              : use std::fs::File;
      28              : use std::io::{Error, ErrorKind, Seek, SeekFrom};
      29              : #[cfg(target_os = "linux")]
      30              : use std::os::unix::fs::OpenOptionsExt;
      31              : use tokio_epoll_uring::{BoundedBuf, IoBuf, IoBufMut, Slice};
      32              : 
      33              : use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
      34              : use std::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering};
      35              : use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
      36              : use tokio::time::Instant;
      37              : 
      38              : pub use pageserver_api::models::virtual_file as api;
      39              : pub(crate) mod io_engine;
      40              : pub use io_engine::feature_test as io_engine_feature_test;
      41              : pub use io_engine::io_engine_for_bench;
      42              : pub use io_engine::FeatureTestResult as IoEngineFeatureTestResult;
      43              : mod metadata;
      44              : mod open_options;
      45              : use self::owned_buffers_io::write::OwnedAsyncWriter;
      46              : pub(crate) use api::IoMode;
      47              : pub(crate) use io_engine::IoEngineKind;
      48              : pub(crate) use metadata::Metadata;
      49              : pub(crate) use open_options::*;
      50              : 
      51              : pub(crate) mod owned_buffers_io {
      52              :     //! Abstractions for IO with owned buffers.
      53              :     //!
      54              :     //! Not actually tied to [`crate::virtual_file`] specifically, but, it's the primary
      55              :     //! reason we need this abstraction.
      56              :     //!
      57              :     //! Over time, this could move into the `tokio-epoll-uring` crate, maybe `uring-common`,
      58              :     //! but for the time being we're proving out the primitives in the neon.git repo
      59              :     //! for faster iteration.
      60              : 
      61              :     pub(crate) mod aligned_buffer;
      62              :     pub(crate) mod io_buf_aligned;
      63              :     pub(crate) mod io_buf_ext;
      64              :     pub(crate) mod slice;
      65              :     pub(crate) mod write;
      66              : }
      67              : 
      68              : #[derive(Debug)]
      69              : pub struct VirtualFile {
      70              :     inner: VirtualFileInner,
      71              :     _mode: IoMode,
      72              : }
      73              : 
      74              : impl VirtualFile {
      75              :     /// Open a file in read-only mode. Like File::open.
      76         2112 :     pub async fn open<P: AsRef<Utf8Path>>(
      77         2112 :         path: P,
      78         2112 :         ctx: &RequestContext,
      79         2112 :     ) -> Result<Self, std::io::Error> {
      80         2112 :         let inner = VirtualFileInner::open(path, ctx).await?;
      81         2112 :         Ok(VirtualFile {
      82         2112 :             inner,
      83         2112 :             _mode: IoMode::Buffered,
      84         2112 :         })
      85         2112 :     }
      86              : 
      87              :     /// Open a file in read-only mode. Like File::open.
      88              :     ///
      89              :     /// `O_DIRECT` will be enabled base on `virtual_file_io_mode`.
      90         2456 :     pub async fn open_v2<P: AsRef<Utf8Path>>(
      91         2456 :         path: P,
      92         2456 :         ctx: &RequestContext,
      93         2456 :     ) -> Result<Self, std::io::Error> {
      94         2456 :         Self::open_with_options_v2(path.as_ref(), OpenOptions::new().read(true), ctx).await
      95         2456 :     }
      96              : 
      97         3030 :     pub async fn create<P: AsRef<Utf8Path>>(
      98         3030 :         path: P,
      99         3030 :         ctx: &RequestContext,
     100         3030 :     ) -> Result<Self, std::io::Error> {
     101         3030 :         let inner = VirtualFileInner::create(path, ctx).await?;
     102         3030 :         Ok(VirtualFile {
     103         3030 :             inner,
     104         3030 :             _mode: IoMode::Buffered,
     105         3030 :         })
     106         3030 :     }
     107              : 
     108            0 :     pub async fn create_v2<P: AsRef<Utf8Path>>(
     109            0 :         path: P,
     110            0 :         ctx: &RequestContext,
     111            0 :     ) -> Result<Self, std::io::Error> {
     112            0 :         VirtualFile::open_with_options_v2(
     113            0 :             path.as_ref(),
     114            0 :             OpenOptions::new().write(true).create(true).truncate(true),
     115            0 :             ctx,
     116            0 :         )
     117            0 :         .await
     118            0 :     }
     119              : 
     120         1592 :     pub async fn open_with_options<P: AsRef<Utf8Path>>(
     121         1592 :         path: P,
     122         1592 :         open_options: &OpenOptions,
     123         1592 :         ctx: &RequestContext, /* TODO: carry a pointer to the metrics in the RequestContext instead of the parsing https://github.com/neondatabase/neon/issues/6107 */
     124         1592 :     ) -> Result<Self, std::io::Error> {
     125         1592 :         let inner = VirtualFileInner::open_with_options(path, open_options, ctx).await?;
     126         1592 :         Ok(VirtualFile {
     127         1592 :             inner,
     128         1592 :             _mode: IoMode::Buffered,
     129         1592 :         })
     130         1592 :     }
     131              : 
     132         5068 :     pub async fn open_with_options_v2<P: AsRef<Utf8Path>>(
     133         5068 :         path: P,
     134         5068 :         open_options: &OpenOptions,
     135         5068 :         ctx: &RequestContext, /* TODO: carry a pointer to the metrics in the RequestContext instead of the parsing https://github.com/neondatabase/neon/issues/6107 */
     136         5068 :     ) -> Result<Self, std::io::Error> {
     137         5068 :         let file = match get_io_mode() {
     138              :             IoMode::Buffered => {
     139         5068 :                 let inner = VirtualFileInner::open_with_options(path, open_options, ctx).await?;
     140         5068 :                 VirtualFile {
     141         5068 :                     inner,
     142         5068 :                     _mode: IoMode::Buffered,
     143         5068 :                 }
     144              :             }
     145              :             #[cfg(target_os = "linux")]
     146              :             IoMode::Direct => {
     147            0 :                 let inner = VirtualFileInner::open_with_options(
     148            0 :                     path,
     149            0 :                     open_options.clone().custom_flags(nix::libc::O_DIRECT),
     150            0 :                     ctx,
     151            0 :                 )
     152            0 :                 .await?;
     153            0 :                 VirtualFile {
     154            0 :                     inner,
     155            0 :                     _mode: IoMode::Direct,
     156            0 :                 }
     157              :             }
     158              :         };
     159         5068 :         Ok(file)
     160         5068 :     }
     161              : 
     162         2376 :     pub fn path(&self) -> &Utf8Path {
     163         2376 :         self.inner.path.as_path()
     164         2376 :     }
     165              : 
     166           44 :     pub async fn crashsafe_overwrite<B: BoundedBuf<Buf = Buf> + Send, Buf: IoBuf + Send>(
     167           44 :         final_path: Utf8PathBuf,
     168           44 :         tmp_path: Utf8PathBuf,
     169           44 :         content: B,
     170           44 :     ) -> std::io::Result<()> {
     171           44 :         VirtualFileInner::crashsafe_overwrite(final_path, tmp_path, content).await
     172           44 :     }
     173              : 
     174         5634 :     pub async fn sync_all(&self) -> Result<(), Error> {
     175         5634 :         if SYNC_MODE.load(std::sync::atomic::Ordering::Relaxed) == SyncMode::UnsafeNoSync as u8 {
     176            0 :             return Ok(());
     177         5634 :         }
     178         5634 :         self.inner.sync_all().await
     179         5634 :     }
     180              : 
     181            0 :     pub async fn sync_data(&self) -> Result<(), Error> {
     182            0 :         if SYNC_MODE.load(std::sync::atomic::Ordering::Relaxed) == SyncMode::UnsafeNoSync as u8 {
     183            0 :             return Ok(());
     184            0 :         }
     185            0 :         self.inner.sync_data().await
     186            0 :     }
     187              : 
     188         3604 :     pub async fn metadata(&self) -> Result<Metadata, Error> {
     189         3604 :         self.inner.metadata().await
     190         3604 :     }
     191              : 
     192          516 :     pub fn remove(self) {
     193          516 :         self.inner.remove();
     194          516 :     }
     195              : 
     196        11352 :     pub async fn seek(&mut self, pos: SeekFrom) -> Result<u64, Error> {
     197        11352 :         self.inner.seek(pos).await
     198        11352 :     }
     199              : 
     200       461970 :     pub async fn read_exact_at<Buf>(
     201       461970 :         &self,
     202       461970 :         slice: Slice<Buf>,
     203       461970 :         offset: u64,
     204       461970 :         ctx: &RequestContext,
     205       461970 :     ) -> Result<Slice<Buf>, Error>
     206       461970 :     where
     207       461970 :         Buf: IoBufAlignedMut + Send,
     208       461970 :     {
     209       461970 :         self.inner.read_exact_at(slice, offset, ctx).await
     210       461970 :     }
     211              : 
     212        63357 :     pub async fn read_exact_at_page(
     213        63357 :         &self,
     214        63357 :         page: PageWriteGuard<'static>,
     215        63357 :         offset: u64,
     216        63357 :         ctx: &RequestContext,
     217        63357 :     ) -> Result<PageWriteGuard<'static>, Error> {
     218        63357 :         self.inner.read_exact_at_page(page, offset, ctx).await
     219        63357 :     }
     220              : 
     221        13222 :     pub async fn write_all_at<Buf: IoBufAligned + Send>(
     222        13222 :         &self,
     223        13222 :         buf: FullSlice<Buf>,
     224        13222 :         offset: u64,
     225        13222 :         ctx: &RequestContext,
     226        13222 :     ) -> (FullSlice<Buf>, Result<(), Error>) {
     227        13222 :         self.inner.write_all_at(buf, offset, ctx).await
     228        13222 :     }
     229              : 
     230      2260790 :     pub async fn write_all<Buf: IoBuf + Send>(
     231      2260790 :         &mut self,
     232      2260790 :         buf: FullSlice<Buf>,
     233      2260790 :         ctx: &RequestContext,
     234      2260790 :     ) -> (FullSlice<Buf>, Result<usize, Error>) {
     235      2260790 :         self.inner.write_all(buf, ctx).await
     236      2260790 :     }
     237              : 
     238          448 :     async fn read_to_end(&mut self, buf: &mut Vec<u8>, ctx: &RequestContext) -> Result<(), Error> {
     239          448 :         self.inner.read_to_end(buf, ctx).await
     240          448 :     }
     241              : 
     242            0 :     pub(crate) async fn read_to_string(
     243            0 :         &mut self,
     244            0 :         ctx: &RequestContext,
     245            0 :     ) -> Result<String, anyhow::Error> {
     246            0 :         let mut buf = Vec::new();
     247            0 :         self.read_to_end(&mut buf, ctx).await?;
     248            0 :         Ok(String::from_utf8(buf)?)
     249            0 :     }
     250              : }
     251              : 
     252              : /// Indicates whether to enable fsync, fdatasync, or O_SYNC/O_DSYNC when writing
     253              : /// files. Switching this off is unsafe and only used for testing on machines
     254              : /// with slow drives.
     255              : #[repr(u8)]
     256              : pub enum SyncMode {
     257              :     Sync,
     258              :     UnsafeNoSync,
     259              : }
     260              : 
     261              : impl TryFrom<u8> for SyncMode {
     262              :     type Error = u8;
     263              : 
     264            0 :     fn try_from(value: u8) -> Result<Self, Self::Error> {
     265            0 :         Ok(match value {
     266            0 :             v if v == (SyncMode::Sync as u8) => SyncMode::Sync,
     267            0 :             v if v == (SyncMode::UnsafeNoSync as u8) => SyncMode::UnsafeNoSync,
     268            0 :             x => return Err(x),
     269              :         })
     270            0 :     }
     271              : }
     272              : 
     273              : ///
     274              : /// A virtual file descriptor. You can use this just like std::fs::File, but internally
     275              : /// the underlying file is closed if the system is low on file descriptors,
     276              : /// and re-opened when it's accessed again.
     277              : ///
     278              : /// Like with std::fs::File, multiple threads can read/write the file concurrently,
     279              : /// holding just a shared reference the same VirtualFile, using the read_at() / write_at()
     280              : /// functions from the FileExt trait. But the functions from the Read/Write/Seek traits
     281              : /// require a mutable reference, because they modify the "current position".
     282              : ///
     283              : /// Each VirtualFile has a physical file descriptor in the global OPEN_FILES array, at the
     284              : /// slot that 'handle points to, if the underlying file is currently open. If it's not
     285              : /// currently open, the 'handle' can still point to the slot where it was last kept. The
     286              : /// 'tag' field is used to detect whether the handle still is valid or not.
     287              : ///
     288              : #[derive(Debug)]
     289              : pub struct VirtualFileInner {
     290              :     /// Lazy handle to the global file descriptor cache. The slot that this points to
     291              :     /// might contain our File, or it may be empty, or it may contain a File that
     292              :     /// belongs to a different VirtualFile.
     293              :     handle: RwLock<SlotHandle>,
     294              : 
     295              :     /// Current file position
     296              :     pos: u64,
     297              : 
     298              :     /// File path and options to use to open it.
     299              :     ///
     300              :     /// Note: this only contains the options needed to re-open it. For example,
     301              :     /// if a new file is created, we only pass the create flag when it's initially
     302              :     /// opened, in the VirtualFile::create() function, and strip the flag before
     303              :     /// storing it here.
     304              :     pub path: Utf8PathBuf,
     305              :     open_options: OpenOptions,
     306              : 
     307              :     // These are strings becase we only use them for metrics, and those expect strings.
     308              :     // It makes no sense for us to constantly turn the `TimelineId` and `TenantId` into
     309              :     // strings.
     310              :     tenant_id: String,
     311              :     shard_id: String,
     312              :     timeline_id: String,
     313              : }
     314              : 
     315              : #[derive(Debug, PartialEq, Clone, Copy)]
     316              : struct SlotHandle {
     317              :     /// Index into OPEN_FILES.slots
     318              :     index: usize,
     319              : 
     320              :     /// Value of 'tag' in the slot. If slot's tag doesn't match, then the slot has
     321              :     /// been recycled and no longer contains the FD for this virtual file.
     322              :     tag: u64,
     323              : }
     324              : 
     325              : /// OPEN_FILES is the global array that holds the physical file descriptors that
     326              : /// are currently open. Each slot in the array is protected by a separate lock,
     327              : /// so that different files can be accessed independently. The lock must be held
     328              : /// in write mode to replace the slot with a different file, but a read mode
     329              : /// is enough to operate on the file, whether you're reading or writing to it.
     330              : ///
     331              : /// OPEN_FILES starts in uninitialized state, and it's initialized by
     332              : /// the virtual_file::init() function. It must be called exactly once at page
     333              : /// server startup.
     334              : static OPEN_FILES: OnceCell<OpenFiles> = OnceCell::new();
     335              : 
     336              : struct OpenFiles {
     337              :     slots: &'static [Slot],
     338              : 
     339              :     /// clock arm for the clock algorithm
     340              :     next: AtomicUsize,
     341              : }
     342              : 
     343              : struct Slot {
     344              :     inner: RwLock<SlotInner>,
     345              : 
     346              :     /// has this file been used since last clock sweep?
     347              :     recently_used: AtomicBool,
     348              : }
     349              : 
     350              : struct SlotInner {
     351              :     /// Counter that's incremented every time a different file is stored here.
     352              :     /// To avoid the ABA problem.
     353              :     tag: u64,
     354              : 
     355              :     /// the underlying file
     356              :     file: Option<OwnedFd>,
     357              : }
     358              : 
     359              : /// Impl of [`tokio_epoll_uring::IoBuf`] and [`tokio_epoll_uring::IoBufMut`] for [`PageWriteGuard`].
     360              : struct PageWriteGuardBuf {
     361              :     page: PageWriteGuard<'static>,
     362              : }
     363              : // Safety: the [`PageWriteGuard`] gives us exclusive ownership of the page cache slot,
     364              : // and the location remains stable even if [`Self`] or the [`PageWriteGuard`] is moved.
     365              : // Page cache pages are zero-initialized, so, wrt uninitialized memory we're good.
     366              : // (Page cache tracks separately whether the contents are valid, see `PageWriteGuard::mark_valid`.)
     367              : unsafe impl tokio_epoll_uring::IoBuf for PageWriteGuardBuf {
     368       253377 :     fn stable_ptr(&self) -> *const u8 {
     369       253377 :         self.page.as_ptr()
     370       253377 :     }
     371       475101 :     fn bytes_init(&self) -> usize {
     372       475101 :         self.page.len()
     373       475101 :     }
     374       190071 :     fn bytes_total(&self) -> usize {
     375       190071 :         self.page.len()
     376       190071 :     }
     377              : }
     378              : // Safety: see above, plus: the ownership of [`PageWriteGuard`] means exclusive access,
     379              : // hence it's safe to hand out the `stable_mut_ptr()`.
     380              : unsafe impl tokio_epoll_uring::IoBufMut for PageWriteGuardBuf {
     381        95010 :     fn stable_mut_ptr(&mut self) -> *mut u8 {
     382        95010 :         self.page.as_mut_ptr()
     383        95010 :     }
     384              : 
     385        63357 :     unsafe fn set_init(&mut self, pos: usize) {
     386        63357 :         // There shouldn't really be any reason to call this API since bytes_init() == bytes_total().
     387        63357 :         assert!(pos <= self.page.len());
     388        63357 :     }
     389              : }
     390              : 
     391              : impl OpenFiles {
     392              :     /// Find a slot to use, evicting an existing file descriptor if needed.
     393              :     ///
     394              :     /// On return, we hold a lock on the slot, and its 'tag' has been updated
     395              :     /// recently_used has been set. It's all ready for reuse.
     396       390284 :     async fn find_victim_slot(&self) -> (SlotHandle, RwLockWriteGuard<SlotInner>) {
     397       390284 :         //
     398       390284 :         // Run the clock algorithm to find a slot to replace.
     399       390284 :         //
     400       390284 :         let num_slots = self.slots.len();
     401       390284 :         let mut retries = 0;
     402              :         let mut slot;
     403              :         let mut slot_guard;
     404              :         let index;
     405              :         loop {
     406      5058137 :             let next = self.next.fetch_add(1, Ordering::AcqRel) % num_slots;
     407      5058137 :             slot = &self.slots[next];
     408      5058137 : 
     409      5058137 :             // If the recently_used flag on this slot is set, continue the clock
     410      5058137 :             // sweep. Otherwise try to use this slot. If we cannot acquire the
     411      5058137 :             // lock, also continue the clock sweep.
     412      5058137 :             //
     413      5058137 :             // We only continue in this manner for a while, though. If we loop
     414      5058137 :             // through the array twice without finding a victim, just pick the
     415      5058137 :             // next slot and wait until we can reuse it. This way, we avoid
     416      5058137 :             // spinning in the extreme case that all the slots are busy with an
     417      5058137 :             // I/O operation.
     418      5058137 :             if retries < num_slots * 2 {
     419      4860430 :                 if !slot.recently_used.swap(false, Ordering::Release) {
     420      4444087 :                     if let Ok(guard) = slot.inner.try_write() {
     421       192577 :                         slot_guard = guard;
     422       192577 :                         index = next;
     423       192577 :                         break;
     424      4251510 :                     }
     425       416343 :                 }
     426      4667853 :                 retries += 1;
     427              :             } else {
     428       197707 :                 slot_guard = slot.inner.write().await;
     429       197707 :                 index = next;
     430       197707 :                 break;
     431              :             }
     432              :         }
     433              : 
     434              :         //
     435              :         // We now have the victim slot locked. If it was in use previously, close the
     436              :         // old file.
     437              :         //
     438       390284 :         if let Some(old_file) = slot_guard.file.take() {
     439       380306 :             // the normal path of dropping VirtualFile uses "close", use "close-by-replace" here to
     440       380306 :             // distinguish the two.
     441       380306 :             STORAGE_IO_TIME_METRIC
     442       380306 :                 .get(StorageIoOperation::CloseByReplace)
     443       380306 :                 .observe_closure_duration(|| drop(old_file));
     444       380306 :         }
     445              : 
     446              :         // Prepare the slot for reuse and return it
     447       390284 :         slot_guard.tag += 1;
     448       390284 :         slot.recently_used.store(true, Ordering::Relaxed);
     449       390284 :         (
     450       390284 :             SlotHandle {
     451       390284 :                 index,
     452       390284 :                 tag: slot_guard.tag,
     453       390284 :             },
     454       390284 :             slot_guard,
     455       390284 :         )
     456       390284 :     }
     457              : }
     458              : 
     459              : /// Identify error types that should alwways terminate the process.  Other
     460              : /// error types may be elegible for retry.
     461            8 : pub(crate) fn is_fatal_io_error(e: &std::io::Error) -> bool {
     462              :     use nix::errno::Errno::*;
     463            8 :     match e.raw_os_error().map(nix::errno::from_i32) {
     464              :         Some(EIO) => {
     465              :             // Terminate on EIO because we no longer trust the device to store
     466              :             // data safely, or to uphold persistence guarantees on fsync.
     467            0 :             true
     468              :         }
     469              :         Some(EROFS) => {
     470              :             // Terminate on EROFS because a filesystem is usually remounted
     471              :             // readonly when it has experienced some critical issue, so the same
     472              :             // logic as EIO applies.
     473            0 :             true
     474              :         }
     475              :         Some(EACCES) => {
     476              :             // Terminate on EACCESS because we should always have permissions
     477              :             // for our own data dir: if we don't, then we can't do our job and
     478              :             // need administrative intervention to fix permissions.  Terminating
     479              :             // is the best way to make sure we stop cleanly rather than going
     480              :             // into infinite retry loops, and will make it clear to the outside
     481              :             // world that we need help.
     482            0 :             true
     483              :         }
     484              :         _ => {
     485              :             // Treat all other local file I/O errors are retryable.  This includes:
     486              :             // - ENOSPC: we stay up and wait for eviction to free some space
     487              :             // - EINVAL, EBADF, EBADFD: this is a code bug, not a filesystem/hardware issue
     488              :             // - WriteZero, Interrupted: these are used internally VirtualFile
     489            8 :             false
     490              :         }
     491              :     }
     492            8 : }
     493              : 
     494              : /// Call this when the local filesystem gives us an error with an external
     495              : /// cause: this includes EIO, EROFS, and EACCESS: all these indicate either
     496              : /// bad storage or bad configuration, and we can't fix that from inside
     497              : /// a running process.
     498            0 : pub(crate) fn on_fatal_io_error(e: &std::io::Error, context: &str) -> ! {
     499            0 :     let backtrace = std::backtrace::Backtrace::force_capture();
     500            0 :     tracing::error!("Fatal I/O error: {e}: {context})\n{backtrace}");
     501            0 :     std::process::abort();
     502              : }
     503              : 
     504              : pub(crate) trait MaybeFatalIo<T> {
     505              :     fn maybe_fatal_err(self, context: &str) -> std::io::Result<T>;
     506              :     fn fatal_err(self, context: &str) -> T;
     507              : }
     508              : 
     509              : impl<T> MaybeFatalIo<T> for std::io::Result<T> {
     510              :     /// Terminate the process if the result is an error of a fatal type, else pass it through
     511              :     ///
     512              :     /// This is appropriate for writes, where we typically want to die on EIO/ACCES etc, but
     513              :     /// not on ENOSPC.
     514      4218606 :     fn maybe_fatal_err(self, context: &str) -> std::io::Result<T> {
     515      4218606 :         if let Err(e) = &self {
     516            8 :             if is_fatal_io_error(e) {
     517            0 :                 on_fatal_io_error(e, context);
     518            8 :             }
     519      4218598 :         }
     520      4218606 :         self
     521      4218606 :     }
     522              : 
     523              :     /// Terminate the process on any I/O error.
     524              :     ///
     525              :     /// This is appropriate for reads on files that we know exist: they should always work.
     526         4116 :     fn fatal_err(self, context: &str) -> T {
     527         4116 :         match self {
     528         4116 :             Ok(v) => v,
     529            0 :             Err(e) => {
     530            0 :                 on_fatal_io_error(&e, context);
     531              :             }
     532              :         }
     533         4116 :     }
     534              : }
     535              : 
     536              : /// Observe duration for the given storage I/O operation
     537              : ///
     538              : /// Unlike `observe_closure_duration`, this supports async,
     539              : /// where "support" means that we measure wall clock time.
     540              : macro_rules! observe_duration {
     541              :     ($op:expr, $($body:tt)*) => {{
     542              :         let instant = Instant::now();
     543              :         let result = $($body)*;
     544              :         let elapsed = instant.elapsed().as_secs_f64();
     545              :         STORAGE_IO_TIME_METRIC
     546              :             .get($op)
     547              :             .observe(elapsed);
     548              :         result
     549              :     }}
     550              : }
     551              : 
     552              : macro_rules! with_file {
     553              :     ($this:expr, $op:expr, | $ident:ident | $($body:tt)*) => {{
     554              :         let $ident = $this.lock_file().await?;
     555              :         observe_duration!($op, $($body)*)
     556              :     }};
     557              :     ($this:expr, $op:expr, | mut $ident:ident | $($body:tt)*) => {{
     558              :         let mut $ident = $this.lock_file().await?;
     559              :         observe_duration!($op, $($body)*)
     560              :     }};
     561              : }
     562              : 
     563              : impl VirtualFileInner {
     564              :     /// Open a file in read-only mode. Like File::open.
     565         2112 :     pub async fn open<P: AsRef<Utf8Path>>(
     566         2112 :         path: P,
     567         2112 :         ctx: &RequestContext,
     568         2112 :     ) -> Result<VirtualFileInner, std::io::Error> {
     569         2112 :         Self::open_with_options(path.as_ref(), OpenOptions::new().read(true), ctx).await
     570         2112 :     }
     571              : 
     572              :     /// Create a new file for writing. If the file exists, it will be truncated.
     573              :     /// Like File::create.
     574         3030 :     pub async fn create<P: AsRef<Utf8Path>>(
     575         3030 :         path: P,
     576         3030 :         ctx: &RequestContext,
     577         3030 :     ) -> Result<VirtualFileInner, std::io::Error> {
     578         3030 :         Self::open_with_options(
     579         3030 :             path.as_ref(),
     580         3030 :             OpenOptions::new().write(true).create(true).truncate(true),
     581         3030 :             ctx,
     582         3030 :         )
     583         3030 :         .await
     584         3030 :     }
     585              : 
     586              :     /// Open a file with given options.
     587              :     ///
     588              :     /// Note: If any custom flags were set in 'open_options' through OpenOptionsExt,
     589              :     /// they will be applied also when the file is subsequently re-opened, not only
     590              :     /// on the first time. Make sure that's sane!
     591        12202 :     pub async fn open_with_options<P: AsRef<Utf8Path>>(
     592        12202 :         path: P,
     593        12202 :         open_options: &OpenOptions,
     594        12202 :         _ctx: &RequestContext, /* TODO: carry a pointer to the metrics in the RequestContext instead of the parsing https://github.com/neondatabase/neon/issues/6107 */
     595        12202 :     ) -> Result<VirtualFileInner, std::io::Error> {
     596        12202 :         let path_ref = path.as_ref();
     597        12202 :         let path_str = path_ref.to_string();
     598        12202 :         let parts = path_str.split('/').collect::<Vec<&str>>();
     599        12202 :         let (tenant_id, shard_id, timeline_id) =
     600        12202 :             if parts.len() > 5 && parts[parts.len() - 5] == TENANTS_SEGMENT_NAME {
     601         9198 :                 let tenant_shard_part = parts[parts.len() - 4];
     602         9198 :                 let (tenant_id, shard_id) = match tenant_shard_part.parse::<TenantShardId>() {
     603         9198 :                     Ok(tenant_shard_id) => (
     604         9198 :                         tenant_shard_id.tenant_id.to_string(),
     605         9198 :                         format!("{}", tenant_shard_id.shard_slug()),
     606         9198 :                     ),
     607              :                     Err(_) => {
     608              :                         // Malformed path: this ID is just for observability, so tolerate it
     609              :                         // and pass through
     610            0 :                         (tenant_shard_part.to_string(), "*".to_string())
     611              :                     }
     612              :                 };
     613         9198 :                 (tenant_id, shard_id, parts[parts.len() - 2].to_string())
     614              :             } else {
     615         3004 :                 ("*".to_string(), "*".to_string(), "*".to_string())
     616              :             };
     617        12202 :         let (handle, mut slot_guard) = get_open_files().find_victim_slot().await;
     618              : 
     619              :         // NB: there is also StorageIoOperation::OpenAfterReplace which is for the case
     620              :         // where our caller doesn't get to use the returned VirtualFile before its
     621              :         // slot gets re-used by someone else.
     622        12202 :         let file = observe_duration!(StorageIoOperation::Open, {
     623        12202 :             open_options.open(path_ref.as_std_path()).await?
     624              :         });
     625              : 
     626              :         // Strip all options other than read and write.
     627              :         //
     628              :         // It would perhaps be nicer to check just for the read and write flags
     629              :         // explicitly, but OpenOptions doesn't contain any functions to read flags,
     630              :         // only to set them.
     631        12202 :         let mut reopen_options = open_options.clone();
     632        12202 :         reopen_options.create(false);
     633        12202 :         reopen_options.create_new(false);
     634        12202 :         reopen_options.truncate(false);
     635        12202 : 
     636        12202 :         let vfile = VirtualFileInner {
     637        12202 :             handle: RwLock::new(handle),
     638        12202 :             pos: 0,
     639        12202 :             path: path_ref.to_path_buf(),
     640        12202 :             open_options: reopen_options,
     641        12202 :             tenant_id,
     642        12202 :             shard_id,
     643        12202 :             timeline_id,
     644        12202 :         };
     645        12202 : 
     646        12202 :         // TODO: Under pressure, it's likely the slot will get re-used and
     647        12202 :         // the underlying file closed before they get around to using it.
     648        12202 :         // => https://github.com/neondatabase/neon/issues/6065
     649        12202 :         slot_guard.file.replace(file);
     650        12202 : 
     651        12202 :         Ok(vfile)
     652        12202 :     }
     653              : 
     654              :     /// Async version of [`::utils::crashsafe::overwrite`].
     655              :     ///
     656              :     /// # NB:
     657              :     ///
     658              :     /// Doesn't actually use the [`VirtualFile`] file descriptor cache, but,
     659              :     /// it did at an earlier time.
     660              :     /// And it will use this module's [`io_engine`] in the near future, so, leaving it here.
     661           56 :     pub async fn crashsafe_overwrite<B: BoundedBuf<Buf = Buf> + Send, Buf: IoBuf + Send>(
     662           56 :         final_path: Utf8PathBuf,
     663           56 :         tmp_path: Utf8PathBuf,
     664           56 :         content: B,
     665           56 :     ) -> std::io::Result<()> {
     666           56 :         // TODO: use tokio_epoll_uring if configured as `io_engine`.
     667           56 :         // See https://github.com/neondatabase/neon/issues/6663
     668           56 : 
     669           56 :         tokio::task::spawn_blocking(move || {
     670           56 :             let slice_storage;
     671           56 :             let content_len = content.bytes_init();
     672           56 :             let content = if content.bytes_init() > 0 {
     673           56 :                 slice_storage = Some(content.slice(0..content_len));
     674           56 :                 slice_storage.as_deref().expect("just set it to Some()")
     675              :             } else {
     676            0 :                 &[]
     677              :             };
     678           56 :             utils::crashsafe::overwrite(&final_path, &tmp_path, content)
     679           56 :                 .maybe_fatal_err("crashsafe_overwrite")
     680           56 :         })
     681           56 :         .await
     682           56 :         .expect("blocking task is never aborted")
     683           56 :     }
     684              : 
     685              :     /// Call File::sync_all() on the underlying File.
     686         5634 :     pub async fn sync_all(&self) -> Result<(), Error> {
     687         5634 :         with_file!(self, StorageIoOperation::Fsync, |file_guard| {
     688         5634 :             let (_file_guard, res) = io_engine::get().sync_all(file_guard).await;
     689         5634 :             res.maybe_fatal_err("sync_all")
     690              :         })
     691         5634 :     }
     692              : 
     693              :     /// Call File::sync_data() on the underlying File.
     694            0 :     pub async fn sync_data(&self) -> Result<(), Error> {
     695            0 :         with_file!(self, StorageIoOperation::Fsync, |file_guard| {
     696            0 :             let (_file_guard, res) = io_engine::get().sync_data(file_guard).await;
     697            0 :             res.maybe_fatal_err("sync_data")
     698              :         })
     699            0 :     }
     700              : 
     701         3604 :     pub async fn metadata(&self) -> Result<Metadata, Error> {
     702         3604 :         with_file!(self, StorageIoOperation::Metadata, |file_guard| {
     703         3604 :             let (_file_guard, res) = io_engine::get().metadata(file_guard).await;
     704         3604 :             res
     705              :         })
     706         3604 :     }
     707              : 
     708              :     /// Helper function internal to `VirtualFile` that looks up the underlying File,
     709              :     /// opens it and evicts some other File if necessary. The passed parameter is
     710              :     /// assumed to be a function available for the physical `File`.
     711              :     ///
     712              :     /// We are doing it via a macro as Rust doesn't support async closures that
     713              :     /// take on parameters with lifetimes.
     714      3250829 :     async fn lock_file(&self) -> Result<FileGuard, Error> {
     715      3250829 :         let open_files = get_open_files();
     716              : 
     717       378082 :         let mut handle_guard = {
     718              :             // Read the cached slot handle, and see if the slot that it points to still
     719              :             // contains our File.
     720              :             //
     721              :             // We only need to hold the handle lock while we read the current handle. If
     722              :             // another thread closes the file and recycles the slot for a different file,
     723              :             // we will notice that the handle we read is no longer valid and retry.
     724      3250829 :             let mut handle = *self.handle.read().await;
     725              :             loop {
     726              :                 // Check if the slot contains our File
     727              :                 {
     728      3456078 :                     let slot = &open_files.slots[handle.index];
     729      3456078 :                     let slot_guard = slot.inner.read().await;
     730      3456078 :                     if slot_guard.tag == handle.tag && slot_guard.file.is_some() {
     731              :                         // Found a cached file descriptor.
     732      2872747 :                         slot.recently_used.store(true, Ordering::Relaxed);
     733      2872747 :                         return Ok(FileGuard { slot_guard });
     734       583331 :                     }
     735              :                 }
     736              : 
     737              :                 // The slot didn't contain our File. We will have to open it ourselves,
     738              :                 // but before that, grab a write lock on handle in the VirtualFile, so
     739              :                 // that no other thread will try to concurrently open the same file.
     740       583331 :                 let handle_guard = self.handle.write().await;
     741              : 
     742              :                 // If another thread changed the handle while we were not holding the lock,
     743              :                 // then the handle might now be valid again. Loop back to retry.
     744       583331 :                 if *handle_guard != handle {
     745       205249 :                     handle = *handle_guard;
     746       205249 :                     continue;
     747       378082 :                 }
     748       378082 :                 break handle_guard;
     749              :             }
     750              :         };
     751              : 
     752              :         // We need to open the file ourselves. The handle in the VirtualFile is
     753              :         // now locked in write-mode. Find a free slot to put it in.
     754       378082 :         let (handle, mut slot_guard) = open_files.find_victim_slot().await;
     755              : 
     756              :         // Re-open the physical file.
     757              :         // NB: we use StorageIoOperation::OpenAferReplace for this to distinguish this
     758              :         // case from StorageIoOperation::Open. This helps with identifying thrashing
     759              :         // of the virtual file descriptor cache.
     760       378082 :         let file = observe_duration!(StorageIoOperation::OpenAfterReplace, {
     761       378082 :             self.open_options.open(self.path.as_std_path()).await?
     762              :         });
     763              : 
     764              :         // Store the File in the slot and update the handle in the VirtualFile
     765              :         // to point to it.
     766       378082 :         slot_guard.file.replace(file);
     767       378082 : 
     768       378082 :         *handle_guard = handle;
     769       378082 : 
     770       378082 :         Ok(FileGuard {
     771       378082 :             slot_guard: slot_guard.downgrade(),
     772       378082 :         })
     773      3250829 :     }
     774              : 
     775          516 :     pub fn remove(self) {
     776          516 :         let path = self.path.clone();
     777          516 :         drop(self);
     778          516 :         std::fs::remove_file(path).expect("failed to remove the virtual file");
     779          516 :     }
     780              : 
     781        11352 :     pub async fn seek(&mut self, pos: SeekFrom) -> Result<u64, Error> {
     782        11352 :         match pos {
     783        11332 :             SeekFrom::Start(offset) => {
     784        11332 :                 self.pos = offset;
     785        11332 :             }
     786            8 :             SeekFrom::End(offset) => {
     787            8 :                 self.pos = with_file!(self, StorageIoOperation::Seek, |mut file_guard| file_guard
     788            8 :                     .with_std_file_mut(|std_file| std_file.seek(SeekFrom::End(offset))))?
     789              :             }
     790           12 :             SeekFrom::Current(offset) => {
     791           12 :                 let pos = self.pos as i128 + offset as i128;
     792           12 :                 if pos < 0 {
     793            4 :                     return Err(Error::new(
     794            4 :                         ErrorKind::InvalidInput,
     795            4 :                         "offset would be negative",
     796            4 :                     ));
     797            8 :                 }
     798            8 :                 if pos > u64::MAX as i128 {
     799            0 :                     return Err(Error::new(ErrorKind::InvalidInput, "offset overflow"));
     800            8 :                 }
     801            8 :                 self.pos = pos as u64;
     802              :             }
     803              :         }
     804        11344 :         Ok(self.pos)
     805        11352 :     }
     806              : 
     807              :     /// Read the file contents in range `offset..(offset + slice.bytes_total())` into `slice[0..slice.bytes_total()]`.
     808              :     ///
     809              :     /// The returned `Slice<Buf>` is equivalent to the input `slice`, i.e., it's the same view into the same buffer.
     810       966759 :     pub async fn read_exact_at<Buf>(
     811       966759 :         &self,
     812       966759 :         slice: Slice<Buf>,
     813       966759 :         offset: u64,
     814       966759 :         ctx: &RequestContext,
     815       966759 :     ) -> Result<Slice<Buf>, Error>
     816       966759 :     where
     817       966759 :         Buf: IoBufAlignedMut + Send,
     818       966759 :     {
     819       966759 :         let assert_we_return_original_bounds = if cfg!(debug_assertions) {
     820       966759 :             Some((slice.stable_ptr() as usize, slice.bytes_total()))
     821              :         } else {
     822            0 :             None
     823              :         };
     824              : 
     825       966759 :         let original_bounds = slice.bounds();
     826       966759 :         let (buf, res) =
     827       966759 :             read_exact_at_impl(slice, offset, |buf, offset| self.read_at(buf, offset, ctx)).await;
     828       966759 :         let res = res.map(|_| buf.slice(original_bounds));
     829              : 
     830       966759 :         if let Some(original_bounds) = assert_we_return_original_bounds {
     831       966759 :             if let Ok(slice) = &res {
     832       966759 :                 let returned_bounds = (slice.stable_ptr() as usize, slice.bytes_total());
     833       966759 :                 assert_eq!(original_bounds, returned_bounds);
     834            0 :             }
     835            0 :         }
     836              : 
     837       966759 :         res
     838       966759 :     }
     839              : 
     840              :     /// Like [`Self::read_exact_at`] but for [`PageWriteGuard`].
     841        63357 :     pub async fn read_exact_at_page(
     842        63357 :         &self,
     843        63357 :         page: PageWriteGuard<'static>,
     844        63357 :         offset: u64,
     845        63357 :         ctx: &RequestContext,
     846        63357 :     ) -> Result<PageWriteGuard<'static>, Error> {
     847        63357 :         let buf = PageWriteGuardBuf { page }.slice_full();
     848        63357 :         debug_assert_eq!(buf.bytes_total(), PAGE_SZ);
     849        63357 :         self.read_exact_at(buf, offset, ctx)
     850        63357 :             .await
     851        63357 :             .map(|slice| slice.into_inner().page)
     852        63357 :     }
     853              : 
     854              :     // Copied from https://doc.rust-lang.org/1.72.0/src/std/os/unix/fs.rs.html#219-235
     855        13222 :     pub async fn write_all_at<Buf: IoBuf + Send>(
     856        13222 :         &self,
     857        13222 :         buf: FullSlice<Buf>,
     858        13222 :         mut offset: u64,
     859        13222 :         ctx: &RequestContext,
     860        13222 :     ) -> (FullSlice<Buf>, Result<(), Error>) {
     861        13222 :         let buf = buf.into_raw_slice();
     862        13222 :         let bounds = buf.bounds();
     863        13222 :         let restore =
     864        13222 :             |buf: Slice<_>| FullSlice::must_new(Slice::from_buf_bounds(buf.into_inner(), bounds));
     865        13222 :         let mut buf = buf;
     866        26444 :         while !buf.is_empty() {
     867        13222 :             let (tmp, res) = self.write_at(FullSlice::must_new(buf), offset, ctx).await;
     868        13222 :             buf = tmp.into_raw_slice();
     869            0 :             match res {
     870              :                 Ok(0) => {
     871            0 :                     return (
     872            0 :                         restore(buf),
     873            0 :                         Err(Error::new(
     874            0 :                             std::io::ErrorKind::WriteZero,
     875            0 :                             "failed to write whole buffer",
     876            0 :                         )),
     877            0 :                     );
     878              :                 }
     879        13222 :                 Ok(n) => {
     880        13222 :                     buf = buf.slice(n..);
     881        13222 :                     offset += n as u64;
     882        13222 :                 }
     883            0 :                 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
     884            0 :                 Err(e) => return (restore(buf), Err(e)),
     885              :             }
     886              :         }
     887        13222 :         (restore(buf), Ok(()))
     888        13222 :     }
     889              : 
     890              :     /// Writes `buf` to the file at the current offset.
     891              :     ///
     892              :     /// Panics if there is an uninitialized range in `buf`, as that is most likely a bug in the caller.
     893      2260790 :     pub async fn write_all<Buf: IoBuf + Send>(
     894      2260790 :         &mut self,
     895      2260790 :         buf: FullSlice<Buf>,
     896      2260790 :         ctx: &RequestContext,
     897      2260790 :     ) -> (FullSlice<Buf>, Result<usize, Error>) {
     898      2260790 :         let buf = buf.into_raw_slice();
     899      2260790 :         let bounds = buf.bounds();
     900      2260790 :         let restore =
     901      2260790 :             |buf: Slice<_>| FullSlice::must_new(Slice::from_buf_bounds(buf.into_inner(), bounds));
     902      2260790 :         let nbytes = buf.len();
     903      2260790 :         let mut buf = buf;
     904      4521500 :         while !buf.is_empty() {
     905      2260714 :             let (tmp, res) = self.write(FullSlice::must_new(buf), ctx).await;
     906      2260714 :             buf = tmp.into_raw_slice();
     907            4 :             match res {
     908              :                 Ok(0) => {
     909            0 :                     return (
     910            0 :                         restore(buf),
     911            0 :                         Err(Error::new(
     912            0 :                             std::io::ErrorKind::WriteZero,
     913            0 :                             "failed to write whole buffer",
     914            0 :                         )),
     915            0 :                     );
     916              :                 }
     917      2260710 :                 Ok(n) => {
     918      2260710 :                     buf = buf.slice(n..);
     919      2260710 :                 }
     920            4 :                 Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {}
     921            4 :                 Err(e) => return (restore(buf), Err(e)),
     922              :             }
     923              :         }
     924      2260786 :         (restore(buf), Ok(nbytes))
     925      2260790 :     }
     926              : 
     927      2260714 :     async fn write<B: IoBuf + Send>(
     928      2260714 :         &mut self,
     929      2260714 :         buf: FullSlice<B>,
     930      2260714 :         ctx: &RequestContext,
     931      2260714 :     ) -> (FullSlice<B>, Result<usize, std::io::Error>) {
     932      2260714 :         let pos = self.pos;
     933      2260714 :         let (buf, res) = self.write_at(buf, pos, ctx).await;
     934      2260714 :         let n = match res {
     935      2260710 :             Ok(n) => n,
     936            4 :             Err(e) => return (buf, Err(e)),
     937              :         };
     938      2260710 :         self.pos += n as u64;
     939      2260710 :         (buf, Ok(n))
     940      2260714 :     }
     941              : 
     942       967647 :     pub(crate) async fn read_at<Buf>(
     943       967647 :         &self,
     944       967647 :         buf: tokio_epoll_uring::Slice<Buf>,
     945       967647 :         offset: u64,
     946       967647 :         _ctx: &RequestContext, /* TODO: use for metrics: https://github.com/neondatabase/neon/issues/6107 */
     947       967647 :     ) -> (tokio_epoll_uring::Slice<Buf>, Result<usize, Error>)
     948       967647 :     where
     949       967647 :         Buf: tokio_epoll_uring::IoBufMut + Send,
     950       967647 :     {
     951       967647 :         let file_guard = match self
     952       967647 :             .lock_file()
     953       967647 :             .await
     954       967647 :             .maybe_fatal_err("lock_file inside VirtualFileInner::read_at")
     955              :         {
     956       967647 :             Ok(file_guard) => file_guard,
     957            0 :             Err(e) => return (buf, Err(e)),
     958              :         };
     959              : 
     960       967647 :         observe_duration!(StorageIoOperation::Read, {
     961       967647 :             let ((_file_guard, buf), res) = io_engine::get().read_at(file_guard, offset, buf).await;
     962       967647 :             let res = res.maybe_fatal_err("io_engine read_at inside VirtualFileInner::read_at");
     963       967647 :             if let Ok(size) = res {
     964       967643 :                 STORAGE_IO_SIZE
     965       967643 :                     .with_label_values(&[
     966       967643 :                         "read",
     967       967643 :                         &self.tenant_id,
     968       967643 :                         &self.shard_id,
     969       967643 :                         &self.timeline_id,
     970       967643 :                     ])
     971       967643 :                     .add(size as i64);
     972       967643 :             }
     973       967647 :             (buf, res)
     974              :         })
     975       967647 :     }
     976              : 
     977              :     /// The function aborts the process if the error is fatal.
     978      2273936 :     async fn write_at<B: IoBuf + Send>(
     979      2273936 :         &self,
     980      2273936 :         buf: FullSlice<B>,
     981      2273936 :         offset: u64,
     982      2273936 :         _ctx: &RequestContext, /* TODO: use for metrics: https://github.com/neondatabase/neon/issues/6107 */
     983      2273936 :     ) -> (FullSlice<B>, Result<usize, Error>) {
     984      2273936 :         let (slice, result) = self.write_at_inner(buf, offset, _ctx).await;
     985      2273936 :         let result = result.maybe_fatal_err("write_at");
     986      2273936 :         (slice, result)
     987      2273936 :     }
     988              : 
     989      2273936 :     async fn write_at_inner<B: IoBuf + Send>(
     990      2273936 :         &self,
     991      2273936 :         buf: FullSlice<B>,
     992      2273936 :         offset: u64,
     993      2273936 :         _ctx: &RequestContext, /* TODO: use for metrics: https://github.com/neondatabase/neon/issues/6107 */
     994      2273936 :     ) -> (FullSlice<B>, Result<usize, Error>) {
     995      2273936 :         let file_guard = match self.lock_file().await {
     996      2273936 :             Ok(file_guard) => file_guard,
     997            0 :             Err(e) => return (buf, Err(e)),
     998              :         };
     999      2273936 :         observe_duration!(StorageIoOperation::Write, {
    1000      2273936 :             let ((_file_guard, buf), result) =
    1001      2273936 :                 io_engine::get().write_at(file_guard, offset, buf).await;
    1002      2273936 :             if let Ok(size) = result {
    1003      2273932 :                 STORAGE_IO_SIZE
    1004      2273932 :                     .with_label_values(&[
    1005      2273932 :                         "write",
    1006      2273932 :                         &self.tenant_id,
    1007      2273932 :                         &self.shard_id,
    1008      2273932 :                         &self.timeline_id,
    1009      2273932 :                     ])
    1010      2273932 :                     .add(size as i64);
    1011      2273932 :             }
    1012      2273936 :             (buf, result)
    1013              :         })
    1014      2273936 :     }
    1015              : 
    1016          448 :     async fn read_to_end(&mut self, buf: &mut Vec<u8>, ctx: &RequestContext) -> Result<(), Error> {
    1017          448 :         let mut tmp = vec![0; 128];
    1018              :         loop {
    1019          888 :             let slice = tmp.slice(..128);
    1020          888 :             let (slice, res) = self.read_at(slice, self.pos, ctx).await;
    1021            4 :             match res {
    1022          444 :                 Ok(0) => return Ok(()),
    1023          440 :                 Ok(n) => {
    1024          440 :                     self.pos += n as u64;
    1025          440 :                     buf.extend_from_slice(&slice[..n]);
    1026          440 :                 }
    1027            4 :                 Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {}
    1028            4 :                 Err(e) => return Err(e),
    1029              :             }
    1030          440 :             tmp = slice.into_inner();
    1031              :         }
    1032          448 :     }
    1033              : }
    1034              : 
    1035              : // Adapted from https://doc.rust-lang.org/1.72.0/src/std/os/unix/fs.rs.html#117-135
    1036       966775 : pub async fn read_exact_at_impl<Buf, F, Fut>(
    1037       966775 :     mut buf: tokio_epoll_uring::Slice<Buf>,
    1038       966775 :     mut offset: u64,
    1039       966775 :     mut read_at: F,
    1040       966775 : ) -> (Buf, std::io::Result<()>)
    1041       966775 : where
    1042       966775 :     Buf: IoBufMut + Send,
    1043       966775 :     F: FnMut(tokio_epoll_uring::Slice<Buf>, u64) -> Fut,
    1044       966775 :     Fut: std::future::Future<Output = (tokio_epoll_uring::Slice<Buf>, std::io::Result<usize>)>,
    1045       966775 : {
    1046      1933554 :     while buf.bytes_total() != 0 {
    1047              :         let res;
    1048       966783 :         (buf, res) = read_at(buf, offset).await;
    1049            0 :         match res {
    1050            4 :             Ok(0) => break,
    1051       966779 :             Ok(n) => {
    1052       966779 :                 buf = buf.slice(n..);
    1053       966779 :                 offset += n as u64;
    1054       966779 :             }
    1055            0 :             Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {}
    1056            0 :             Err(e) => return (buf.into_inner(), Err(e)),
    1057              :         }
    1058              :     }
    1059              :     // NB: don't use `buf.is_empty()` here; it is from the
    1060              :     // `impl Deref for Slice { Target = [u8] }`; the &[u8]
    1061              :     // returned by it only covers the initialized portion of `buf`.
    1062              :     // Whereas we're interested in ensuring that we filled the entire
    1063              :     // buffer that the user passed in.
    1064       966775 :     if buf.bytes_total() != 0 {
    1065            4 :         (
    1066            4 :             buf.into_inner(),
    1067            4 :             Err(std::io::Error::new(
    1068            4 :                 std::io::ErrorKind::UnexpectedEof,
    1069            4 :                 "failed to fill whole buffer",
    1070            4 :             )),
    1071            4 :         )
    1072              :     } else {
    1073       966771 :         assert_eq!(buf.len(), buf.bytes_total());
    1074       966771 :         (buf.into_inner(), Ok(()))
    1075              :     }
    1076       966775 : }
    1077              : 
    1078              : #[cfg(test)]
    1079              : mod test_read_exact_at_impl {
    1080              : 
    1081              :     use std::{collections::VecDeque, sync::Arc};
    1082              : 
    1083              :     use tokio_epoll_uring::{BoundedBuf, BoundedBufMut};
    1084              : 
    1085              :     use super::read_exact_at_impl;
    1086              : 
    1087              :     struct Expectation {
    1088              :         offset: u64,
    1089              :         bytes_total: usize,
    1090              :         result: std::io::Result<Vec<u8>>,
    1091              :     }
    1092              :     struct MockReadAt {
    1093              :         expectations: VecDeque<Expectation>,
    1094              :     }
    1095              : 
    1096              :     impl MockReadAt {
    1097           24 :         async fn read_at(
    1098           24 :             &mut self,
    1099           24 :             mut buf: tokio_epoll_uring::Slice<Vec<u8>>,
    1100           24 :             offset: u64,
    1101           24 :         ) -> (tokio_epoll_uring::Slice<Vec<u8>>, std::io::Result<usize>) {
    1102           24 :             let exp = self
    1103           24 :                 .expectations
    1104           24 :                 .pop_front()
    1105           24 :                 .expect("read_at called but we have no expectations left");
    1106           24 :             assert_eq!(exp.offset, offset);
    1107           24 :             assert_eq!(exp.bytes_total, buf.bytes_total());
    1108           24 :             match exp.result {
    1109           24 :                 Ok(bytes) => {
    1110           24 :                     assert!(bytes.len() <= buf.bytes_total());
    1111           24 :                     buf.put_slice(&bytes);
    1112           24 :                     (buf, Ok(bytes.len()))
    1113              :                 }
    1114            0 :                 Err(e) => (buf, Err(e)),
    1115              :             }
    1116           24 :         }
    1117              :     }
    1118              : 
    1119              :     impl Drop for MockReadAt {
    1120           16 :         fn drop(&mut self) {
    1121           16 :             assert_eq!(self.expectations.len(), 0);
    1122           16 :         }
    1123              :     }
    1124              : 
    1125              :     #[tokio::test]
    1126            4 :     async fn test_basic() {
    1127            4 :         let buf = Vec::with_capacity(5).slice_full();
    1128            4 :         let mock_read_at = Arc::new(tokio::sync::Mutex::new(MockReadAt {
    1129            4 :             expectations: VecDeque::from(vec![Expectation {
    1130            4 :                 offset: 0,
    1131            4 :                 bytes_total: 5,
    1132            4 :                 result: Ok(vec![b'a', b'b', b'c', b'd', b'e']),
    1133            4 :             }]),
    1134            4 :         }));
    1135            4 :         let (buf, res) = read_exact_at_impl(buf, 0, |buf, offset| {
    1136            4 :             let mock_read_at = Arc::clone(&mock_read_at);
    1137            4 :             async move { mock_read_at.lock().await.read_at(buf, offset).await }
    1138            4 :         })
    1139            4 :         .await;
    1140            4 :         assert!(res.is_ok());
    1141            4 :         assert_eq!(buf, vec![b'a', b'b', b'c', b'd', b'e']);
    1142            4 :     }
    1143              : 
    1144              :     #[tokio::test]
    1145            4 :     async fn test_empty_buf_issues_no_syscall() {
    1146            4 :         let buf = Vec::new().slice_full();
    1147            4 :         let mock_read_at = Arc::new(tokio::sync::Mutex::new(MockReadAt {
    1148            4 :             expectations: VecDeque::new(),
    1149            4 :         }));
    1150            4 :         let (_buf, res) = read_exact_at_impl(buf, 0, |buf, offset| {
    1151            0 :             let mock_read_at = Arc::clone(&mock_read_at);
    1152            4 :             async move { mock_read_at.lock().await.read_at(buf, offset).await }
    1153            4 :         })
    1154            4 :         .await;
    1155            4 :         assert!(res.is_ok());
    1156            4 :     }
    1157              : 
    1158              :     #[tokio::test]
    1159            4 :     async fn test_two_read_at_calls_needed_until_buf_filled() {
    1160            4 :         let buf = Vec::with_capacity(4).slice_full();
    1161            4 :         let mock_read_at = Arc::new(tokio::sync::Mutex::new(MockReadAt {
    1162            4 :             expectations: VecDeque::from(vec![
    1163            4 :                 Expectation {
    1164            4 :                     offset: 0,
    1165            4 :                     bytes_total: 4,
    1166            4 :                     result: Ok(vec![b'a', b'b']),
    1167            4 :                 },
    1168            4 :                 Expectation {
    1169            4 :                     offset: 2,
    1170            4 :                     bytes_total: 2,
    1171            4 :                     result: Ok(vec![b'c', b'd']),
    1172            4 :                 },
    1173            4 :             ]),
    1174            4 :         }));
    1175            8 :         let (buf, res) = read_exact_at_impl(buf, 0, |buf, offset| {
    1176            8 :             let mock_read_at = Arc::clone(&mock_read_at);
    1177            8 :             async move { mock_read_at.lock().await.read_at(buf, offset).await }
    1178            8 :         })
    1179            4 :         .await;
    1180            4 :         assert!(res.is_ok());
    1181            4 :         assert_eq!(buf, vec![b'a', b'b', b'c', b'd']);
    1182            4 :     }
    1183              : 
    1184              :     #[tokio::test]
    1185            4 :     async fn test_eof_before_buffer_full() {
    1186            4 :         let buf = Vec::with_capacity(3).slice_full();
    1187            4 :         let mock_read_at = Arc::new(tokio::sync::Mutex::new(MockReadAt {
    1188            4 :             expectations: VecDeque::from(vec![
    1189            4 :                 Expectation {
    1190            4 :                     offset: 0,
    1191            4 :                     bytes_total: 3,
    1192            4 :                     result: Ok(vec![b'a']),
    1193            4 :                 },
    1194            4 :                 Expectation {
    1195            4 :                     offset: 1,
    1196            4 :                     bytes_total: 2,
    1197            4 :                     result: Ok(vec![b'b']),
    1198            4 :                 },
    1199            4 :                 Expectation {
    1200            4 :                     offset: 2,
    1201            4 :                     bytes_total: 1,
    1202            4 :                     result: Ok(vec![]),
    1203            4 :                 },
    1204            4 :             ]),
    1205            4 :         }));
    1206           12 :         let (_buf, res) = read_exact_at_impl(buf, 0, |buf, offset| {
    1207           12 :             let mock_read_at = Arc::clone(&mock_read_at);
    1208           12 :             async move { mock_read_at.lock().await.read_at(buf, offset).await }
    1209           12 :         })
    1210            4 :         .await;
    1211            4 :         let Err(err) = res else {
    1212            4 :             panic!("should return an error");
    1213            4 :         };
    1214            4 :         assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
    1215            4 :         assert_eq!(format!("{err}"), "failed to fill whole buffer");
    1216            4 :         // buffer contents on error are unspecified
    1217            4 :     }
    1218              : }
    1219              : 
    1220              : struct FileGuard {
    1221              :     slot_guard: RwLockReadGuard<'static, SlotInner>,
    1222              : }
    1223              : 
    1224              : impl AsRef<OwnedFd> for FileGuard {
    1225      3250829 :     fn as_ref(&self) -> &OwnedFd {
    1226      3250829 :         // This unwrap is safe because we only create `FileGuard`s
    1227      3250829 :         // if we know that the file is Some.
    1228      3250829 :         self.slot_guard.file.as_ref().unwrap()
    1229      3250829 :     }
    1230              : }
    1231              : 
    1232              : impl FileGuard {
    1233              :     /// Soft deprecation: we'll move VirtualFile to async APIs and remove this function eventually.
    1234      1625124 :     fn with_std_file<F, R>(&self, with: F) -> R
    1235      1625124 :     where
    1236      1625124 :         F: FnOnce(&File) -> R,
    1237      1625124 :     {
    1238      1625124 :         // SAFETY:
    1239      1625124 :         // - lifetime of the fd: `file` doesn't outlive the OwnedFd stored in `self`.
    1240      1625124 :         // - `&` usage below: `self` is `&`, hence Rust typesystem guarantees there are is no `&mut`
    1241      1625124 :         let file = unsafe { File::from_raw_fd(self.as_ref().as_raw_fd()) };
    1242      1625124 :         let res = with(&file);
    1243      1625124 :         let _ = file.into_raw_fd();
    1244      1625124 :         res
    1245      1625124 :     }
    1246              :     /// Soft deprecation: we'll move VirtualFile to async APIs and remove this function eventually.
    1247            8 :     fn with_std_file_mut<F, R>(&mut self, with: F) -> R
    1248            8 :     where
    1249            8 :         F: FnOnce(&mut File) -> R,
    1250            8 :     {
    1251            8 :         // SAFETY:
    1252            8 :         // - lifetime of the fd: `file` doesn't outlive the OwnedFd stored in `self`.
    1253            8 :         // - &mut usage below: `self` is `&mut`, hence this call is the only task/thread that has control over the underlying fd
    1254            8 :         let mut file = unsafe { File::from_raw_fd(self.as_ref().as_raw_fd()) };
    1255            8 :         let res = with(&mut file);
    1256            8 :         let _ = file.into_raw_fd();
    1257            8 :         res
    1258            8 :     }
    1259              : }
    1260              : 
    1261              : impl tokio_epoll_uring::IoFd for FileGuard {
    1262      1625697 :     unsafe fn as_fd(&self) -> RawFd {
    1263      1625697 :         let owned_fd: &OwnedFd = self.as_ref();
    1264      1625697 :         owned_fd.as_raw_fd()
    1265      1625697 :     }
    1266              : }
    1267              : 
    1268              : #[cfg(test)]
    1269              : impl VirtualFile {
    1270        41832 :     pub(crate) async fn read_blk(
    1271        41832 :         &self,
    1272        41832 :         blknum: u32,
    1273        41832 :         ctx: &RequestContext,
    1274        41832 :     ) -> Result<crate::tenant::block_io::BlockLease<'_>, std::io::Error> {
    1275        41832 :         self.inner.read_blk(blknum, ctx).await
    1276        41832 :     }
    1277              : }
    1278              : 
    1279              : #[cfg(test)]
    1280              : impl VirtualFileInner {
    1281        41832 :     pub(crate) async fn read_blk(
    1282        41832 :         &self,
    1283        41832 :         blknum: u32,
    1284        41832 :         ctx: &RequestContext,
    1285        41832 :     ) -> Result<crate::tenant::block_io::BlockLease<'_>, std::io::Error> {
    1286              :         use crate::page_cache::PAGE_SZ;
    1287        41832 :         let slice = IoBufferMut::with_capacity(PAGE_SZ).slice_full();
    1288        41832 :         assert_eq!(slice.bytes_total(), PAGE_SZ);
    1289        41832 :         let slice = self
    1290        41832 :             .read_exact_at(slice, blknum as u64 * (PAGE_SZ as u64), ctx)
    1291        41832 :             .await?;
    1292        41832 :         Ok(crate::tenant::block_io::BlockLease::IoBufferMut(
    1293        41832 :             slice.into_inner(),
    1294        41832 :         ))
    1295        41832 :     }
    1296              : }
    1297              : 
    1298              : impl Drop for VirtualFileInner {
    1299              :     /// If a VirtualFile is dropped, close the underlying file if it was open.
    1300        10590 :     fn drop(&mut self) {
    1301        10590 :         let handle = self.handle.get_mut();
    1302              : 
    1303        10590 :         fn clean_slot(slot: &Slot, mut slot_guard: RwLockWriteGuard<'_, SlotInner>, tag: u64) {
    1304        10590 :             if slot_guard.tag == tag {
    1305         9436 :                 slot.recently_used.store(false, Ordering::Relaxed);
    1306              :                 // there is also operation "close-by-replace" for closes done on eviction for
    1307              :                 // comparison.
    1308         9436 :                 if let Some(fd) = slot_guard.file.take() {
    1309         9436 :                     STORAGE_IO_TIME_METRIC
    1310         9436 :                         .get(StorageIoOperation::Close)
    1311         9436 :                         .observe_closure_duration(|| drop(fd));
    1312         9436 :                 }
    1313         1154 :             }
    1314        10590 :         }
    1315              : 
    1316              :         // We don't have async drop so we cannot directly await the lock here.
    1317              :         // Instead, first do a best-effort attempt at closing the underlying
    1318              :         // file descriptor by using `try_write`, and if that fails, spawn
    1319              :         // a tokio task to do it asynchronously: we just want it to be
    1320              :         // cleaned up eventually.
    1321              :         // Most of the time, the `try_lock` should succeed though,
    1322              :         // as we have `&mut self` access. In other words, if the slot
    1323              :         // is still occupied by our file, there should be no access from
    1324              :         // other I/O operations; the only other possible place to lock
    1325              :         // the slot is the lock algorithm looking for free slots.
    1326        10590 :         let slot = &get_open_files().slots[handle.index];
    1327        10590 :         if let Ok(slot_guard) = slot.inner.try_write() {
    1328        10590 :             clean_slot(slot, slot_guard, handle.tag);
    1329        10590 :         } else {
    1330            0 :             let tag = handle.tag;
    1331            0 :             tokio::spawn(async move {
    1332            0 :                 let slot_guard = slot.inner.write().await;
    1333            0 :                 clean_slot(slot, slot_guard, tag);
    1334            0 :             });
    1335            0 :         };
    1336        10590 :     }
    1337              : }
    1338              : 
    1339              : impl OwnedAsyncWriter for VirtualFile {
    1340        13214 :     async fn write_all_at<Buf: IoBufAligned + Send>(
    1341        13214 :         &self,
    1342        13214 :         buf: FullSlice<Buf>,
    1343        13214 :         offset: u64,
    1344        13214 :         ctx: &RequestContext,
    1345        13214 :     ) -> std::io::Result<FullSlice<Buf>> {
    1346        13214 :         let (buf, res) = VirtualFile::write_all_at(self, buf, offset, ctx).await;
    1347        13214 :         res.map(|_| buf)
    1348        13214 :     }
    1349              : }
    1350              : 
    1351              : impl OpenFiles {
    1352          464 :     fn new(num_slots: usize) -> OpenFiles {
    1353          464 :         let mut slots = Box::new(Vec::with_capacity(num_slots));
    1354         4640 :         for _ in 0..num_slots {
    1355         4640 :             let slot = Slot {
    1356         4640 :                 recently_used: AtomicBool::new(false),
    1357         4640 :                 inner: RwLock::new(SlotInner { tag: 0, file: None }),
    1358         4640 :             };
    1359         4640 :             slots.push(slot);
    1360         4640 :         }
    1361              : 
    1362          464 :         OpenFiles {
    1363          464 :             next: AtomicUsize::new(0),
    1364          464 :             slots: Box::leak(slots),
    1365          464 :         }
    1366          464 :     }
    1367              : }
    1368              : 
    1369              : ///
    1370              : /// Initialize the virtual file module. This must be called once at page
    1371              : /// server startup.
    1372              : ///
    1373              : #[cfg(not(test))]
    1374            0 : pub fn init(num_slots: usize, engine: IoEngineKind, mode: IoMode, sync_mode: SyncMode) {
    1375            0 :     if OPEN_FILES.set(OpenFiles::new(num_slots)).is_err() {
    1376            0 :         panic!("virtual_file::init called twice");
    1377            0 :     }
    1378            0 :     set_io_mode(mode);
    1379            0 :     io_engine::init(engine);
    1380            0 :     SYNC_MODE.store(sync_mode as u8, std::sync::atomic::Ordering::Relaxed);
    1381            0 :     crate::metrics::virtual_file_descriptor_cache::SIZE_MAX.set(num_slots as u64);
    1382            0 : }
    1383              : 
    1384              : const TEST_MAX_FILE_DESCRIPTORS: usize = 10;
    1385              : 
    1386              : // Get a handle to the global slots array.
    1387      3273621 : fn get_open_files() -> &'static OpenFiles {
    1388      3273621 :     //
    1389      3273621 :     // In unit tests, page server startup doesn't happen and no one calls
    1390      3273621 :     // virtual_file::init(). Initialize it here, with a small array.
    1391      3273621 :     //
    1392      3273621 :     // This applies to the virtual file tests below, but all other unit
    1393      3273621 :     // tests too, so the virtual file facility is always usable in
    1394      3273621 :     // unit tests.
    1395      3273621 :     //
    1396      3273621 :     if cfg!(test) {
    1397      3273621 :         OPEN_FILES.get_or_init(|| OpenFiles::new(TEST_MAX_FILE_DESCRIPTORS))
    1398              :     } else {
    1399            0 :         OPEN_FILES.get().expect("virtual_file::init not called yet")
    1400              :     }
    1401      3273621 : }
    1402              : 
    1403              : /// Gets the io buffer alignment.
    1404            0 : pub(crate) const fn get_io_buffer_alignment() -> usize {
    1405            0 :     DEFAULT_IO_BUFFER_ALIGNMENT
    1406            0 : }
    1407              : 
    1408              : pub(crate) type IoBufferMut = AlignedBufferMut<ConstAlign<{ get_io_buffer_alignment() }>>;
    1409              : pub(crate) type IoBuffer = AlignedBuffer<ConstAlign<{ get_io_buffer_alignment() }>>;
    1410              : pub(crate) type IoPageSlice<'a> =
    1411              :     AlignedSlice<'a, PAGE_SZ, ConstAlign<{ get_io_buffer_alignment() }>>;
    1412              : 
    1413              : static IO_MODE: AtomicU8 = AtomicU8::new(IoMode::preferred() as u8);
    1414              : 
    1415            0 : pub(crate) fn set_io_mode(mode: IoMode) {
    1416            0 :     IO_MODE.store(mode as u8, std::sync::atomic::Ordering::Relaxed);
    1417            0 : }
    1418              : 
    1419         5068 : pub(crate) fn get_io_mode() -> IoMode {
    1420         5068 :     IoMode::try_from(IO_MODE.load(Ordering::Relaxed)).unwrap()
    1421         5068 : }
    1422              : 
    1423              : static SYNC_MODE: AtomicU8 = AtomicU8::new(SyncMode::Sync as u8);
    1424              : 
    1425              : #[cfg(test)]
    1426              : mod tests {
    1427              :     use crate::context::DownloadBehavior;
    1428              :     use crate::task_mgr::TaskKind;
    1429              : 
    1430              :     use super::*;
    1431              :     use owned_buffers_io::io_buf_ext::IoBufExt;
    1432              :     use owned_buffers_io::slice::SliceMutExt;
    1433              :     use rand::seq::SliceRandom;
    1434              :     use rand::thread_rng;
    1435              :     use rand::Rng;
    1436              :     use std::io::Write;
    1437              :     use std::os::unix::fs::FileExt;
    1438              :     use std::sync::Arc;
    1439              : 
    1440              :     enum MaybeVirtualFile {
    1441              :         VirtualFile(VirtualFile),
    1442              :         File(File),
    1443              :     }
    1444              : 
    1445              :     impl From<VirtualFile> for MaybeVirtualFile {
    1446           12 :         fn from(vf: VirtualFile) -> Self {
    1447           12 :             MaybeVirtualFile::VirtualFile(vf)
    1448           12 :         }
    1449              :     }
    1450              : 
    1451              :     impl MaybeVirtualFile {
    1452          808 :         async fn read_exact_at(
    1453          808 :             &self,
    1454          808 :             mut slice: tokio_epoll_uring::Slice<IoBufferMut>,
    1455          808 :             offset: u64,
    1456          808 :             ctx: &RequestContext,
    1457          808 :         ) -> Result<tokio_epoll_uring::Slice<IoBufferMut>, Error> {
    1458          808 :             match self {
    1459          404 :                 MaybeVirtualFile::VirtualFile(file) => file.read_exact_at(slice, offset, ctx).await,
    1460          404 :                 MaybeVirtualFile::File(file) => {
    1461          404 :                     let rust_slice: &mut [u8] = slice.as_mut_rust_slice_full_zeroed();
    1462          404 :                     file.read_exact_at(rust_slice, offset).map(|()| slice)
    1463              :                 }
    1464              :             }
    1465          808 :         }
    1466           16 :         async fn write_all_at<Buf: IoBufAligned + Send>(
    1467           16 :             &self,
    1468           16 :             buf: FullSlice<Buf>,
    1469           16 :             offset: u64,
    1470           16 :             ctx: &RequestContext,
    1471           16 :         ) -> Result<(), Error> {
    1472           16 :             match self {
    1473            8 :                 MaybeVirtualFile::VirtualFile(file) => {
    1474            8 :                     let (_buf, res) = file.write_all_at(buf, offset, ctx).await;
    1475            8 :                     res
    1476              :                 }
    1477            8 :                 MaybeVirtualFile::File(file) => file.write_all_at(&buf[..], offset),
    1478              :             }
    1479           16 :         }
    1480           72 :         async fn seek(&mut self, pos: SeekFrom) -> Result<u64, Error> {
    1481           72 :             match self {
    1482           36 :                 MaybeVirtualFile::VirtualFile(file) => file.seek(pos).await,
    1483           36 :                 MaybeVirtualFile::File(file) => file.seek(pos),
    1484              :             }
    1485           72 :         }
    1486           16 :         async fn write_all<Buf: IoBuf + Send>(
    1487           16 :             &mut self,
    1488           16 :             buf: FullSlice<Buf>,
    1489           16 :             ctx: &RequestContext,
    1490           16 :         ) -> Result<(), Error> {
    1491           16 :             match self {
    1492            8 :                 MaybeVirtualFile::VirtualFile(file) => {
    1493            8 :                     let (_buf, res) = file.write_all(buf, ctx).await;
    1494            8 :                     res.map(|_| ())
    1495              :                 }
    1496            8 :                 MaybeVirtualFile::File(file) => file.write_all(&buf[..]),
    1497              :             }
    1498           16 :         }
    1499              : 
    1500              :         // Helper function to slurp contents of a file, starting at the current position,
    1501              :         // into a string
    1502          884 :         async fn read_string(&mut self, ctx: &RequestContext) -> Result<String, Error> {
    1503              :             use std::io::Read;
    1504          884 :             let mut buf = String::new();
    1505          884 :             match self {
    1506          448 :                 MaybeVirtualFile::VirtualFile(file) => {
    1507          448 :                     let mut buf = Vec::new();
    1508          448 :                     file.read_to_end(&mut buf, ctx).await?;
    1509          444 :                     return Ok(String::from_utf8(buf).unwrap());
    1510              :                 }
    1511          436 :                 MaybeVirtualFile::File(file) => {
    1512          436 :                     file.read_to_string(&mut buf)?;
    1513              :                 }
    1514              :             }
    1515          432 :             Ok(buf)
    1516          884 :         }
    1517              : 
    1518              :         // Helper function to slurp a portion of a file into a string
    1519          808 :         async fn read_string_at(
    1520          808 :             &mut self,
    1521          808 :             pos: u64,
    1522          808 :             len: usize,
    1523          808 :             ctx: &RequestContext,
    1524          808 :         ) -> Result<String, Error> {
    1525          808 :             let slice = IoBufferMut::with_capacity(len).slice_full();
    1526          808 :             assert_eq!(slice.bytes_total(), len);
    1527          808 :             let slice = self.read_exact_at(slice, pos, ctx).await?;
    1528          808 :             let buf = slice.into_inner();
    1529          808 :             assert_eq!(buf.len(), len);
    1530              : 
    1531          808 :             Ok(String::from_utf8(buf.to_vec()).unwrap())
    1532          808 :         }
    1533              :     }
    1534              : 
    1535              :     #[tokio::test]
    1536            4 :     async fn test_virtual_files() -> anyhow::Result<()> {
    1537            4 :         // The real work is done in the test_files() helper function. This
    1538            4 :         // allows us to run the same set of tests against a native File, and
    1539            4 :         // VirtualFile. We trust the native Files and wouldn't need to test them,
    1540            4 :         // but this allows us to verify that the operations return the same
    1541            4 :         // results with VirtualFiles as with native Files. (Except that with
    1542            4 :         // native files, you will run out of file descriptors if the ulimit
    1543            4 :         // is low enough.)
    1544            4 :         struct A;
    1545            4 : 
    1546            4 :         impl Adapter for A {
    1547          412 :             async fn open(
    1548          412 :                 path: Utf8PathBuf,
    1549          412 :                 opts: OpenOptions,
    1550          412 :                 ctx: &RequestContext,
    1551          412 :             ) -> Result<MaybeVirtualFile, anyhow::Error> {
    1552          412 :                 let vf = VirtualFile::open_with_options(&path, &opts, ctx).await?;
    1553          412 :                 Ok(MaybeVirtualFile::VirtualFile(vf))
    1554          412 :             }
    1555            4 :         }
    1556            4 :         test_files::<A>("virtual_files").await
    1557            4 :     }
    1558              : 
    1559              :     #[tokio::test]
    1560            4 :     async fn test_physical_files() -> anyhow::Result<()> {
    1561            4 :         struct B;
    1562            4 : 
    1563            4 :         impl Adapter for B {
    1564          412 :             async fn open(
    1565          412 :                 path: Utf8PathBuf,
    1566          412 :                 opts: OpenOptions,
    1567          412 :                 _ctx: &RequestContext,
    1568          412 :             ) -> Result<MaybeVirtualFile, anyhow::Error> {
    1569            4 :                 Ok(MaybeVirtualFile::File({
    1570          412 :                     let owned_fd = opts.open(path.as_std_path()).await?;
    1571          412 :                     File::from(owned_fd)
    1572            4 :                 }))
    1573          412 :             }
    1574            4 :         }
    1575            4 : 
    1576            4 :         test_files::<B>("physical_files").await
    1577            4 :     }
    1578              : 
    1579              :     /// This is essentially a closure which returns a MaybeVirtualFile, but because rust edition
    1580              :     /// 2024 is not yet out with new lifetime capture or outlives rules, this is a async function
    1581              :     /// in trait which benefits from the new lifetime capture rules already.
    1582              :     trait Adapter {
    1583              :         async fn open(
    1584              :             path: Utf8PathBuf,
    1585              :             opts: OpenOptions,
    1586              :             ctx: &RequestContext,
    1587              :         ) -> Result<MaybeVirtualFile, anyhow::Error>;
    1588              :     }
    1589              : 
    1590            8 :     async fn test_files<A>(testname: &str) -> anyhow::Result<()>
    1591            8 :     where
    1592            8 :         A: Adapter,
    1593            8 :     {
    1594            8 :         let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
    1595            8 :         let testdir = crate::config::PageServerConf::test_repo_dir(testname);
    1596            8 :         std::fs::create_dir_all(&testdir)?;
    1597              : 
    1598            8 :         let path_a = testdir.join("file_a");
    1599            8 :         let mut file_a = A::open(
    1600            8 :             path_a.clone(),
    1601            8 :             OpenOptions::new()
    1602            8 :                 .write(true)
    1603            8 :                 .create(true)
    1604            8 :                 .truncate(true)
    1605            8 :                 .to_owned(),
    1606            8 :             &ctx,
    1607            8 :         )
    1608            8 :         .await?;
    1609              : 
    1610            8 :         file_a
    1611            8 :             .write_all(b"foobar".to_vec().slice_len(), &ctx)
    1612            8 :             .await?;
    1613              : 
    1614              :         // cannot read from a file opened in write-only mode
    1615            8 :         let _ = file_a.read_string(&ctx).await.unwrap_err();
    1616              : 
    1617              :         // Close the file and re-open for reading
    1618            8 :         let mut file_a = A::open(path_a, OpenOptions::new().read(true).to_owned(), &ctx).await?;
    1619              : 
    1620              :         // cannot write to a file opened in read-only mode
    1621            8 :         let _ = file_a
    1622            8 :             .write_all(b"bar".to_vec().slice_len(), &ctx)
    1623            8 :             .await
    1624            8 :             .unwrap_err();
    1625            8 : 
    1626            8 :         // Try simple read
    1627            8 :         assert_eq!("foobar", file_a.read_string(&ctx).await?);
    1628              : 
    1629              :         // It's positioned at the EOF now.
    1630            8 :         assert_eq!("", file_a.read_string(&ctx).await?);
    1631              : 
    1632              :         // Test seeks.
    1633            8 :         assert_eq!(file_a.seek(SeekFrom::Start(1)).await?, 1);
    1634            8 :         assert_eq!("oobar", file_a.read_string(&ctx).await?);
    1635              : 
    1636            8 :         assert_eq!(file_a.seek(SeekFrom::End(-2)).await?, 4);
    1637            8 :         assert_eq!("ar", file_a.read_string(&ctx).await?);
    1638              : 
    1639            8 :         assert_eq!(file_a.seek(SeekFrom::Start(1)).await?, 1);
    1640            8 :         assert_eq!(file_a.seek(SeekFrom::Current(2)).await?, 3);
    1641            8 :         assert_eq!("bar", file_a.read_string(&ctx).await?);
    1642              : 
    1643            8 :         assert_eq!(file_a.seek(SeekFrom::Current(-5)).await?, 1);
    1644            8 :         assert_eq!("oobar", file_a.read_string(&ctx).await?);
    1645              : 
    1646              :         // Test erroneous seeks to before byte 0
    1647            8 :         file_a.seek(SeekFrom::End(-7)).await.unwrap_err();
    1648            8 :         assert_eq!(file_a.seek(SeekFrom::Start(1)).await?, 1);
    1649            8 :         file_a.seek(SeekFrom::Current(-2)).await.unwrap_err();
    1650            8 : 
    1651            8 :         // the erroneous seek should have left the position unchanged
    1652            8 :         assert_eq!("oobar", file_a.read_string(&ctx).await?);
    1653              : 
    1654              :         // Create another test file, and try FileExt functions on it.
    1655            8 :         let path_b = testdir.join("file_b");
    1656            8 :         let mut file_b = A::open(
    1657            8 :             path_b.clone(),
    1658            8 :             OpenOptions::new()
    1659            8 :                 .read(true)
    1660            8 :                 .write(true)
    1661            8 :                 .create(true)
    1662            8 :                 .truncate(true)
    1663            8 :                 .to_owned(),
    1664            8 :             &ctx,
    1665            8 :         )
    1666            8 :         .await?;
    1667            8 :         file_b
    1668            8 :             .write_all_at(IoBuffer::from(b"BAR").slice_len(), 3, &ctx)
    1669            8 :             .await?;
    1670            8 :         file_b
    1671            8 :             .write_all_at(IoBuffer::from(b"FOO").slice_len(), 0, &ctx)
    1672            8 :             .await?;
    1673              : 
    1674            8 :         assert_eq!(file_b.read_string_at(2, 3, &ctx).await?, "OBA");
    1675              : 
    1676              :         // Open a lot of files, enough to cause some evictions. (Or to be precise,
    1677              :         // open the same file many times. The effect is the same.)
    1678              :         //
    1679              :         // leave file_a positioned at offset 1 before we start
    1680            8 :         assert_eq!(file_a.seek(SeekFrom::Start(1)).await?, 1);
    1681              : 
    1682            8 :         let mut vfiles = Vec::new();
    1683          808 :         for _ in 0..100 {
    1684          800 :             let mut vfile = A::open(
    1685          800 :                 path_b.clone(),
    1686          800 :                 OpenOptions::new().read(true).to_owned(),
    1687          800 :                 &ctx,
    1688          800 :             )
    1689          800 :             .await?;
    1690          800 :             assert_eq!("FOOBAR", vfile.read_string(&ctx).await?);
    1691          800 :             vfiles.push(vfile);
    1692              :         }
    1693              : 
    1694              :         // make sure we opened enough files to definitely cause evictions.
    1695            8 :         assert!(vfiles.len() > TEST_MAX_FILE_DESCRIPTORS * 2);
    1696              : 
    1697              :         // The underlying file descriptor for 'file_a' should be closed now. Try to read
    1698              :         // from it again. We left the file positioned at offset 1 above.
    1699            8 :         assert_eq!("oobar", file_a.read_string(&ctx).await?);
    1700              : 
    1701              :         // Check that all the other FDs still work too. Use them in random order for
    1702              :         // good measure.
    1703            8 :         vfiles.as_mut_slice().shuffle(&mut thread_rng());
    1704          800 :         for vfile in vfiles.iter_mut() {
    1705          800 :             assert_eq!("OOBAR", vfile.read_string_at(1, 5, &ctx).await?);
    1706              :         }
    1707              : 
    1708            8 :         Ok(())
    1709            8 :     }
    1710              : 
    1711              :     /// Test using VirtualFiles from many threads concurrently. This tests both using
    1712              :     /// a lot of VirtualFiles concurrently, causing evictions, and also using the same
    1713              :     /// VirtualFile from multiple threads concurrently.
    1714              :     #[tokio::test]
    1715            4 :     async fn test_vfile_concurrency() -> Result<(), Error> {
    1716            4 :         const SIZE: usize = 8 * 1024;
    1717            4 :         const VIRTUAL_FILES: usize = 100;
    1718            4 :         const THREADS: usize = 100;
    1719            4 :         const SAMPLE: [u8; SIZE] = [0xADu8; SIZE];
    1720            4 : 
    1721            4 :         let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
    1722            4 :         let testdir = crate::config::PageServerConf::test_repo_dir("vfile_concurrency");
    1723            4 :         std::fs::create_dir_all(&testdir)?;
    1724            4 : 
    1725            4 :         // Create a test file.
    1726            4 :         let test_file_path = testdir.join("concurrency_test_file");
    1727            4 :         {
    1728            4 :             let file = File::create(&test_file_path)?;
    1729            4 :             file.write_all_at(&SAMPLE, 0)?;
    1730            4 :         }
    1731            4 : 
    1732            4 :         // Open the file many times.
    1733            4 :         let mut files = Vec::new();
    1734          404 :         for _ in 0..VIRTUAL_FILES {
    1735          400 :             let f = VirtualFileInner::open_with_options(
    1736          400 :                 &test_file_path,
    1737          400 :                 OpenOptions::new().read(true),
    1738          400 :                 &ctx,
    1739          400 :             )
    1740          400 :             .await?;
    1741          400 :             files.push(f);
    1742            4 :         }
    1743            4 :         let files = Arc::new(files);
    1744            4 : 
    1745            4 :         // Launch many threads, and use the virtual files concurrently in random order.
    1746            4 :         let rt = tokio::runtime::Builder::new_multi_thread()
    1747            4 :             .worker_threads(THREADS)
    1748            4 :             .thread_name("test_vfile_concurrency thread")
    1749            4 :             .build()
    1750            4 :             .unwrap();
    1751            4 :         let mut hdls = Vec::new();
    1752          404 :         for _threadno in 0..THREADS {
    1753          400 :             let files = files.clone();
    1754          400 :             let ctx = ctx.detached_child(TaskKind::UnitTest, DownloadBehavior::Error);
    1755          400 :             let hdl = rt.spawn(async move {
    1756          400 :                 let mut buf = IoBufferMut::with_capacity_zeroed(SIZE);
    1757          400 :                 let mut rng = rand::rngs::OsRng;
    1758       400000 :                 for _ in 1..1000 {
    1759       399600 :                     let f = &files[rng.gen_range(0..files.len())];
    1760       399600 :                     buf = f
    1761       399600 :                         .read_exact_at(buf.slice_full(), 0, &ctx)
    1762       399600 :                         .await
    1763       399600 :                         .unwrap()
    1764       399600 :                         .into_inner();
    1765       399600 :                     assert!(buf[..] == SAMPLE);
    1766            4 :                 }
    1767          400 :             });
    1768          400 :             hdls.push(hdl);
    1769          400 :         }
    1770          404 :         for hdl in hdls {
    1771          400 :             hdl.await?;
    1772            4 :         }
    1773            4 :         std::mem::forget(rt);
    1774            4 : 
    1775            4 :         Ok(())
    1776            4 :     }
    1777              : 
    1778              :     #[tokio::test]
    1779            4 :     async fn test_atomic_overwrite_basic() {
    1780            4 :         let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
    1781            4 :         let testdir = crate::config::PageServerConf::test_repo_dir("test_atomic_overwrite_basic");
    1782            4 :         std::fs::create_dir_all(&testdir).unwrap();
    1783            4 : 
    1784            4 :         let path = testdir.join("myfile");
    1785            4 :         let tmp_path = testdir.join("myfile.tmp");
    1786            4 : 
    1787            4 :         VirtualFileInner::crashsafe_overwrite(path.clone(), tmp_path.clone(), b"foo".to_vec())
    1788            4 :             .await
    1789            4 :             .unwrap();
    1790            4 :         let mut file = MaybeVirtualFile::from(VirtualFile::open(&path, &ctx).await.unwrap());
    1791            4 :         let post = file.read_string(&ctx).await.unwrap();
    1792            4 :         assert_eq!(post, "foo");
    1793            4 :         assert!(!tmp_path.exists());
    1794            4 :         drop(file);
    1795            4 : 
    1796            4 :         VirtualFileInner::crashsafe_overwrite(path.clone(), tmp_path.clone(), b"bar".to_vec())
    1797            4 :             .await
    1798            4 :             .unwrap();
    1799            4 :         let mut file = MaybeVirtualFile::from(VirtualFile::open(&path, &ctx).await.unwrap());
    1800            4 :         let post = file.read_string(&ctx).await.unwrap();
    1801            4 :         assert_eq!(post, "bar");
    1802            4 :         assert!(!tmp_path.exists());
    1803            4 :         drop(file);
    1804            4 :     }
    1805              : 
    1806              :     #[tokio::test]
    1807            4 :     async fn test_atomic_overwrite_preexisting_tmp() {
    1808            4 :         let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
    1809            4 :         let testdir =
    1810            4 :             crate::config::PageServerConf::test_repo_dir("test_atomic_overwrite_preexisting_tmp");
    1811            4 :         std::fs::create_dir_all(&testdir).unwrap();
    1812            4 : 
    1813            4 :         let path = testdir.join("myfile");
    1814            4 :         let tmp_path = testdir.join("myfile.tmp");
    1815            4 : 
    1816            4 :         std::fs::write(&tmp_path, "some preexisting junk that should be removed").unwrap();
    1817            4 :         assert!(tmp_path.exists());
    1818            4 : 
    1819            4 :         VirtualFileInner::crashsafe_overwrite(path.clone(), tmp_path.clone(), b"foo".to_vec())
    1820            4 :             .await
    1821            4 :             .unwrap();
    1822            4 : 
    1823            4 :         let mut file = MaybeVirtualFile::from(VirtualFile::open(&path, &ctx).await.unwrap());
    1824            4 :         let post = file.read_string(&ctx).await.unwrap();
    1825            4 :         assert_eq!(post, "foo");
    1826            4 :         assert!(!tmp_path.exists());
    1827            4 :         drop(file);
    1828            4 :     }
    1829              : }
        

Generated by: LCOV version 2.1-beta