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

Generated by: LCOV version 2.1-beta