LCOV - code coverage report
Current view: top level - pageserver/src/tenant/storage_layer - image_layer.rs (source / functions) Coverage Total Hit
Test: 6df3fc19ec669bcfbbf9aba41d1338898d24eaa0.info Lines: 78.7 % 966 760
Test Date: 2025-03-12 18:28:53 Functions: 57.5 % 87 50

            Line data    Source code
       1              : //! An ImageLayer represents an image or a snapshot of a key-range at
       2              : //! one particular LSN.
       3              : //!
       4              : //! It contains an image of all key-value pairs in its key-range. Any key
       5              : //! that falls into the image layer's range but does not exist in the layer,
       6              : //! does not exist.
       7              : //!
       8              : //! An image layer is stored in a file on disk. The file is stored in
       9              : //! timelines/<timeline_id> directory.  Currently, there are no
      10              : //! subdirectories, and each image layer file is named like this:
      11              : //!
      12              : //! ```text
      13              : //!    <key start>-<key end>__<LSN>
      14              : //! ```
      15              : //!
      16              : //! For example:
      17              : //!
      18              : //! ```text
      19              : //!    000000067F000032BE0000400000000070B6-000000067F000032BE0000400000000080B6__00000000346BC568
      20              : //! ```
      21              : //!
      22              : //! Every image layer file consists of three parts: "summary",
      23              : //! "index", and "values".  The summary is a fixed size header at the
      24              : //! beginning of the file, and it contains basic information about the
      25              : //! layer, and offsets to the other parts. The "index" is a B-tree,
      26              : //! mapping from Key to an offset in the "values" part.  The
      27              : //! actual page images are stored in the "values" part.
      28              : use std::collections::{HashMap, VecDeque};
      29              : use std::fs::File;
      30              : use std::io::SeekFrom;
      31              : use std::ops::Range;
      32              : use std::os::unix::prelude::FileExt;
      33              : use std::str::FromStr;
      34              : use std::sync::Arc;
      35              : 
      36              : use anyhow::{Context, Result, bail, ensure};
      37              : use bytes::Bytes;
      38              : use camino::{Utf8Path, Utf8PathBuf};
      39              : use hex;
      40              : use itertools::Itertools;
      41              : use pageserver_api::config::MaxVectoredReadBytes;
      42              : use pageserver_api::key::{DBDIR_KEY, KEY_SIZE, Key};
      43              : use pageserver_api::keyspace::KeySpace;
      44              : use pageserver_api::shard::{ShardIdentity, TenantShardId};
      45              : use pageserver_api::value::Value;
      46              : use rand::Rng;
      47              : use rand::distributions::Alphanumeric;
      48              : use serde::{Deserialize, Serialize};
      49              : use tokio::sync::OnceCell;
      50              : use tokio_stream::StreamExt;
      51              : use tracing::*;
      52              : use utils::bin_ser::BeSer;
      53              : use utils::id::{TenantId, TimelineId};
      54              : use utils::lsn::Lsn;
      55              : 
      56              : use super::layer_name::ImageLayerName;
      57              : use super::{
      58              :     AsLayerDesc, LayerName, OnDiskValue, OnDiskValueIo, PersistentLayerDesc, ResidentLayer,
      59              :     ValuesReconstructState,
      60              : };
      61              : use crate::config::PageServerConf;
      62              : use crate::context::{PageContentKind, RequestContext, RequestContextBuilder};
      63              : use crate::page_cache::{self, FileId, PAGE_SZ};
      64              : use crate::tenant::blob_io::BlobWriter;
      65              : use crate::tenant::block_io::{BlockBuf, FileBlockReader};
      66              : use crate::tenant::disk_btree::{
      67              :     DiskBtreeBuilder, DiskBtreeIterator, DiskBtreeReader, VisitDirection,
      68              : };
      69              : use crate::tenant::timeline::GetVectoredError;
      70              : use crate::tenant::vectored_blob_io::{
      71              :     BlobFlag, BufView, StreamingVectoredReadPlanner, VectoredBlobReader, VectoredRead,
      72              :     VectoredReadPlanner,
      73              : };
      74              : use crate::virtual_file::owned_buffers_io::io_buf_ext::IoBufExt;
      75              : use crate::virtual_file::{self, IoBufferMut, MaybeFatalIo, VirtualFile};
      76              : use crate::{IMAGE_FILE_MAGIC, STORAGE_FORMAT_VERSION, TEMP_FILE_SUFFIX};
      77              : 
      78              : ///
      79              : /// Header stored in the beginning of the file
      80              : ///
      81              : /// After this comes the 'values' part, starting on block 1. After that,
      82              : /// the 'index' starts at the block indicated by 'index_start_blk'
      83              : ///
      84            0 : #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
      85              : pub struct Summary {
      86              :     /// Magic value to identify this as a neon image file. Always IMAGE_FILE_MAGIC.
      87              :     pub magic: u16,
      88              :     pub format_version: u16,
      89              : 
      90              :     pub tenant_id: TenantId,
      91              :     pub timeline_id: TimelineId,
      92              :     pub key_range: Range<Key>,
      93              :     pub lsn: Lsn,
      94              : 
      95              :     /// Block number where the 'index' part of the file begins.
      96              :     pub index_start_blk: u32,
      97              :     /// Block within the 'index', where the B-tree root page is stored
      98              :     pub index_root_blk: u32,
      99              :     // the 'values' part starts after the summary header, on block 1.
     100              : }
     101              : 
     102              : impl From<&ImageLayer> for Summary {
     103            0 :     fn from(layer: &ImageLayer) -> Self {
     104            0 :         Self::expected(
     105            0 :             layer.desc.tenant_shard_id.tenant_id,
     106            0 :             layer.desc.timeline_id,
     107            0 :             layer.desc.key_range.clone(),
     108            0 :             layer.lsn,
     109            0 :         )
     110            0 :     }
     111              : }
     112              : 
     113              : impl Summary {
     114          276 :     pub(super) fn expected(
     115          276 :         tenant_id: TenantId,
     116          276 :         timeline_id: TimelineId,
     117          276 :         key_range: Range<Key>,
     118          276 :         lsn: Lsn,
     119          276 :     ) -> Self {
     120          276 :         Self {
     121          276 :             magic: IMAGE_FILE_MAGIC,
     122          276 :             format_version: STORAGE_FORMAT_VERSION,
     123          276 :             tenant_id,
     124          276 :             timeline_id,
     125          276 :             key_range,
     126          276 :             lsn,
     127          276 : 
     128          276 :             index_start_blk: 0,
     129          276 :             index_root_blk: 0,
     130          276 :         }
     131          276 :     }
     132              : }
     133              : 
     134              : /// This is used only from `pagectl`. Within pageserver, all layers are
     135              : /// [`crate::tenant::storage_layer::Layer`], which can hold an [`ImageLayerInner`].
     136              : pub struct ImageLayer {
     137              :     path: Utf8PathBuf,
     138              :     pub desc: PersistentLayerDesc,
     139              :     // This entry contains an image of all pages as of this LSN, should be the same as desc.lsn
     140              :     pub lsn: Lsn,
     141              :     inner: OnceCell<ImageLayerInner>,
     142              : }
     143              : 
     144              : impl std::fmt::Debug for ImageLayer {
     145            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     146              :         use super::RangeDisplayDebug;
     147              : 
     148            0 :         f.debug_struct("ImageLayer")
     149            0 :             .field("key_range", &RangeDisplayDebug(&self.desc.key_range))
     150            0 :             .field("file_size", &self.desc.file_size)
     151            0 :             .field("lsn", &self.lsn)
     152            0 :             .field("inner", &self.inner)
     153            0 :             .finish()
     154            0 :     }
     155              : }
     156              : 
     157              : /// ImageLayer is the in-memory data structure associated with an on-disk image
     158              : /// file.
     159              : pub struct ImageLayerInner {
     160              :     // values copied from summary
     161              :     index_start_blk: u32,
     162              :     index_root_blk: u32,
     163              : 
     164              :     key_range: Range<Key>,
     165              :     lsn: Lsn,
     166              : 
     167              :     file: Arc<VirtualFile>,
     168              :     file_id: FileId,
     169              : 
     170              :     max_vectored_read_bytes: Option<MaxVectoredReadBytes>,
     171              : }
     172              : 
     173              : impl ImageLayerInner {
     174            0 :     pub(crate) fn layer_dbg_info(&self) -> String {
     175            0 :         format!(
     176            0 :             "image {}..{} {}",
     177            0 :             self.key_range().start,
     178            0 :             self.key_range().end,
     179            0 :             self.lsn()
     180            0 :         )
     181            0 :     }
     182              : }
     183              : 
     184              : impl std::fmt::Debug for ImageLayerInner {
     185            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     186            0 :         f.debug_struct("ImageLayerInner")
     187            0 :             .field("index_start_blk", &self.index_start_blk)
     188            0 :             .field("index_root_blk", &self.index_root_blk)
     189            0 :             .finish()
     190            0 :     }
     191              : }
     192              : 
     193              : impl ImageLayerInner {
     194            0 :     pub(super) async fn dump(&self, ctx: &RequestContext) -> anyhow::Result<()> {
     195            0 :         let block_reader = FileBlockReader::new(&self.file, self.file_id);
     196            0 :         let tree_reader = DiskBtreeReader::<_, KEY_SIZE>::new(
     197            0 :             self.index_start_blk,
     198            0 :             self.index_root_blk,
     199            0 :             block_reader,
     200            0 :         );
     201            0 : 
     202            0 :         tree_reader.dump(ctx).await?;
     203              : 
     204            0 :         tree_reader
     205            0 :             .visit(
     206            0 :                 &[0u8; KEY_SIZE],
     207            0 :                 VisitDirection::Forwards,
     208            0 :                 |key, value| {
     209            0 :                     println!("key: {} offset {}", hex::encode(key), value);
     210            0 :                     true
     211            0 :                 },
     212            0 :                 ctx,
     213            0 :             )
     214            0 :             .await?;
     215              : 
     216            0 :         Ok(())
     217            0 :     }
     218              : }
     219              : 
     220              : /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
     221              : impl std::fmt::Display for ImageLayer {
     222            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     223            0 :         write!(f, "{}", self.layer_desc().short_id())
     224            0 :     }
     225              : }
     226              : 
     227              : impl AsLayerDesc for ImageLayer {
     228            0 :     fn layer_desc(&self) -> &PersistentLayerDesc {
     229            0 :         &self.desc
     230            0 :     }
     231              : }
     232              : 
     233              : impl ImageLayer {
     234            0 :     pub async fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
     235            0 :         self.desc.dump();
     236            0 : 
     237            0 :         if !verbose {
     238            0 :             return Ok(());
     239            0 :         }
     240              : 
     241            0 :         let inner = self.load(ctx).await?;
     242              : 
     243            0 :         inner.dump(ctx).await?;
     244              : 
     245            0 :         Ok(())
     246            0 :     }
     247              : 
     248         1204 :     fn temp_path_for(
     249         1204 :         conf: &PageServerConf,
     250         1204 :         timeline_id: TimelineId,
     251         1204 :         tenant_shard_id: TenantShardId,
     252         1204 :         fname: &ImageLayerName,
     253         1204 :     ) -> Utf8PathBuf {
     254         1204 :         let rand_string: String = rand::thread_rng()
     255         1204 :             .sample_iter(&Alphanumeric)
     256         1204 :             .take(8)
     257         1204 :             .map(char::from)
     258         1204 :             .collect();
     259         1204 : 
     260         1204 :         conf.timeline_path(&tenant_shard_id, &timeline_id)
     261         1204 :             .join(format!("{fname}.{rand_string}.{TEMP_FILE_SUFFIX}"))
     262         1204 :     }
     263              : 
     264              :     ///
     265              :     /// Open the underlying file and read the metadata into memory, if it's
     266              :     /// not loaded already.
     267              :     ///
     268            0 :     async fn load(&self, ctx: &RequestContext) -> Result<&ImageLayerInner> {
     269            0 :         self.inner
     270            0 :             .get_or_try_init(|| self.load_inner(ctx))
     271            0 :             .await
     272            0 :             .with_context(|| format!("Failed to load image layer {}", self.path()))
     273            0 :     }
     274              : 
     275            0 :     async fn load_inner(&self, ctx: &RequestContext) -> Result<ImageLayerInner> {
     276            0 :         let path = self.path();
     277              : 
     278            0 :         let loaded =
     279            0 :             ImageLayerInner::load(&path, self.desc.image_layer_lsn(), None, None, ctx).await?;
     280              : 
     281              :         // not production code
     282            0 :         let actual_layer_name = LayerName::from_str(path.file_name().unwrap()).unwrap();
     283            0 :         let expected_layer_name = self.layer_desc().layer_name();
     284            0 : 
     285            0 :         if actual_layer_name != expected_layer_name {
     286            0 :             println!("warning: filename does not match what is expected from in-file summary");
     287            0 :             println!("actual: {:?}", actual_layer_name.to_string());
     288            0 :             println!("expected: {:?}", expected_layer_name.to_string());
     289            0 :         }
     290              : 
     291            0 :         Ok(loaded)
     292            0 :     }
     293              : 
     294              :     /// Create an ImageLayer struct representing an existing file on disk.
     295              :     ///
     296              :     /// This variant is only used for debugging purposes, by the 'pagectl' binary.
     297            0 :     pub fn new_for_path(path: &Utf8Path, file: File) -> Result<ImageLayer> {
     298            0 :         let mut summary_buf = vec![0; PAGE_SZ];
     299            0 :         file.read_exact_at(&mut summary_buf, 0)?;
     300            0 :         let summary = Summary::des_prefix(&summary_buf)?;
     301            0 :         let metadata = file
     302            0 :             .metadata()
     303            0 :             .context("get file metadata to determine size")?;
     304              : 
     305              :         // This function is never used for constructing layers in a running pageserver,
     306              :         // so it does not need an accurate TenantShardId.
     307            0 :         let tenant_shard_id = TenantShardId::unsharded(summary.tenant_id);
     308            0 : 
     309            0 :         Ok(ImageLayer {
     310            0 :             path: path.to_path_buf(),
     311            0 :             desc: PersistentLayerDesc::new_img(
     312            0 :                 tenant_shard_id,
     313            0 :                 summary.timeline_id,
     314            0 :                 summary.key_range,
     315            0 :                 summary.lsn,
     316            0 :                 metadata.len(),
     317            0 :             ), // Now we assume image layer ALWAYS covers the full range. This may change in the future.
     318            0 :             lsn: summary.lsn,
     319            0 :             inner: OnceCell::new(),
     320            0 :         })
     321            0 :     }
     322              : 
     323            0 :     fn path(&self) -> Utf8PathBuf {
     324            0 :         self.path.clone()
     325            0 :     }
     326              : }
     327              : 
     328              : #[derive(thiserror::Error, Debug)]
     329              : pub enum RewriteSummaryError {
     330              :     #[error("magic mismatch")]
     331              :     MagicMismatch,
     332              :     #[error(transparent)]
     333              :     Other(#[from] anyhow::Error),
     334              : }
     335              : 
     336              : impl From<std::io::Error> for RewriteSummaryError {
     337            0 :     fn from(e: std::io::Error) -> Self {
     338            0 :         Self::Other(anyhow::anyhow!(e))
     339            0 :     }
     340              : }
     341              : 
     342              : impl ImageLayer {
     343            0 :     pub async fn rewrite_summary<F>(
     344            0 :         path: &Utf8Path,
     345            0 :         rewrite: F,
     346            0 :         ctx: &RequestContext,
     347            0 :     ) -> Result<(), RewriteSummaryError>
     348            0 :     where
     349            0 :         F: Fn(Summary) -> Summary,
     350            0 :     {
     351            0 :         let mut file = VirtualFile::open_with_options(
     352            0 :             path,
     353            0 :             virtual_file::OpenOptions::new().read(true).write(true),
     354            0 :             ctx,
     355            0 :         )
     356            0 :         .await
     357            0 :         .with_context(|| format!("Failed to open file '{}'", path))?;
     358            0 :         let file_id = page_cache::next_file_id();
     359            0 :         let block_reader = FileBlockReader::new(&file, file_id);
     360            0 :         let summary_blk = block_reader.read_blk(0, ctx).await?;
     361            0 :         let actual_summary = Summary::des_prefix(summary_blk.as_ref()).context("deserialize")?;
     362            0 :         if actual_summary.magic != IMAGE_FILE_MAGIC {
     363            0 :             return Err(RewriteSummaryError::MagicMismatch);
     364            0 :         }
     365            0 : 
     366            0 :         let new_summary = rewrite(actual_summary);
     367            0 : 
     368            0 :         let mut buf = Vec::with_capacity(PAGE_SZ);
     369            0 :         // TODO: could use smallvec here but it's a pain with Slice<T>
     370            0 :         Summary::ser_into(&new_summary, &mut buf).context("serialize")?;
     371            0 :         file.seek(SeekFrom::Start(0)).await?;
     372            0 :         let (_buf, res) = file.write_all(buf.slice_len(), ctx).await;
     373            0 :         res?;
     374            0 :         Ok(())
     375            0 :     }
     376              : }
     377              : 
     378              : impl ImageLayerInner {
     379          264 :     pub(crate) fn key_range(&self) -> &Range<Key> {
     380          264 :         &self.key_range
     381          264 :     }
     382              : 
     383          264 :     pub(crate) fn lsn(&self) -> Lsn {
     384          264 :         self.lsn
     385          264 :     }
     386              : 
     387          276 :     pub(super) async fn load(
     388          276 :         path: &Utf8Path,
     389          276 :         lsn: Lsn,
     390          276 :         summary: Option<Summary>,
     391          276 :         max_vectored_read_bytes: Option<MaxVectoredReadBytes>,
     392          276 :         ctx: &RequestContext,
     393          276 :     ) -> anyhow::Result<Self> {
     394          276 :         let file = Arc::new(
     395          276 :             VirtualFile::open_v2(path, ctx)
     396          276 :                 .await
     397          276 :                 .context("open layer file")?,
     398              :         );
     399          276 :         let file_id = page_cache::next_file_id();
     400          276 :         let block_reader = FileBlockReader::new(&file, file_id);
     401          276 :         let summary_blk = block_reader
     402          276 :             .read_blk(0, ctx)
     403          276 :             .await
     404          276 :             .context("read first block")?;
     405              : 
     406              :         // length is the only way how this could fail, so it's not actually likely at all unless
     407              :         // read_blk returns wrong sized block.
     408              :         //
     409              :         // TODO: confirm and make this into assertion
     410          276 :         let actual_summary =
     411          276 :             Summary::des_prefix(summary_blk.as_ref()).context("deserialize first block")?;
     412              : 
     413          276 :         if let Some(mut expected_summary) = summary {
     414              :             // production code path
     415          276 :             expected_summary.index_start_blk = actual_summary.index_start_blk;
     416          276 :             expected_summary.index_root_blk = actual_summary.index_root_blk;
     417          276 :             // mask out the timeline_id, but still require the layers to be from the same tenant
     418          276 :             expected_summary.timeline_id = actual_summary.timeline_id;
     419          276 : 
     420          276 :             if actual_summary != expected_summary {
     421            0 :                 bail!(
     422            0 :                     "in-file summary does not match expected summary. actual = {:?} expected = {:?}",
     423            0 :                     actual_summary,
     424            0 :                     expected_summary
     425            0 :                 );
     426          276 :             }
     427            0 :         }
     428              : 
     429          276 :         Ok(ImageLayerInner {
     430          276 :             index_start_blk: actual_summary.index_start_blk,
     431          276 :             index_root_blk: actual_summary.index_root_blk,
     432          276 :             lsn,
     433          276 :             file,
     434          276 :             file_id,
     435          276 :             max_vectored_read_bytes,
     436          276 :             key_range: actual_summary.key_range,
     437          276 :         })
     438          276 :     }
     439              : 
     440              :     // Look up the keys in the provided keyspace and update
     441              :     // the reconstruct state with whatever is found.
     442        45124 :     pub(super) async fn get_values_reconstruct_data(
     443        45124 :         &self,
     444        45124 :         this: ResidentLayer,
     445        45124 :         keyspace: KeySpace,
     446        45124 :         reconstruct_state: &mut ValuesReconstructState,
     447        45124 :         ctx: &RequestContext,
     448        45124 :     ) -> Result<(), GetVectoredError> {
     449        45124 :         let reads = self
     450        45124 :             .plan_reads(keyspace, None, ctx)
     451        45124 :             .await
     452        45124 :             .map_err(GetVectoredError::Other)?;
     453              : 
     454        45124 :         self.do_reads_and_update_state(this, reads, reconstruct_state, ctx)
     455        45124 :             .await;
     456              : 
     457        45124 :         reconstruct_state.on_image_layer_visited(&self.key_range);
     458        45124 : 
     459        45124 :         Ok(())
     460        45124 :     }
     461              : 
     462              :     /// Traverse the layer's index to build read operations on the overlap of the input keyspace
     463              :     /// and the keys in this layer.
     464              :     ///
     465              :     /// If shard_identity is provided, it will be used to filter keys down to those stored on
     466              :     /// this shard.
     467        45140 :     async fn plan_reads(
     468        45140 :         &self,
     469        45140 :         keyspace: KeySpace,
     470        45140 :         shard_identity: Option<&ShardIdentity>,
     471        45140 :         ctx: &RequestContext,
     472        45140 :     ) -> anyhow::Result<Vec<VectoredRead>> {
     473        45140 :         let mut planner = VectoredReadPlanner::new(
     474        45140 :             self.max_vectored_read_bytes
     475        45140 :                 .expect("Layer is loaded with max vectored bytes config")
     476        45140 :                 .0
     477        45140 :                 .into(),
     478        45140 :         );
     479        45140 : 
     480        45140 :         let block_reader = FileBlockReader::new(&self.file, self.file_id);
     481        45140 :         let tree_reader =
     482        45140 :             DiskBtreeReader::new(self.index_start_blk, self.index_root_blk, block_reader);
     483        45140 : 
     484        45140 :         let ctx = RequestContextBuilder::extend(ctx)
     485        45140 :             .page_content_kind(PageContentKind::ImageLayerBtreeNode)
     486        45140 :             .build();
     487              : 
     488        45176 :         for range in keyspace.ranges.iter() {
     489        45176 :             let mut range_end_handled = false;
     490        45176 :             let mut search_key: [u8; KEY_SIZE] = [0u8; KEY_SIZE];
     491        45176 :             range.start.write_to_byte_slice(&mut search_key);
     492        45176 : 
     493        45176 :             let index_stream = tree_reader.clone().into_stream(&search_key, &ctx);
     494        45176 :             let mut index_stream = std::pin::pin!(index_stream);
     495              : 
     496      2251528 :             while let Some(index_entry) = index_stream.next().await {
     497      2250900 :                 let (raw_key, offset) = index_entry?;
     498              : 
     499      2250900 :                 let key = Key::from_slice(&raw_key[..KEY_SIZE]);
     500      2250900 :                 assert!(key >= range.start);
     501              : 
     502      2250900 :                 let flag = if let Some(shard_identity) = shard_identity {
     503      2097152 :                     if shard_identity.is_key_disposable(&key) {
     504      1572864 :                         BlobFlag::Ignore
     505              :                     } else {
     506       524288 :                         BlobFlag::None
     507              :                     }
     508              :                 } else {
     509       153748 :                     BlobFlag::None
     510              :                 };
     511              : 
     512      2250900 :                 if key >= range.end {
     513        44548 :                     planner.handle_range_end(offset);
     514        44548 :                     range_end_handled = true;
     515        44548 :                     break;
     516      2206352 :                 } else {
     517      2206352 :                     planner.handle(key, self.lsn, offset, flag);
     518      2206352 :                 }
     519              :             }
     520              : 
     521        45176 :             if !range_end_handled {
     522          628 :                 let payload_end = self.index_start_blk as u64 * PAGE_SZ as u64;
     523          628 :                 planner.handle_range_end(payload_end);
     524        44548 :             }
     525              :         }
     526              : 
     527        45140 :         Ok(planner.finish())
     528        45140 :     }
     529              : 
     530              :     /// Given a key range, select the parts of that range that should be retained by the ShardIdentity,
     531              :     /// then execute vectored GET operations, passing the results of all read keys into the writer.
     532           16 :     pub(super) async fn filter(
     533           16 :         &self,
     534           16 :         shard_identity: &ShardIdentity,
     535           16 :         writer: &mut ImageLayerWriter,
     536           16 :         ctx: &RequestContext,
     537           16 :     ) -> anyhow::Result<usize> {
     538              :         // Fragment the range into the regions owned by this ShardIdentity
     539           16 :         let plan = self
     540           16 :             .plan_reads(
     541           16 :                 KeySpace {
     542           16 :                     // If asked for the total key space, plan_reads will give us all the keys in the layer
     543           16 :                     ranges: vec![Key::MIN..Key::MAX],
     544           16 :                 },
     545           16 :                 Some(shard_identity),
     546           16 :                 ctx,
     547           16 :             )
     548           16 :             .await?;
     549              : 
     550           16 :         let vectored_blob_reader = VectoredBlobReader::new(&self.file);
     551           16 :         let mut key_count = 0;
     552           32 :         for read in plan.into_iter() {
     553           32 :             let buf_size = read.size();
     554           32 : 
     555           32 :             let buf = IoBufferMut::with_capacity(buf_size);
     556           32 :             let blobs_buf = vectored_blob_reader.read_blobs(&read, buf, ctx).await?;
     557              : 
     558           32 :             let view = BufView::new_slice(&blobs_buf.buf);
     559              : 
     560       524288 :             for meta in blobs_buf.blobs.iter() {
     561       524288 :                 let img_buf = meta.read(&view).await?;
     562              : 
     563       524288 :                 key_count += 1;
     564       524288 :                 writer
     565       524288 :                     .put_image(meta.meta.key, img_buf.into_bytes(), ctx)
     566       524288 :                     .await
     567       524288 :                     .context(format!("Storing key {}", meta.meta.key))?;
     568              :             }
     569              :         }
     570              : 
     571           16 :         Ok(key_count)
     572           16 :     }
     573              : 
     574        45124 :     async fn do_reads_and_update_state(
     575        45124 :         &self,
     576        45124 :         this: ResidentLayer,
     577        45124 :         reads: Vec<VectoredRead>,
     578        45124 :         reconstruct_state: &mut ValuesReconstructState,
     579        45124 :         ctx: &RequestContext,
     580        45124 :     ) {
     581        45124 :         let max_vectored_read_bytes = self
     582        45124 :             .max_vectored_read_bytes
     583        45124 :             .expect("Layer is loaded with max vectored bytes config")
     584        45124 :             .0
     585        45124 :             .into();
     586              : 
     587        45124 :         for read in reads.into_iter() {
     588        45048 :             let mut ios: HashMap<(Key, Lsn), OnDiskValueIo> = Default::default();
     589       109200 :             for (_, blob_meta) in read.blobs_at.as_slice() {
     590       109200 :                 let io = reconstruct_state.update_key(&blob_meta.key, blob_meta.lsn, true);
     591       109200 :                 ios.insert((blob_meta.key, blob_meta.lsn), io);
     592       109200 :             }
     593              : 
     594        45048 :             let buf_size = read.size();
     595        45048 : 
     596        45048 :             if buf_size > max_vectored_read_bytes {
     597              :                 // If the read is oversized, it should only contain one key.
     598            0 :                 let offenders = read
     599            0 :                     .blobs_at
     600            0 :                     .as_slice()
     601            0 :                     .iter()
     602            0 :                     .filter_map(|(_, blob_meta)| {
     603            0 :                         if blob_meta.key.is_rel_dir_key()
     604            0 :                             || blob_meta.key == DBDIR_KEY
     605            0 :                             || blob_meta.key.is_aux_file_key()
     606              :                         {
     607              :                             // The size of values for these keys is unbounded and can
     608              :                             // grow very large in pathological cases.
     609            0 :                             None
     610              :                         } else {
     611            0 :                             Some(format!("{}@{}", blob_meta.key, blob_meta.lsn))
     612              :                         }
     613            0 :                     })
     614            0 :                     .join(", ");
     615            0 : 
     616            0 :                 if !offenders.is_empty() {
     617            0 :                     tracing::warn!(
     618            0 :                         "Oversized vectored read ({} > {}) for keys {}",
     619              :                         buf_size,
     620              :                         max_vectored_read_bytes,
     621              :                         offenders
     622              :                     );
     623            0 :                 }
     624        45048 :             }
     625              : 
     626        45048 :             let read_extend_residency = this.clone();
     627        45048 :             let read_from = self.file.clone();
     628        45048 :             let read_ctx = ctx.attached_child();
     629        45048 :             reconstruct_state
     630        45048 :                 .spawn_io(async move {
     631        45048 :                     let buf = IoBufferMut::with_capacity(buf_size);
     632        45048 :                     let vectored_blob_reader = VectoredBlobReader::new(&read_from);
     633        45048 :                     let res = vectored_blob_reader.read_blobs(&read, buf, &read_ctx).await;
     634              : 
     635        45048 :                     match res {
     636        45048 :                         Ok(blobs_buf) => {
     637        45048 :                             let view = BufView::new_slice(&blobs_buf.buf);
     638       109200 :                             for meta in blobs_buf.blobs.iter() {
     639       109200 :                                 let io: OnDiskValueIo =
     640       109200 :                                     ios.remove(&(meta.meta.key, meta.meta.lsn)).unwrap();
     641       109200 :                                 let img_buf = meta.read(&view).await;
     642              : 
     643       109200 :                                 let img_buf = match img_buf {
     644       109200 :                                     Ok(img_buf) => img_buf,
     645            0 :                                     Err(e) => {
     646            0 :                                         io.complete(Err(e));
     647            0 :                                         continue;
     648              :                                     }
     649              :                                 };
     650              : 
     651       109200 :                                 io.complete(Ok(OnDiskValue::RawImage(img_buf.into_bytes())));
     652              :                             }
     653              : 
     654        45048 :                             assert!(ios.is_empty());
     655              :                         }
     656            0 :                         Err(err) => {
     657            0 :                             for (_, io) in ios {
     658            0 :                                 io.complete(Err(std::io::Error::new(
     659            0 :                                     err.kind(),
     660            0 :                                     "vec read failed",
     661            0 :                                 )));
     662            0 :                             }
     663              :                         }
     664              :                     }
     665              : 
     666              :                     // keep layer resident until this IO is done; this spawned IO future generally outlives the
     667              :                     // call to `self` / the `Arc<DownloadedLayer>` / the `ResidentLayer` that guarantees residency
     668        45048 :                     drop(read_extend_residency);
     669        45048 :                 })
     670        45048 :                 .await;
     671              :         }
     672        45124 :     }
     673              : 
     674          244 :     pub(crate) fn iter<'a>(&'a self, ctx: &'a RequestContext) -> ImageLayerIterator<'a> {
     675          244 :         let block_reader = FileBlockReader::new(&self.file, self.file_id);
     676          244 :         let tree_reader =
     677          244 :             DiskBtreeReader::new(self.index_start_blk, self.index_root_blk, block_reader);
     678          244 :         ImageLayerIterator {
     679          244 :             image_layer: self,
     680          244 :             ctx,
     681          244 :             index_iter: tree_reader.iter(&[0; KEY_SIZE], ctx),
     682          244 :             key_values_batch: VecDeque::new(),
     683          244 :             is_end: false,
     684          244 :             planner: StreamingVectoredReadPlanner::new(
     685          244 :                 1024 * 8192, // The default value. Unit tests might use a different value. 1024 * 8K = 8MB buffer.
     686          244 :                 1024,        // The default value. Unit tests might use a different value
     687          244 :             ),
     688          244 :         }
     689          244 :     }
     690              : 
     691              :     /// NB: not super efficient, but not terrible either. Should prob be an iterator.
     692              :     //
     693              :     // We're reusing the index traversal logical in plan_reads; would be nice to
     694              :     // factor that out.
     695            0 :     pub(crate) async fn load_keys(&self, ctx: &RequestContext) -> anyhow::Result<Vec<Key>> {
     696            0 :         let plan = self
     697            0 :             .plan_reads(KeySpace::single(self.key_range.clone()), None, ctx)
     698            0 :             .await?;
     699            0 :         Ok(plan
     700            0 :             .into_iter()
     701            0 :             .flat_map(|read| read.blobs_at)
     702            0 :             .map(|(_, blob_meta)| blob_meta.key)
     703            0 :             .collect())
     704            0 :     }
     705              : }
     706              : 
     707              : /// A builder object for constructing a new image layer.
     708              : ///
     709              : /// Usage:
     710              : ///
     711              : /// 1. Create the ImageLayerWriter by calling ImageLayerWriter::new(...)
     712              : ///
     713              : /// 2. Write the contents by calling `put_page_image` for every key-value
     714              : ///    pair in the key range.
     715              : ///
     716              : /// 3. Call `finish`.
     717              : ///
     718              : struct ImageLayerWriterInner {
     719              :     conf: &'static PageServerConf,
     720              :     path: Utf8PathBuf,
     721              :     timeline_id: TimelineId,
     722              :     tenant_shard_id: TenantShardId,
     723              :     key_range: Range<Key>,
     724              :     lsn: Lsn,
     725              : 
     726              :     // Total uncompressed bytes passed into put_image
     727              :     uncompressed_bytes: u64,
     728              : 
     729              :     // Like `uncompressed_bytes`,
     730              :     // but only of images we might consider for compression
     731              :     uncompressed_bytes_eligible: u64,
     732              : 
     733              :     // Like `uncompressed_bytes`, but only of images
     734              :     // where we have chosen their compressed form
     735              :     uncompressed_bytes_chosen: u64,
     736              : 
     737              :     // Number of keys in the layer.
     738              :     num_keys: usize,
     739              : 
     740              :     blob_writer: BlobWriter<false>,
     741              :     tree: DiskBtreeBuilder<BlockBuf, KEY_SIZE>,
     742              : 
     743              :     #[cfg(feature = "testing")]
     744              :     last_written_key: Key,
     745              : }
     746              : 
     747              : impl ImageLayerWriterInner {
     748              :     ///
     749              :     /// Start building a new image layer.
     750              :     ///
     751         1204 :     async fn new(
     752         1204 :         conf: &'static PageServerConf,
     753         1204 :         timeline_id: TimelineId,
     754         1204 :         tenant_shard_id: TenantShardId,
     755         1204 :         key_range: &Range<Key>,
     756         1204 :         lsn: Lsn,
     757         1204 :         ctx: &RequestContext,
     758         1204 :     ) -> anyhow::Result<Self> {
     759         1204 :         // Create the file initially with a temporary filename.
     760         1204 :         // We'll atomically rename it to the final name when we're done.
     761         1204 :         let path = ImageLayer::temp_path_for(
     762         1204 :             conf,
     763         1204 :             timeline_id,
     764         1204 :             tenant_shard_id,
     765         1204 :             &ImageLayerName {
     766         1204 :                 key_range: key_range.clone(),
     767         1204 :                 lsn,
     768         1204 :             },
     769         1204 :         );
     770         1204 :         trace!("creating image layer {}", path);
     771         1204 :         let mut file = {
     772         1204 :             VirtualFile::open_with_options(
     773         1204 :                 &path,
     774         1204 :                 virtual_file::OpenOptions::new()
     775         1204 :                     .write(true)
     776         1204 :                     .create_new(true),
     777         1204 :                 ctx,
     778         1204 :             )
     779         1204 :             .await?
     780              :         };
     781              :         // make room for the header block
     782         1204 :         file.seek(SeekFrom::Start(PAGE_SZ as u64)).await?;
     783         1204 :         let blob_writer = BlobWriter::new(file, PAGE_SZ as u64);
     784         1204 : 
     785         1204 :         // Initialize the b-tree index builder
     786         1204 :         let block_buf = BlockBuf::new();
     787         1204 :         let tree_builder = DiskBtreeBuilder::new(block_buf);
     788         1204 : 
     789         1204 :         let writer = Self {
     790         1204 :             conf,
     791         1204 :             path,
     792         1204 :             timeline_id,
     793         1204 :             tenant_shard_id,
     794         1204 :             key_range: key_range.clone(),
     795         1204 :             lsn,
     796         1204 :             tree: tree_builder,
     797         1204 :             blob_writer,
     798         1204 :             uncompressed_bytes: 0,
     799         1204 :             uncompressed_bytes_eligible: 0,
     800         1204 :             uncompressed_bytes_chosen: 0,
     801         1204 :             num_keys: 0,
     802         1204 :             #[cfg(feature = "testing")]
     803         1204 :             last_written_key: Key::MIN,
     804         1204 :         };
     805         1204 : 
     806         1204 :         Ok(writer)
     807         1204 :     }
     808              : 
     809              :     ///
     810              :     /// Write next value to the file.
     811              :     ///
     812              :     /// The page versions must be appended in blknum order.
     813              :     ///
     814      1093372 :     async fn put_image(
     815      1093372 :         &mut self,
     816      1093372 :         key: Key,
     817      1093372 :         img: Bytes,
     818      1093372 :         ctx: &RequestContext,
     819      1093372 :     ) -> anyhow::Result<()> {
     820      1093372 :         ensure!(self.key_range.contains(&key));
     821      1093372 :         let compression = self.conf.image_compression;
     822      1093372 :         let uncompressed_len = img.len() as u64;
     823      1093372 :         self.uncompressed_bytes += uncompressed_len;
     824      1093372 :         self.num_keys += 1;
     825      1093372 :         let (_img, res) = self
     826      1093372 :             .blob_writer
     827      1093372 :             .write_blob_maybe_compressed(img.slice_len(), ctx, compression)
     828      1093372 :             .await;
     829              :         // TODO: re-use the buffer for `img` further upstack
     830      1093372 :         let (off, compression_info) = res?;
     831      1093372 :         if compression_info.compressed_size.is_some() {
     832        16004 :             // The image has been considered for compression at least
     833        16004 :             self.uncompressed_bytes_eligible += uncompressed_len;
     834      1077368 :         }
     835      1093372 :         if compression_info.written_compressed {
     836            0 :             // The image has been compressed
     837            0 :             self.uncompressed_bytes_chosen += uncompressed_len;
     838      1093372 :         }
     839              : 
     840      1093372 :         let mut keybuf: [u8; KEY_SIZE] = [0u8; KEY_SIZE];
     841      1093372 :         key.write_to_byte_slice(&mut keybuf);
     842      1093372 :         self.tree.append(&keybuf, off)?;
     843              : 
     844              :         #[cfg(feature = "testing")]
     845      1093372 :         {
     846      1093372 :             self.last_written_key = key;
     847      1093372 :         }
     848      1093372 : 
     849      1093372 :         Ok(())
     850      1093372 :     }
     851              : 
     852              :     ///
     853              :     /// Finish writing the image layer.
     854              :     ///
     855          732 :     async fn finish(
     856          732 :         self,
     857          732 :         ctx: &RequestContext,
     858          732 :         end_key: Option<Key>,
     859          732 :     ) -> anyhow::Result<(PersistentLayerDesc, Utf8PathBuf)> {
     860          732 :         let temp_path = self.path.clone();
     861          732 :         let result = self.finish0(ctx, end_key).await;
     862          732 :         if let Err(ref e) = result {
     863            0 :             tracing::info!(%temp_path, "cleaning up temporary file after error during writing: {e}");
     864            0 :             if let Err(e) = std::fs::remove_file(&temp_path) {
     865            0 :                 tracing::warn!(error=%e, %temp_path, "error cleaning up temporary layer file after error during writing");
     866            0 :             }
     867          732 :         }
     868          732 :         result
     869          732 :     }
     870              : 
     871              :     ///
     872              :     /// Finish writing the image layer.
     873              :     ///
     874          732 :     async fn finish0(
     875          732 :         self,
     876          732 :         ctx: &RequestContext,
     877          732 :         end_key: Option<Key>,
     878          732 :     ) -> anyhow::Result<(PersistentLayerDesc, Utf8PathBuf)> {
     879          732 :         let index_start_blk = self.blob_writer.size().div_ceil(PAGE_SZ as u64) as u32;
     880          732 : 
     881          732 :         // Calculate compression ratio
     882          732 :         let compressed_size = self.blob_writer.size() - PAGE_SZ as u64; // Subtract PAGE_SZ for header
     883          732 :         crate::metrics::COMPRESSION_IMAGE_INPUT_BYTES.inc_by(self.uncompressed_bytes);
     884          732 :         crate::metrics::COMPRESSION_IMAGE_INPUT_BYTES_CONSIDERED
     885          732 :             .inc_by(self.uncompressed_bytes_eligible);
     886          732 :         crate::metrics::COMPRESSION_IMAGE_INPUT_BYTES_CHOSEN.inc_by(self.uncompressed_bytes_chosen);
     887          732 :         crate::metrics::COMPRESSION_IMAGE_OUTPUT_BYTES.inc_by(compressed_size);
     888          732 : 
     889          732 :         let mut file = self.blob_writer.into_inner();
     890          732 : 
     891          732 :         // Write out the index
     892          732 :         file.seek(SeekFrom::Start(index_start_blk as u64 * PAGE_SZ as u64))
     893          732 :             .await?;
     894          732 :         let (index_root_blk, block_buf) = self.tree.finish()?;
     895         2376 :         for buf in block_buf.blocks {
     896         1644 :             let (_buf, res) = file.write_all(buf.slice_len(), ctx).await;
     897         1644 :             res?;
     898              :         }
     899              : 
     900          732 :         let final_key_range = if let Some(end_key) = end_key {
     901          584 :             self.key_range.start..end_key
     902              :         } else {
     903          148 :             self.key_range.clone()
     904              :         };
     905              : 
     906              :         // Fill in the summary on blk 0
     907          732 :         let summary = Summary {
     908          732 :             magic: IMAGE_FILE_MAGIC,
     909          732 :             format_version: STORAGE_FORMAT_VERSION,
     910          732 :             tenant_id: self.tenant_shard_id.tenant_id,
     911          732 :             timeline_id: self.timeline_id,
     912          732 :             key_range: final_key_range.clone(),
     913          732 :             lsn: self.lsn,
     914          732 :             index_start_blk,
     915          732 :             index_root_blk,
     916          732 :         };
     917          732 : 
     918          732 :         let mut buf = Vec::with_capacity(PAGE_SZ);
     919          732 :         // TODO: could use smallvec here but it's a pain with Slice<T>
     920          732 :         Summary::ser_into(&summary, &mut buf)?;
     921          732 :         file.seek(SeekFrom::Start(0)).await?;
     922          732 :         let (_buf, res) = file.write_all(buf.slice_len(), ctx).await;
     923          732 :         res?;
     924              : 
     925          732 :         let metadata = file
     926          732 :             .metadata()
     927          732 :             .await
     928          732 :             .context("get metadata to determine file size")?;
     929              : 
     930          732 :         let desc = PersistentLayerDesc::new_img(
     931          732 :             self.tenant_shard_id,
     932          732 :             self.timeline_id,
     933          732 :             final_key_range,
     934          732 :             self.lsn,
     935          732 :             metadata.len(),
     936          732 :         );
     937              : 
     938              :         #[cfg(feature = "testing")]
     939          732 :         if let Some(end_key) = end_key {
     940          584 :             assert!(
     941          584 :                 self.last_written_key < end_key,
     942            0 :                 "written key violates end_key range"
     943              :             );
     944          148 :         }
     945              : 
     946              :         // Note: Because we open the file in write-only mode, we cannot
     947              :         // reuse the same VirtualFile for reading later. That's why we don't
     948              :         // set inner.file here. The first read will have to re-open it.
     949              : 
     950              :         // fsync the file
     951          732 :         file.sync_all()
     952          732 :             .await
     953          732 :             .maybe_fatal_err("image_layer sync_all")?;
     954              : 
     955          732 :         trace!("created image layer {}", self.path);
     956              : 
     957          732 :         Ok((desc, self.path))
     958          732 :     }
     959              : }
     960              : 
     961              : /// A builder object for constructing a new image layer.
     962              : ///
     963              : /// Usage:
     964              : ///
     965              : /// 1. Create the ImageLayerWriter by calling ImageLayerWriter::new(...)
     966              : ///
     967              : /// 2. Write the contents by calling `put_page_image` for every key-value
     968              : ///    pair in the key range.
     969              : ///
     970              : /// 3. Call `finish`.
     971              : ///
     972              : /// # Note
     973              : ///
     974              : /// As described in <https://github.com/neondatabase/neon/issues/2650>, it's
     975              : /// possible for the writer to drop before `finish` is actually called. So this
     976              : /// could lead to odd temporary files in the directory, exhausting file system.
     977              : /// This structure wraps `ImageLayerWriterInner` and also contains `Drop`
     978              : /// implementation that cleans up the temporary file in failure. It's not
     979              : /// possible to do this directly in `ImageLayerWriterInner` since `finish` moves
     980              : /// out some fields, making it impossible to implement `Drop`.
     981              : ///
     982              : #[must_use]
     983              : pub struct ImageLayerWriter {
     984              :     inner: Option<ImageLayerWriterInner>,
     985              : }
     986              : 
     987              : impl ImageLayerWriter {
     988              :     ///
     989              :     /// Start building a new image layer.
     990              :     ///
     991         1204 :     pub async fn new(
     992         1204 :         conf: &'static PageServerConf,
     993         1204 :         timeline_id: TimelineId,
     994         1204 :         tenant_shard_id: TenantShardId,
     995         1204 :         key_range: &Range<Key>,
     996         1204 :         lsn: Lsn,
     997         1204 :         ctx: &RequestContext,
     998         1204 :     ) -> anyhow::Result<ImageLayerWriter> {
     999         1204 :         Ok(Self {
    1000         1204 :             inner: Some(
    1001         1204 :                 ImageLayerWriterInner::new(conf, timeline_id, tenant_shard_id, key_range, lsn, ctx)
    1002         1204 :                     .await?,
    1003              :             ),
    1004              :         })
    1005         1204 :     }
    1006              : 
    1007              :     ///
    1008              :     /// Write next value to the file.
    1009              :     ///
    1010              :     /// The page versions must be appended in blknum order.
    1011              :     ///
    1012      1093372 :     pub async fn put_image(
    1013      1093372 :         &mut self,
    1014      1093372 :         key: Key,
    1015      1093372 :         img: Bytes,
    1016      1093372 :         ctx: &RequestContext,
    1017      1093372 :     ) -> anyhow::Result<()> {
    1018      1093372 :         self.inner.as_mut().unwrap().put_image(key, img, ctx).await
    1019      1093372 :     }
    1020              : 
    1021              :     /// Estimated size of the image layer.
    1022        17080 :     pub(crate) fn estimated_size(&self) -> u64 {
    1023        17080 :         let inner = self.inner.as_ref().unwrap();
    1024        17080 :         inner.blob_writer.size() + inner.tree.borrow_writer().size() + PAGE_SZ as u64
    1025        17080 :     }
    1026              : 
    1027        17272 :     pub(crate) fn num_keys(&self) -> usize {
    1028        17272 :         self.inner.as_ref().unwrap().num_keys
    1029        17272 :     }
    1030              : 
    1031              :     ///
    1032              :     /// Finish writing the image layer.
    1033              :     ///
    1034          148 :     pub(crate) async fn finish(
    1035          148 :         mut self,
    1036          148 :         ctx: &RequestContext,
    1037          148 :     ) -> anyhow::Result<(PersistentLayerDesc, Utf8PathBuf)> {
    1038          148 :         self.inner.take().unwrap().finish(ctx, None).await
    1039          148 :     }
    1040              : 
    1041              :     /// Finish writing the image layer with an end key, used in [`super::batch_split_writer::SplitImageLayerWriter`]. The end key determines the end of the image layer's covered range and is exclusive.
    1042          584 :     pub(super) async fn finish_with_end_key(
    1043          584 :         mut self,
    1044          584 :         end_key: Key,
    1045          584 :         ctx: &RequestContext,
    1046          584 :     ) -> anyhow::Result<(PersistentLayerDesc, Utf8PathBuf)> {
    1047          584 :         self.inner.take().unwrap().finish(ctx, Some(end_key)).await
    1048          584 :     }
    1049              : }
    1050              : 
    1051              : impl Drop for ImageLayerWriter {
    1052         1204 :     fn drop(&mut self) {
    1053         1204 :         if let Some(inner) = self.inner.take() {
    1054          472 :             inner.blob_writer.into_inner().remove();
    1055          732 :         }
    1056         1204 :     }
    1057              : }
    1058              : 
    1059              : pub struct ImageLayerIterator<'a> {
    1060              :     image_layer: &'a ImageLayerInner,
    1061              :     ctx: &'a RequestContext,
    1062              :     planner: StreamingVectoredReadPlanner,
    1063              :     index_iter: DiskBtreeIterator<'a>,
    1064              :     key_values_batch: VecDeque<(Key, Lsn, Value)>,
    1065              :     is_end: bool,
    1066              : }
    1067              : 
    1068              : impl ImageLayerIterator<'_> {
    1069            0 :     pub(crate) fn layer_dbg_info(&self) -> String {
    1070            0 :         self.image_layer.layer_dbg_info()
    1071            0 :     }
    1072              : 
    1073              :     /// Retrieve a batch of key-value pairs into the iterator buffer.
    1074        38252 :     async fn next_batch(&mut self) -> anyhow::Result<()> {
    1075        38252 :         assert!(self.key_values_batch.is_empty());
    1076        38252 :         assert!(!self.is_end);
    1077              : 
    1078        38252 :         let plan = loop {
    1079        58072 :             if let Some(res) = self.index_iter.next().await {
    1080        57884 :                 let (raw_key, offset) = res?;
    1081        57884 :                 if let Some(batch_plan) = self.planner.handle(
    1082        57884 :                     Key::from_slice(&raw_key[..KEY_SIZE]),
    1083        57884 :                     self.image_layer.lsn,
    1084        57884 :                     offset,
    1085        57884 :                     true,
    1086        57884 :                 ) {
    1087        38064 :                     break batch_plan;
    1088        19820 :                 }
    1089              :             } else {
    1090          188 :                 self.is_end = true;
    1091          188 :                 let payload_end = self.image_layer.index_start_blk as u64 * PAGE_SZ as u64;
    1092          188 :                 if let Some(item) = self.planner.handle_range_end(payload_end) {
    1093          188 :                     break item;
    1094              :                 } else {
    1095            0 :                     return Ok(()); // TODO: a test case on empty iterator
    1096              :                 }
    1097              :             }
    1098              :         };
    1099        38252 :         let vectored_blob_reader = VectoredBlobReader::new(&self.image_layer.file);
    1100        38252 :         let mut next_batch = std::collections::VecDeque::new();
    1101        38252 :         let buf_size = plan.size();
    1102        38252 :         let buf = IoBufferMut::with_capacity(buf_size);
    1103        38252 :         let blobs_buf = vectored_blob_reader
    1104        38252 :             .read_blobs(&plan, buf, self.ctx)
    1105        38252 :             .await?;
    1106        38252 :         let view = BufView::new_slice(&blobs_buf.buf);
    1107        57828 :         for meta in blobs_buf.blobs.iter() {
    1108        57828 :             let img_buf = meta.read(&view).await?;
    1109        57828 :             next_batch.push_back((
    1110        57828 :                 meta.meta.key,
    1111        57828 :                 self.image_layer.lsn,
    1112        57828 :                 Value::Image(img_buf.into_bytes()),
    1113        57828 :             ));
    1114              :         }
    1115        38252 :         self.key_values_batch = next_batch;
    1116        38252 :         Ok(())
    1117        38252 :     }
    1118              : 
    1119        57608 :     pub async fn next(&mut self) -> anyhow::Result<Option<(Key, Lsn, Value)>> {
    1120        57608 :         if self.key_values_batch.is_empty() {
    1121        38404 :             if self.is_end {
    1122          320 :                 return Ok(None);
    1123        38084 :             }
    1124        38084 :             self.next_batch().await?;
    1125        19204 :         }
    1126        57288 :         Ok(Some(
    1127        57288 :             self.key_values_batch
    1128        57288 :                 .pop_front()
    1129        57288 :                 .expect("should not be empty"),
    1130        57288 :         ))
    1131        57608 :     }
    1132              : }
    1133              : 
    1134              : #[cfg(test)]
    1135              : mod test {
    1136              :     use std::sync::Arc;
    1137              :     use std::time::Duration;
    1138              : 
    1139              :     use bytes::Bytes;
    1140              :     use itertools::Itertools;
    1141              :     use pageserver_api::key::Key;
    1142              :     use pageserver_api::shard::{ShardCount, ShardIdentity, ShardNumber, ShardStripeSize};
    1143              :     use pageserver_api::value::Value;
    1144              :     use utils::generation::Generation;
    1145              :     use utils::id::{TenantId, TimelineId};
    1146              :     use utils::lsn::Lsn;
    1147              : 
    1148              :     use super::{ImageLayerIterator, ImageLayerWriter};
    1149              :     use crate::DEFAULT_PG_VERSION;
    1150              :     use crate::context::RequestContext;
    1151              :     use crate::tenant::config::TenantConf;
    1152              :     use crate::tenant::harness::{TIMELINE_ID, TenantHarness};
    1153              :     use crate::tenant::storage_layer::{Layer, ResidentLayer};
    1154              :     use crate::tenant::vectored_blob_io::StreamingVectoredReadPlanner;
    1155              :     use crate::tenant::{Tenant, Timeline};
    1156              : 
    1157              :     #[tokio::test]
    1158            4 :     async fn image_layer_rewrite() {
    1159            4 :         let tenant_conf = TenantConf {
    1160            4 :             gc_period: Duration::ZERO,
    1161            4 :             compaction_period: Duration::ZERO,
    1162            4 :             ..TenantConf::default()
    1163            4 :         };
    1164            4 :         let tenant_id = TenantId::generate();
    1165            4 :         let mut gen_ = Generation::new(0xdead0001);
    1166           20 :         let mut get_next_gen = || {
    1167           20 :             let ret = gen_;
    1168           20 :             gen_ = gen_.next();
    1169           20 :             ret
    1170           20 :         };
    1171            4 :         // The LSN at which we will create an image layer to filter
    1172            4 :         let lsn = Lsn(0xdeadbeef0000);
    1173            4 :         let timeline_id = TimelineId::generate();
    1174            4 : 
    1175            4 :         //
    1176            4 :         // Create an unsharded parent with a layer.
    1177            4 :         //
    1178            4 : 
    1179            4 :         let harness = TenantHarness::create_custom(
    1180            4 :             "test_image_layer_rewrite--parent",
    1181            4 :             tenant_conf.clone(),
    1182            4 :             tenant_id,
    1183            4 :             ShardIdentity::unsharded(),
    1184            4 :             get_next_gen(),
    1185            4 :         )
    1186            4 :         .await
    1187            4 :         .unwrap();
    1188            4 :         let (tenant, ctx) = harness.load().await;
    1189            4 :         let timeline = tenant
    1190            4 :             .create_test_timeline(timeline_id, lsn, DEFAULT_PG_VERSION, &ctx)
    1191            4 :             .await
    1192            4 :             .unwrap();
    1193            4 : 
    1194            4 :         // This key range contains several 0x8000 page stripes, only one of which belongs to shard zero
    1195            4 :         let input_start = Key::from_hex("000000067f00000001000000ae0000000000").unwrap();
    1196            4 :         let input_end = Key::from_hex("000000067f00000001000000ae0000020000").unwrap();
    1197            4 :         let range = input_start..input_end;
    1198            4 : 
    1199            4 :         // Build an image layer to filter
    1200            4 :         let resident = {
    1201            4 :             let mut writer = ImageLayerWriter::new(
    1202            4 :                 harness.conf,
    1203            4 :                 timeline_id,
    1204            4 :                 harness.tenant_shard_id,
    1205            4 :                 &range,
    1206            4 :                 lsn,
    1207            4 :                 &ctx,
    1208            4 :             )
    1209            4 :             .await
    1210            4 :             .unwrap();
    1211            4 : 
    1212            4 :             let foo_img = Bytes::from_static(&[1, 2, 3, 4]);
    1213            4 :             let mut key = range.start;
    1214       524292 :             while key < range.end {
    1215       524288 :                 writer.put_image(key, foo_img.clone(), &ctx).await.unwrap();
    1216       524288 : 
    1217       524288 :                 key = key.next();
    1218            4 :             }
    1219            4 :             let (desc, path) = writer.finish(&ctx).await.unwrap();
    1220            4 :             Layer::finish_creating(tenant.conf, &timeline, desc, &path).unwrap()
    1221            4 :         };
    1222            4 :         let original_size = resident.metadata().file_size;
    1223            4 : 
    1224            4 :         //
    1225            4 :         // Create child shards and do the rewrite, exercising filter().
    1226            4 :         // TODO: abstraction in TenantHarness for splits.
    1227            4 :         //
    1228            4 : 
    1229            4 :         // Filter for various shards: this exercises cases like values at start of key range, end of key
    1230            4 :         // range, middle of key range.
    1231            4 :         let shard_count = ShardCount::new(4);
    1232           16 :         for shard_number in 0..shard_count.count() {
    1233            4 :             //
    1234            4 :             // mimic the shard split
    1235            4 :             //
    1236           16 :             let shard_identity = ShardIdentity::new(
    1237           16 :                 ShardNumber(shard_number),
    1238           16 :                 shard_count,
    1239           16 :                 ShardStripeSize(0x8000),
    1240           16 :             )
    1241           16 :             .unwrap();
    1242           16 :             let harness = TenantHarness::create_custom(
    1243           16 :                 Box::leak(Box::new(format!(
    1244           16 :                     "test_image_layer_rewrite--child{}",
    1245           16 :                     shard_identity.shard_slug()
    1246           16 :                 ))),
    1247           16 :                 tenant_conf.clone(),
    1248           16 :                 tenant_id,
    1249           16 :                 shard_identity,
    1250           16 :                 // NB: in reality, the shards would each fork off their own gen number sequence from the parent.
    1251           16 :                 // But here, all we care about is that the gen number is unique.
    1252           16 :                 get_next_gen(),
    1253           16 :             )
    1254           16 :             .await
    1255           16 :             .unwrap();
    1256           16 :             let (tenant, ctx) = harness.load().await;
    1257           16 :             let timeline = tenant
    1258           16 :                 .create_test_timeline(timeline_id, lsn, DEFAULT_PG_VERSION, &ctx)
    1259           16 :                 .await
    1260           16 :                 .unwrap();
    1261            4 : 
    1262            4 :             //
    1263            4 :             // use filter() and make assertions
    1264            4 :             //
    1265            4 : 
    1266           16 :             let mut filtered_writer = ImageLayerWriter::new(
    1267           16 :                 harness.conf,
    1268           16 :                 timeline_id,
    1269           16 :                 harness.tenant_shard_id,
    1270           16 :                 &range,
    1271           16 :                 lsn,
    1272           16 :                 &ctx,
    1273           16 :             )
    1274           16 :             .await
    1275           16 :             .unwrap();
    1276            4 : 
    1277           16 :             let wrote_keys = resident
    1278           16 :                 .filter(&shard_identity, &mut filtered_writer, &ctx)
    1279           16 :                 .await
    1280           16 :                 .unwrap();
    1281           16 :             let replacement = if wrote_keys > 0 {
    1282           12 :                 let (desc, path) = filtered_writer.finish(&ctx).await.unwrap();
    1283           12 :                 let resident = Layer::finish_creating(tenant.conf, &timeline, desc, &path).unwrap();
    1284           12 :                 Some(resident)
    1285            4 :             } else {
    1286            4 :                 None
    1287            4 :             };
    1288            4 : 
    1289            4 :             // This exact size and those below will need updating as/when the layer encoding changes, but
    1290            4 :             // should be deterministic for a given version of the format, as we used no randomness generating the input.
    1291           16 :             assert_eq!(original_size, 1597440);
    1292            4 : 
    1293           16 :             match shard_number {
    1294            4 :                 0 => {
    1295            4 :                     // We should have written out just one stripe for our shard identity
    1296            4 :                     assert_eq!(wrote_keys, 0x8000);
    1297            4 :                     let replacement = replacement.unwrap();
    1298            4 : 
    1299            4 :                     // We should have dropped some of the data
    1300            4 :                     assert!(replacement.metadata().file_size < original_size);
    1301            4 :                     assert!(replacement.metadata().file_size > 0);
    1302            4 : 
    1303            4 :                     // Assert that we dropped ~3/4 of the data.
    1304            4 :                     assert_eq!(replacement.metadata().file_size, 417792);
    1305            4 :                 }
    1306            4 :                 1 => {
    1307            4 :                     // Shard 1 has no keys in our input range
    1308            4 :                     assert_eq!(wrote_keys, 0x0);
    1309            4 :                     assert!(replacement.is_none());
    1310            4 :                 }
    1311            4 :                 2 => {
    1312            4 :                     // Shard 2 has one stripes in the input range
    1313            4 :                     assert_eq!(wrote_keys, 0x8000);
    1314            4 :                     let replacement = replacement.unwrap();
    1315            4 :                     assert!(replacement.metadata().file_size < original_size);
    1316            4 :                     assert!(replacement.metadata().file_size > 0);
    1317            4 :                     assert_eq!(replacement.metadata().file_size, 417792);
    1318            4 :                 }
    1319            4 :                 3 => {
    1320            4 :                     // Shard 3 has two stripes in the input range
    1321            4 :                     assert_eq!(wrote_keys, 0x10000);
    1322            4 :                     let replacement = replacement.unwrap();
    1323            4 :                     assert!(replacement.metadata().file_size < original_size);
    1324            4 :                     assert!(replacement.metadata().file_size > 0);
    1325            4 :                     assert_eq!(replacement.metadata().file_size, 811008);
    1326            4 :                 }
    1327            4 :                 _ => unreachable!(),
    1328            4 :             }
    1329            4 :         }
    1330            4 :     }
    1331              : 
    1332            4 :     async fn produce_image_layer(
    1333            4 :         tenant: &Tenant,
    1334            4 :         tline: &Arc<Timeline>,
    1335            4 :         mut images: Vec<(Key, Bytes)>,
    1336            4 :         lsn: Lsn,
    1337            4 :         ctx: &RequestContext,
    1338            4 :     ) -> anyhow::Result<ResidentLayer> {
    1339            4 :         images.sort();
    1340            4 :         let (key_start, _) = images.first().unwrap();
    1341            4 :         let (key_last, _) = images.last().unwrap();
    1342            4 :         let key_end = key_last.next();
    1343            4 :         let key_range = *key_start..key_end;
    1344            4 :         let mut writer = ImageLayerWriter::new(
    1345            4 :             tenant.conf,
    1346            4 :             tline.timeline_id,
    1347            4 :             tenant.tenant_shard_id,
    1348            4 :             &key_range,
    1349            4 :             lsn,
    1350            4 :             ctx,
    1351            4 :         )
    1352            4 :         .await?;
    1353              : 
    1354         4004 :         for (key, img) in images {
    1355         4000 :             writer.put_image(key, img, ctx).await?;
    1356              :         }
    1357            4 :         let (desc, path) = writer.finish(ctx).await?;
    1358            4 :         let img_layer = Layer::finish_creating(tenant.conf, tline, desc, &path)?;
    1359              : 
    1360            4 :         Ok::<_, anyhow::Error>(img_layer)
    1361            4 :     }
    1362              : 
    1363           56 :     async fn assert_img_iter_equal(
    1364           56 :         img_iter: &mut ImageLayerIterator<'_>,
    1365           56 :         expect: &[(Key, Bytes)],
    1366           56 :         expect_lsn: Lsn,
    1367           56 :     ) {
    1368           56 :         let mut expect_iter = expect.iter();
    1369              :         loop {
    1370        56056 :             let o1 = img_iter.next().await.unwrap();
    1371        56056 :             let o2 = expect_iter.next();
    1372        56056 :             match (o1, o2) {
    1373           56 :                 (None, None) => break,
    1374        56000 :                 (Some((k1, l1, v1)), Some((k2, i2))) => {
    1375        56000 :                     let Value::Image(i1) = v1 else {
    1376            0 :                         panic!("expect Value::Image")
    1377              :                     };
    1378        56000 :                     assert_eq!(&k1, k2);
    1379        56000 :                     assert_eq!(l1, expect_lsn);
    1380        56000 :                     assert_eq!(&i1, i2);
    1381              :                 }
    1382            0 :                 (o1, o2) => panic!("iterators length mismatch: {:?}, {:?}", o1, o2),
    1383              :             }
    1384              :         }
    1385           56 :     }
    1386              : 
    1387              :     #[tokio::test]
    1388            4 :     async fn image_layer_iterator() {
    1389            4 :         let harness = TenantHarness::create("image_layer_iterator").await.unwrap();
    1390            4 :         let (tenant, ctx) = harness.load().await;
    1391            4 : 
    1392            4 :         let tline = tenant
    1393            4 :             .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
    1394            4 :             .await
    1395            4 :             .unwrap();
    1396            4 : 
    1397         4000 :         fn get_key(id: u32) -> Key {
    1398         4000 :             let mut key = Key::from_hex("000000000033333333444444445500000000").unwrap();
    1399         4000 :             key.field6 = id;
    1400         4000 :             key
    1401         4000 :         }
    1402            4 :         const N: usize = 1000;
    1403            4 :         let test_imgs = (0..N)
    1404         4000 :             .map(|idx| (get_key(idx as u32), Bytes::from(format!("img{idx:05}"))))
    1405            4 :             .collect_vec();
    1406            4 :         let resident_layer =
    1407            4 :             produce_image_layer(&tenant, &tline, test_imgs.clone(), Lsn(0x10), &ctx)
    1408            4 :                 .await
    1409            4 :                 .unwrap();
    1410            4 :         let img_layer = resident_layer.get_as_image(&ctx).await.unwrap();
    1411           12 :         for max_read_size in [1, 1024] {
    1412           64 :             for batch_size in [1, 2, 4, 8, 3, 7, 13] {
    1413           56 :                 println!("running with batch_size={batch_size} max_read_size={max_read_size}");
    1414           56 :                 // Test if the batch size is correctly determined
    1415           56 :                 let mut iter = img_layer.iter(&ctx);
    1416           56 :                 iter.planner = StreamingVectoredReadPlanner::new(max_read_size, batch_size);
    1417           56 :                 let mut num_items = 0;
    1418          224 :                 for _ in 0..3 {
    1419          168 :                     iter.next_batch().await.unwrap();
    1420          168 :                     num_items += iter.key_values_batch.len();
    1421          168 :                     if max_read_size == 1 {
    1422            4 :                         // every key should be a batch b/c the value is larger than max_read_size
    1423           84 :                         assert_eq!(iter.key_values_batch.len(), 1);
    1424            4 :                     } else {
    1425           84 :                         assert!(iter.key_values_batch.len() <= batch_size);
    1426            4 :                     }
    1427          168 :                     if num_items >= N {
    1428            4 :                         break;
    1429          168 :                     }
    1430          168 :                     iter.key_values_batch.clear();
    1431            4 :                 }
    1432            4 :                 // Test if the result is correct
    1433           56 :                 let mut iter = img_layer.iter(&ctx);
    1434           56 :                 iter.planner = StreamingVectoredReadPlanner::new(max_read_size, batch_size);
    1435           56 :                 assert_img_iter_equal(&mut iter, &test_imgs, Lsn(0x10)).await;
    1436            4 :             }
    1437            4 :         }
    1438            4 :     }
    1439              : }
        

Generated by: LCOV version 2.1-beta