LCOV - code coverage report
Current view: top level - pageserver/src - pgdatadir_mapping.rs (source / functions) Coverage Total Hit
Test: bb45db3982713bfd5bec075773079136e362195e.info Lines: 58.9 % 1668 982
Test Date: 2024-12-11 15:53:32 Functions: 43.5 % 200 87

            Line data    Source code
       1              : //!
       2              : //! This provides an abstraction to store PostgreSQL relations and other files
       3              : //! in the key-value store that implements the Repository interface.
       4              : //!
       5              : //! (TODO: The line between PUT-functions here and walingest.rs is a bit blurry, as
       6              : //! walingest.rs handles a few things like implicit relation creation and extension.
       7              : //! Clarify that)
       8              : //!
       9              : use super::tenant::{PageReconstructError, Timeline};
      10              : use crate::aux_file;
      11              : use crate::context::RequestContext;
      12              : use crate::keyspace::{KeySpace, KeySpaceAccum};
      13              : use crate::metrics::{
      14              :     RELSIZE_CACHE_ENTRIES, RELSIZE_CACHE_HITS, RELSIZE_CACHE_MISSES, RELSIZE_CACHE_MISSES_OLD,
      15              : };
      16              : use crate::span::{
      17              :     debug_assert_current_span_has_tenant_and_timeline_id,
      18              :     debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id,
      19              : };
      20              : use crate::tenant::timeline::GetVectoredError;
      21              : use anyhow::{ensure, Context};
      22              : use bytes::{Buf, Bytes, BytesMut};
      23              : use enum_map::Enum;
      24              : use itertools::Itertools;
      25              : use pageserver_api::key::Key;
      26              : use pageserver_api::key::{
      27              :     dbdir_key_range, rel_block_to_key, rel_dir_to_key, rel_key_range, rel_size_to_key,
      28              :     relmap_file_key, repl_origin_key, repl_origin_key_range, slru_block_to_key, slru_dir_to_key,
      29              :     slru_segment_key_range, slru_segment_size_to_key, twophase_file_key, twophase_key_range,
      30              :     CompactKey, AUX_FILES_KEY, CHECKPOINT_KEY, CONTROLFILE_KEY, DBDIR_KEY, TWOPHASEDIR_KEY,
      31              : };
      32              : use pageserver_api::keyspace::SparseKeySpace;
      33              : use pageserver_api::record::NeonWalRecord;
      34              : use pageserver_api::reltag::{BlockNumber, RelTag, SlruKind};
      35              : use pageserver_api::shard::ShardIdentity;
      36              : use pageserver_api::value::Value;
      37              : use postgres_ffi::relfile_utils::{FSM_FORKNUM, VISIBILITYMAP_FORKNUM};
      38              : use postgres_ffi::BLCKSZ;
      39              : use postgres_ffi::{Oid, RepOriginId, TimestampTz, TransactionId};
      40              : use serde::{Deserialize, Serialize};
      41              : use std::collections::{hash_map, BTreeMap, HashMap, HashSet};
      42              : use std::ops::ControlFlow;
      43              : use std::ops::Range;
      44              : use strum::IntoEnumIterator;
      45              : use tokio_util::sync::CancellationToken;
      46              : use tracing::{debug, trace, warn};
      47              : use utils::bin_ser::DeserializeError;
      48              : use utils::pausable_failpoint;
      49              : use utils::{bin_ser::BeSer, lsn::Lsn};
      50              : use wal_decoder::serialized_batch::SerializedValueBatch;
      51              : 
      52              : /// Max delta records appended to the AUX_FILES_KEY (for aux v1). The write path will write a full image once this threshold is reached.
      53              : pub const MAX_AUX_FILE_DELTAS: usize = 1024;
      54              : 
      55              : /// Max number of aux-file-related delta layers. The compaction will create a new image layer once this threshold is reached.
      56              : pub const MAX_AUX_FILE_V2_DELTAS: usize = 16;
      57              : 
      58              : #[derive(Debug)]
      59              : pub enum LsnForTimestamp {
      60              :     /// Found commits both before and after the given timestamp
      61              :     Present(Lsn),
      62              : 
      63              :     /// Found no commits after the given timestamp, this means
      64              :     /// that the newest data in the branch is older than the given
      65              :     /// timestamp.
      66              :     ///
      67              :     /// All commits <= LSN happened before the given timestamp
      68              :     Future(Lsn),
      69              : 
      70              :     /// The queried timestamp is past our horizon we look back at (PITR)
      71              :     ///
      72              :     /// All commits > LSN happened after the given timestamp,
      73              :     /// but any commits < LSN might have happened before or after
      74              :     /// the given timestamp. We don't know because no data before
      75              :     /// the given lsn is available.
      76              :     Past(Lsn),
      77              : 
      78              :     /// We have found no commit with a timestamp,
      79              :     /// so we can't return anything meaningful.
      80              :     ///
      81              :     /// The associated LSN is the lower bound value we can safely
      82              :     /// create branches on, but no statement is made if it is
      83              :     /// older or newer than the timestamp.
      84              :     ///
      85              :     /// This variant can e.g. be returned right after a
      86              :     /// cluster import.
      87              :     NoData(Lsn),
      88              : }
      89              : 
      90              : #[derive(Debug, thiserror::Error)]
      91              : pub(crate) enum CalculateLogicalSizeError {
      92              :     #[error("cancelled")]
      93              :     Cancelled,
      94              : 
      95              :     /// Something went wrong while reading the metadata we use to calculate logical size
      96              :     /// Note that cancellation variants of `PageReconstructError` are transformed to [`Self::Cancelled`]
      97              :     /// in the `From` implementation for this variant.
      98              :     #[error(transparent)]
      99              :     PageRead(PageReconstructError),
     100              : 
     101              :     /// Something went wrong deserializing metadata that we read to calculate logical size
     102              :     #[error("decode error: {0}")]
     103              :     Decode(#[from] DeserializeError),
     104              : }
     105              : 
     106              : #[derive(Debug, thiserror::Error)]
     107              : pub(crate) enum CollectKeySpaceError {
     108              :     #[error(transparent)]
     109              :     Decode(#[from] DeserializeError),
     110              :     #[error(transparent)]
     111              :     PageRead(PageReconstructError),
     112              :     #[error("cancelled")]
     113              :     Cancelled,
     114              : }
     115              : 
     116              : impl From<PageReconstructError> for CollectKeySpaceError {
     117            0 :     fn from(err: PageReconstructError) -> Self {
     118            0 :         match err {
     119            0 :             PageReconstructError::Cancelled => Self::Cancelled,
     120            0 :             err => Self::PageRead(err),
     121              :         }
     122            0 :     }
     123              : }
     124              : 
     125              : impl From<PageReconstructError> for CalculateLogicalSizeError {
     126            0 :     fn from(pre: PageReconstructError) -> Self {
     127            0 :         match pre {
     128            0 :             PageReconstructError::Cancelled => Self::Cancelled,
     129            0 :             _ => Self::PageRead(pre),
     130              :         }
     131            0 :     }
     132              : }
     133              : 
     134              : #[derive(Debug, thiserror::Error)]
     135              : pub enum RelationError {
     136              :     #[error("Relation Already Exists")]
     137              :     AlreadyExists,
     138              :     #[error("invalid relnode")]
     139              :     InvalidRelnode,
     140              :     #[error(transparent)]
     141              :     Other(#[from] anyhow::Error),
     142              : }
     143              : 
     144              : ///
     145              : /// This impl provides all the functionality to store PostgreSQL relations, SLRUs,
     146              : /// and other special kinds of files, in a versioned key-value store. The
     147              : /// Timeline struct provides the key-value store.
     148              : ///
     149              : /// This is a separate impl, so that we can easily include all these functions in a Timeline
     150              : /// implementation, and might be moved into a separate struct later.
     151              : impl Timeline {
     152              :     /// Start ingesting a WAL record, or other atomic modification of
     153              :     /// the timeline.
     154              :     ///
     155              :     /// This provides a transaction-like interface to perform a bunch
     156              :     /// of modifications atomically.
     157              :     ///
     158              :     /// To ingest a WAL record, call begin_modification(lsn) to get a
     159              :     /// DatadirModification object. Use the functions in the object to
     160              :     /// modify the repository state, updating all the pages and metadata
     161              :     /// that the WAL record affects. When you're done, call commit() to
     162              :     /// commit the changes.
     163              :     ///
     164              :     /// Lsn stored in modification is advanced by `ingest_record` and
     165              :     /// is used by `commit()` to update `last_record_lsn`.
     166              :     ///
     167              :     /// Calling commit() will flush all the changes and reset the state,
     168              :     /// so the `DatadirModification` struct can be reused to perform the next modification.
     169              :     ///
     170              :     /// Note that any pending modifications you make through the
     171              :     /// modification object won't be visible to calls to the 'get' and list
     172              :     /// functions of the timeline until you finish! And if you update the
     173              :     /// same page twice, the last update wins.
     174              :     ///
     175       268380 :     pub fn begin_modification(&self, lsn: Lsn) -> DatadirModification
     176       268380 :     where
     177       268380 :         Self: Sized,
     178       268380 :     {
     179       268380 :         DatadirModification {
     180       268380 :             tline: self,
     181       268380 :             pending_lsns: Vec::new(),
     182       268380 :             pending_metadata_pages: HashMap::new(),
     183       268380 :             pending_data_batch: None,
     184       268380 :             pending_deletions: Vec::new(),
     185       268380 :             pending_nblocks: 0,
     186       268380 :             pending_directory_entries: Vec::new(),
     187       268380 :             pending_metadata_bytes: 0,
     188       268380 :             lsn,
     189       268380 :         }
     190       268380 :     }
     191              : 
     192              :     //------------------------------------------------------------------------------
     193              :     // Public GET functions
     194              :     //------------------------------------------------------------------------------
     195              : 
     196              :     /// Look up given page version.
     197        18384 :     pub(crate) async fn get_rel_page_at_lsn(
     198        18384 :         &self,
     199        18384 :         tag: RelTag,
     200        18384 :         blknum: BlockNumber,
     201        18384 :         version: Version<'_>,
     202        18384 :         ctx: &RequestContext,
     203        18384 :     ) -> Result<Bytes, PageReconstructError> {
     204        18384 :         match version {
     205        18384 :             Version::Lsn(effective_lsn) => {
     206        18384 :                 let pages: smallvec::SmallVec<[_; 1]> = smallvec::smallvec![(tag, blknum)];
     207        18384 :                 let res = self
     208        18384 :                     .get_rel_page_at_lsn_batched(
     209        18384 :                         pages.iter().map(|(tag, blknum)| (tag, blknum)),
     210        18384 :                         effective_lsn,
     211        18384 :                         ctx,
     212        18384 :                     )
     213        18384 :                     .await;
     214        18384 :                 assert_eq!(res.len(), 1);
     215        18384 :                 res.into_iter().next().unwrap()
     216              :             }
     217            0 :             Version::Modified(modification) => {
     218            0 :                 if tag.relnode == 0 {
     219            0 :                     return Err(PageReconstructError::Other(
     220            0 :                         RelationError::InvalidRelnode.into(),
     221            0 :                     ));
     222            0 :                 }
     223              : 
     224            0 :                 let nblocks = self.get_rel_size(tag, version, ctx).await?;
     225            0 :                 if blknum >= nblocks {
     226            0 :                     debug!(
     227            0 :                         "read beyond EOF at {} blk {} at {}, size is {}: returning all-zeros page",
     228            0 :                         tag,
     229            0 :                         blknum,
     230            0 :                         version.get_lsn(),
     231              :                         nblocks
     232              :                     );
     233            0 :                     return Ok(ZERO_PAGE.clone());
     234            0 :                 }
     235            0 : 
     236            0 :                 let key = rel_block_to_key(tag, blknum);
     237            0 :                 modification.get(key, ctx).await
     238              :             }
     239              :         }
     240        18384 :     }
     241              : 
     242              :     /// Like [`Self::get_rel_page_at_lsn`], but returns a batch of pages.
     243              :     ///
     244              :     /// The ordering of the returned vec corresponds to the ordering of `pages`.
     245        18384 :     pub(crate) async fn get_rel_page_at_lsn_batched(
     246        18384 :         &self,
     247        18384 :         pages: impl ExactSizeIterator<Item = (&RelTag, &BlockNumber)>,
     248        18384 :         effective_lsn: Lsn,
     249        18384 :         ctx: &RequestContext,
     250        18384 :     ) -> Vec<Result<Bytes, PageReconstructError>> {
     251        18384 :         debug_assert_current_span_has_tenant_and_timeline_id();
     252        18384 : 
     253        18384 :         let mut slots_filled = 0;
     254        18384 :         let page_count = pages.len();
     255        18384 : 
     256        18384 :         // Would be nice to use smallvec here but it doesn't provide the spare_capacity_mut() API.
     257        18384 :         let mut result = Vec::with_capacity(pages.len());
     258        18384 :         let result_slots = result.spare_capacity_mut();
     259        18384 : 
     260        18384 :         let mut keys_slots: BTreeMap<Key, smallvec::SmallVec<[usize; 1]>> = BTreeMap::default();
     261        18384 :         for (response_slot_idx, (tag, blknum)) in pages.enumerate() {
     262        18384 :             if tag.relnode == 0 {
     263            0 :                 result_slots[response_slot_idx].write(Err(PageReconstructError::Other(
     264            0 :                     RelationError::InvalidRelnode.into(),
     265            0 :                 )));
     266            0 : 
     267            0 :                 slots_filled += 1;
     268            0 :                 continue;
     269        18384 :             }
     270              : 
     271        18384 :             let nblocks = match self
     272        18384 :                 .get_rel_size(*tag, Version::Lsn(effective_lsn), ctx)
     273        18384 :                 .await
     274              :             {
     275        18384 :                 Ok(nblocks) => nblocks,
     276            0 :                 Err(err) => {
     277            0 :                     result_slots[response_slot_idx].write(Err(err));
     278            0 :                     slots_filled += 1;
     279            0 :                     continue;
     280              :                 }
     281              :             };
     282              : 
     283        18384 :             if *blknum >= nblocks {
     284            0 :                 debug!(
     285            0 :                     "read beyond EOF at {} blk {} at {}, size is {}: returning all-zeros page",
     286              :                     tag, blknum, effective_lsn, nblocks
     287              :                 );
     288            0 :                 result_slots[response_slot_idx].write(Ok(ZERO_PAGE.clone()));
     289            0 :                 slots_filled += 1;
     290            0 :                 continue;
     291        18384 :             }
     292        18384 : 
     293        18384 :             let key = rel_block_to_key(*tag, *blknum);
     294        18384 : 
     295        18384 :             let key_slots = keys_slots.entry(key).or_default();
     296        18384 :             key_slots.push(response_slot_idx);
     297              :         }
     298              : 
     299        18384 :         let keyspace = {
     300              :             // add_key requires monotonicity
     301        18384 :             let mut acc = KeySpaceAccum::new();
     302        18384 :             for key in keys_slots
     303        18384 :                 .keys()
     304        18384 :                 // in fact it requires strong monotonicity
     305        18384 :                 .dedup()
     306        18384 :             {
     307        18384 :                 acc.add_key(*key);
     308        18384 :             }
     309        18384 :             acc.to_keyspace()
     310        18384 :         };
     311        18384 : 
     312        18384 :         match self.get_vectored(keyspace, effective_lsn, ctx).await {
     313        18384 :             Ok(results) => {
     314        36768 :                 for (key, res) in results {
     315        18384 :                     let mut key_slots = keys_slots.remove(&key).unwrap().into_iter();
     316        18384 :                     let first_slot = key_slots.next().unwrap();
     317              : 
     318        18384 :                     for slot in key_slots {
     319            0 :                         let clone = match &res {
     320            0 :                             Ok(buf) => Ok(buf.clone()),
     321            0 :                             Err(err) => Err(match err {
     322              :                                 PageReconstructError::Cancelled => {
     323            0 :                                     PageReconstructError::Cancelled
     324              :                                 }
     325              : 
     326            0 :                                 x @ PageReconstructError::Other(_) |
     327            0 :                                 x @ PageReconstructError::AncestorLsnTimeout(_) |
     328            0 :                                 x @ PageReconstructError::WalRedo(_) |
     329            0 :                                 x @ PageReconstructError::MissingKey(_) => {
     330            0 :                                     PageReconstructError::Other(anyhow::anyhow!("there was more than one request for this key in the batch, error logged once: {x:?}"))
     331              :                                 },
     332              :                             }),
     333              :                         };
     334              : 
     335            0 :                         result_slots[slot].write(clone);
     336            0 :                         slots_filled += 1;
     337              :                     }
     338              : 
     339        18384 :                     result_slots[first_slot].write(res);
     340        18384 :                     slots_filled += 1;
     341              :                 }
     342              :             }
     343            0 :             Err(err) => {
     344              :                 // this cannot really happen because get_vectored only errors globally on invalid LSN or too large batch size
     345              :                 // (We enforce the max batch size outside of this function, in the code that constructs the batch request.)
     346            0 :                 for slot in keys_slots.values().flatten() {
     347              :                     // this whole `match` is a lot like `From<GetVectoredError> for PageReconstructError`
     348              :                     // but without taking ownership of the GetVectoredError
     349            0 :                     let err = match &err {
     350              :                         GetVectoredError::Cancelled => {
     351            0 :                             Err(PageReconstructError::Cancelled)
     352              :                         }
     353              :                         // TODO: restructure get_vectored API to make this error per-key
     354            0 :                         GetVectoredError::MissingKey(err) => {
     355            0 :                             Err(PageReconstructError::Other(anyhow::anyhow!("whole vectored get request failed because one or more of the requested keys were missing: {err:?}")))
     356              :                         }
     357              :                         // TODO: restructure get_vectored API to make this error per-key
     358            0 :                         GetVectoredError::GetReadyAncestorError(err) => {
     359            0 :                             Err(PageReconstructError::Other(anyhow::anyhow!("whole vectored get request failed because one or more key required ancestor that wasn't ready: {err:?}")))
     360              :                         }
     361              :                         // TODO: restructure get_vectored API to make this error per-key
     362            0 :                         GetVectoredError::Other(err) => {
     363            0 :                             Err(PageReconstructError::Other(
     364            0 :                                 anyhow::anyhow!("whole vectored get request failed: {err:?}"),
     365            0 :                             ))
     366              :                         }
     367              :                         // TODO: we can prevent this error class by moving this check into the type system
     368            0 :                         GetVectoredError::InvalidLsn(e) => {
     369            0 :                             Err(anyhow::anyhow!("invalid LSN: {e:?}").into())
     370              :                         }
     371              :                         // NB: this should never happen in practice because we limit MAX_GET_VECTORED_KEYS
     372              :                         // TODO: we can prevent this error class by moving this check into the type system
     373            0 :                         GetVectoredError::Oversized(err) => {
     374            0 :                             Err(anyhow::anyhow!(
     375            0 :                                 "batching oversized: {err:?}"
     376            0 :                             )
     377            0 :                             .into())
     378              :                         }
     379              :                     };
     380              : 
     381            0 :                     result_slots[*slot].write(err);
     382              :                 }
     383              : 
     384            0 :                 slots_filled += keys_slots.values().map(|slots| slots.len()).sum::<usize>();
     385            0 :             }
     386              :         };
     387              : 
     388        18384 :         assert_eq!(slots_filled, page_count);
     389              :         // SAFETY:
     390              :         // 1. `result` and any of its uninint members are not read from until this point
     391              :         // 2. The length below is tracked at run-time and matches the number of requested pages.
     392        18384 :         unsafe {
     393        18384 :             result.set_len(page_count);
     394        18384 :         }
     395        18384 : 
     396        18384 :         result
     397        18384 :     }
     398              : 
     399              :     /// Get size of a database in blocks. This is only accurate on shard 0. It will undercount on
     400              :     /// other shards, by only accounting for relations the shard has pages for, and only accounting
     401              :     /// for pages up to the highest page number it has stored.
     402            0 :     pub(crate) async fn get_db_size(
     403            0 :         &self,
     404            0 :         spcnode: Oid,
     405            0 :         dbnode: Oid,
     406            0 :         version: Version<'_>,
     407            0 :         ctx: &RequestContext,
     408            0 :     ) -> Result<usize, PageReconstructError> {
     409            0 :         let mut total_blocks = 0;
     410              : 
     411            0 :         let rels = self.list_rels(spcnode, dbnode, version, ctx).await?;
     412              : 
     413            0 :         for rel in rels {
     414            0 :             let n_blocks = self.get_rel_size(rel, version, ctx).await?;
     415            0 :             total_blocks += n_blocks as usize;
     416              :         }
     417            0 :         Ok(total_blocks)
     418            0 :     }
     419              : 
     420              :     /// Get size of a relation file. The relation must exist, otherwise an error is returned.
     421              :     ///
     422              :     /// This is only accurate on shard 0. On other shards, it will return the size up to the highest
     423              :     /// page number stored in the shard.
     424        24434 :     pub(crate) async fn get_rel_size(
     425        24434 :         &self,
     426        24434 :         tag: RelTag,
     427        24434 :         version: Version<'_>,
     428        24434 :         ctx: &RequestContext,
     429        24434 :     ) -> Result<BlockNumber, PageReconstructError> {
     430        24434 :         if tag.relnode == 0 {
     431            0 :             return Err(PageReconstructError::Other(
     432            0 :                 RelationError::InvalidRelnode.into(),
     433            0 :             ));
     434        24434 :         }
     435              : 
     436        24434 :         if let Some(nblocks) = self.get_cached_rel_size(&tag, version.get_lsn()) {
     437        19294 :             return Ok(nblocks);
     438         5140 :         }
     439         5140 : 
     440         5140 :         if (tag.forknum == FSM_FORKNUM || tag.forknum == VISIBILITYMAP_FORKNUM)
     441            0 :             && !self.get_rel_exists(tag, version, ctx).await?
     442              :         {
     443              :             // FIXME: Postgres sometimes calls smgrcreate() to create
     444              :             // FSM, and smgrnblocks() on it immediately afterwards,
     445              :             // without extending it.  Tolerate that by claiming that
     446              :             // any non-existent FSM fork has size 0.
     447            0 :             return Ok(0);
     448         5140 :         }
     449         5140 : 
     450         5140 :         let key = rel_size_to_key(tag);
     451         5140 :         let mut buf = version.get(self, key, ctx).await?;
     452         5136 :         let nblocks = buf.get_u32_le();
     453         5136 : 
     454         5136 :         self.update_cached_rel_size(tag, version.get_lsn(), nblocks);
     455         5136 : 
     456         5136 :         Ok(nblocks)
     457        24434 :     }
     458              : 
     459              :     /// Does the relation exist?
     460              :     ///
     461              :     /// Only shard 0 has a full view of the relations. Other shards only know about relations that
     462              :     /// the shard stores pages for.
     463         6050 :     pub(crate) async fn get_rel_exists(
     464         6050 :         &self,
     465         6050 :         tag: RelTag,
     466         6050 :         version: Version<'_>,
     467         6050 :         ctx: &RequestContext,
     468         6050 :     ) -> Result<bool, PageReconstructError> {
     469         6050 :         if tag.relnode == 0 {
     470            0 :             return Err(PageReconstructError::Other(
     471            0 :                 RelationError::InvalidRelnode.into(),
     472            0 :             ));
     473         6050 :         }
     474              : 
     475              :         // first try to lookup relation in cache
     476         6050 :         if let Some(_nblocks) = self.get_cached_rel_size(&tag, version.get_lsn()) {
     477         6032 :             return Ok(true);
     478           18 :         }
     479              :         // then check if the database was already initialized.
     480              :         // get_rel_exists can be called before dbdir is created.
     481           18 :         let buf = version.get(self, DBDIR_KEY, ctx).await?;
     482           18 :         let dbdirs = DbDirectory::des(&buf)?.dbdirs;
     483           18 :         if !dbdirs.contains_key(&(tag.spcnode, tag.dbnode)) {
     484            0 :             return Ok(false);
     485           18 :         }
     486           18 :         // fetch directory listing
     487           18 :         let key = rel_dir_to_key(tag.spcnode, tag.dbnode);
     488           18 :         let buf = version.get(self, key, ctx).await?;
     489              : 
     490           18 :         let dir = RelDirectory::des(&buf)?;
     491           18 :         Ok(dir.rels.contains(&(tag.relnode, tag.forknum)))
     492         6050 :     }
     493              : 
     494              :     /// Get a list of all existing relations in given tablespace and database.
     495              :     ///
     496              :     /// Only shard 0 has a full view of the relations. Other shards only know about relations that
     497              :     /// the shard stores pages for.
     498              :     ///
     499              :     /// # Cancel-Safety
     500              :     ///
     501              :     /// This method is cancellation-safe.
     502            0 :     pub(crate) async fn list_rels(
     503            0 :         &self,
     504            0 :         spcnode: Oid,
     505            0 :         dbnode: Oid,
     506            0 :         version: Version<'_>,
     507            0 :         ctx: &RequestContext,
     508            0 :     ) -> Result<HashSet<RelTag>, PageReconstructError> {
     509            0 :         // fetch directory listing
     510            0 :         let key = rel_dir_to_key(spcnode, dbnode);
     511            0 :         let buf = version.get(self, key, ctx).await?;
     512              : 
     513            0 :         let dir = RelDirectory::des(&buf)?;
     514            0 :         let rels: HashSet<RelTag> =
     515            0 :             HashSet::from_iter(dir.rels.iter().map(|(relnode, forknum)| RelTag {
     516            0 :                 spcnode,
     517            0 :                 dbnode,
     518            0 :                 relnode: *relnode,
     519            0 :                 forknum: *forknum,
     520            0 :             }));
     521            0 : 
     522            0 :         Ok(rels)
     523            0 :     }
     524              : 
     525              :     /// Get the whole SLRU segment
     526            0 :     pub(crate) async fn get_slru_segment(
     527            0 :         &self,
     528            0 :         kind: SlruKind,
     529            0 :         segno: u32,
     530            0 :         lsn: Lsn,
     531            0 :         ctx: &RequestContext,
     532            0 :     ) -> Result<Bytes, PageReconstructError> {
     533            0 :         assert!(self.tenant_shard_id.is_shard_zero());
     534            0 :         let n_blocks = self
     535            0 :             .get_slru_segment_size(kind, segno, Version::Lsn(lsn), ctx)
     536            0 :             .await?;
     537            0 :         let mut segment = BytesMut::with_capacity(n_blocks as usize * BLCKSZ as usize);
     538            0 :         for blkno in 0..n_blocks {
     539            0 :             let block = self
     540            0 :                 .get_slru_page_at_lsn(kind, segno, blkno, lsn, ctx)
     541            0 :                 .await?;
     542            0 :             segment.extend_from_slice(&block[..BLCKSZ as usize]);
     543              :         }
     544            0 :         Ok(segment.freeze())
     545            0 :     }
     546              : 
     547              :     /// Look up given SLRU page version.
     548            0 :     pub(crate) async fn get_slru_page_at_lsn(
     549            0 :         &self,
     550            0 :         kind: SlruKind,
     551            0 :         segno: u32,
     552            0 :         blknum: BlockNumber,
     553            0 :         lsn: Lsn,
     554            0 :         ctx: &RequestContext,
     555            0 :     ) -> Result<Bytes, PageReconstructError> {
     556            0 :         assert!(self.tenant_shard_id.is_shard_zero());
     557            0 :         let key = slru_block_to_key(kind, segno, blknum);
     558            0 :         self.get(key, lsn, ctx).await
     559            0 :     }
     560              : 
     561              :     /// Get size of an SLRU segment
     562            0 :     pub(crate) async fn get_slru_segment_size(
     563            0 :         &self,
     564            0 :         kind: SlruKind,
     565            0 :         segno: u32,
     566            0 :         version: Version<'_>,
     567            0 :         ctx: &RequestContext,
     568            0 :     ) -> Result<BlockNumber, PageReconstructError> {
     569            0 :         assert!(self.tenant_shard_id.is_shard_zero());
     570            0 :         let key = slru_segment_size_to_key(kind, segno);
     571            0 :         let mut buf = version.get(self, key, ctx).await?;
     572            0 :         Ok(buf.get_u32_le())
     573            0 :     }
     574              : 
     575              :     /// Get size of an SLRU segment
     576            0 :     pub(crate) async fn get_slru_segment_exists(
     577            0 :         &self,
     578            0 :         kind: SlruKind,
     579            0 :         segno: u32,
     580            0 :         version: Version<'_>,
     581            0 :         ctx: &RequestContext,
     582            0 :     ) -> Result<bool, PageReconstructError> {
     583            0 :         assert!(self.tenant_shard_id.is_shard_zero());
     584              :         // fetch directory listing
     585            0 :         let key = slru_dir_to_key(kind);
     586            0 :         let buf = version.get(self, key, ctx).await?;
     587              : 
     588            0 :         let dir = SlruSegmentDirectory::des(&buf)?;
     589            0 :         Ok(dir.segments.contains(&segno))
     590            0 :     }
     591              : 
     592              :     /// Locate LSN, such that all transactions that committed before
     593              :     /// 'search_timestamp' are visible, but nothing newer is.
     594              :     ///
     595              :     /// This is not exact. Commit timestamps are not guaranteed to be ordered,
     596              :     /// so it's not well defined which LSN you get if there were multiple commits
     597              :     /// "in flight" at that point in time.
     598              :     ///
     599            0 :     pub(crate) async fn find_lsn_for_timestamp(
     600            0 :         &self,
     601            0 :         search_timestamp: TimestampTz,
     602            0 :         cancel: &CancellationToken,
     603            0 :         ctx: &RequestContext,
     604            0 :     ) -> Result<LsnForTimestamp, PageReconstructError> {
     605            0 :         pausable_failpoint!("find-lsn-for-timestamp-pausable");
     606              : 
     607            0 :         let gc_cutoff_lsn_guard = self.get_latest_gc_cutoff_lsn();
     608            0 :         // We use this method to figure out the branching LSN for the new branch, but the
     609            0 :         // GC cutoff could be before the branching point and we cannot create a new branch
     610            0 :         // with LSN < `ancestor_lsn`. Thus, pick the maximum of these two to be
     611            0 :         // on the safe side.
     612            0 :         let min_lsn = std::cmp::max(*gc_cutoff_lsn_guard, self.get_ancestor_lsn());
     613            0 :         let max_lsn = self.get_last_record_lsn();
     614            0 : 
     615            0 :         // LSNs are always 8-byte aligned. low/mid/high represent the
     616            0 :         // LSN divided by 8.
     617            0 :         let mut low = min_lsn.0 / 8;
     618            0 :         let mut high = max_lsn.0 / 8 + 1;
     619            0 : 
     620            0 :         let mut found_smaller = false;
     621            0 :         let mut found_larger = false;
     622              : 
     623            0 :         while low < high {
     624            0 :             if cancel.is_cancelled() {
     625            0 :                 return Err(PageReconstructError::Cancelled);
     626            0 :             }
     627            0 :             // cannot overflow, high and low are both smaller than u64::MAX / 2
     628            0 :             let mid = (high + low) / 2;
     629              : 
     630            0 :             let cmp = self
     631            0 :                 .is_latest_commit_timestamp_ge_than(
     632            0 :                     search_timestamp,
     633            0 :                     Lsn(mid * 8),
     634            0 :                     &mut found_smaller,
     635            0 :                     &mut found_larger,
     636            0 :                     ctx,
     637            0 :                 )
     638            0 :                 .await?;
     639              : 
     640            0 :             if cmp {
     641            0 :                 high = mid;
     642            0 :             } else {
     643            0 :                 low = mid + 1;
     644            0 :             }
     645              :         }
     646              :         // If `found_smaller == true`, `low = t + 1` where `t` is the target LSN,
     647              :         // so the LSN of the last commit record before or at `search_timestamp`.
     648              :         // Remove one from `low` to get `t`.
     649              :         //
     650              :         // FIXME: it would be better to get the LSN of the previous commit.
     651              :         // Otherwise, if you restore to the returned LSN, the database will
     652              :         // include physical changes from later commits that will be marked
     653              :         // as aborted, and will need to be vacuumed away.
     654            0 :         let commit_lsn = Lsn((low - 1) * 8);
     655            0 :         match (found_smaller, found_larger) {
     656              :             (false, false) => {
     657              :                 // This can happen if no commit records have been processed yet, e.g.
     658              :                 // just after importing a cluster.
     659            0 :                 Ok(LsnForTimestamp::NoData(min_lsn))
     660              :             }
     661              :             (false, true) => {
     662              :                 // Didn't find any commit timestamps smaller than the request
     663            0 :                 Ok(LsnForTimestamp::Past(min_lsn))
     664              :             }
     665            0 :             (true, _) if commit_lsn < min_lsn => {
     666            0 :                 // the search above did set found_smaller to true but it never increased the lsn.
     667            0 :                 // Then, low is still the old min_lsn, and the subtraction above gave a value
     668            0 :                 // below the min_lsn. We should never do that.
     669            0 :                 Ok(LsnForTimestamp::Past(min_lsn))
     670              :             }
     671              :             (true, false) => {
     672              :                 // Only found commits with timestamps smaller than the request.
     673              :                 // It's still a valid case for branch creation, return it.
     674              :                 // And `update_gc_info()` ignores LSN for a `LsnForTimestamp::Future`
     675              :                 // case, anyway.
     676            0 :                 Ok(LsnForTimestamp::Future(commit_lsn))
     677              :             }
     678            0 :             (true, true) => Ok(LsnForTimestamp::Present(commit_lsn)),
     679              :         }
     680            0 :     }
     681              : 
     682              :     /// Subroutine of find_lsn_for_timestamp(). Returns true, if there are any
     683              :     /// commits that committed after 'search_timestamp', at LSN 'probe_lsn'.
     684              :     ///
     685              :     /// Additionally, sets 'found_smaller'/'found_Larger, if encounters any commits
     686              :     /// with a smaller/larger timestamp.
     687              :     ///
     688            0 :     pub(crate) async fn is_latest_commit_timestamp_ge_than(
     689            0 :         &self,
     690            0 :         search_timestamp: TimestampTz,
     691            0 :         probe_lsn: Lsn,
     692            0 :         found_smaller: &mut bool,
     693            0 :         found_larger: &mut bool,
     694            0 :         ctx: &RequestContext,
     695            0 :     ) -> Result<bool, PageReconstructError> {
     696            0 :         self.map_all_timestamps(probe_lsn, ctx, |timestamp| {
     697            0 :             if timestamp >= search_timestamp {
     698            0 :                 *found_larger = true;
     699            0 :                 return ControlFlow::Break(true);
     700            0 :             } else {
     701            0 :                 *found_smaller = true;
     702            0 :             }
     703            0 :             ControlFlow::Continue(())
     704            0 :         })
     705            0 :         .await
     706            0 :     }
     707              : 
     708              :     /// Obtain the possible timestamp range for the given lsn.
     709              :     ///
     710              :     /// If the lsn has no timestamps, returns None. returns `(min, max, median)` if it has timestamps.
     711            0 :     pub(crate) async fn get_timestamp_for_lsn(
     712            0 :         &self,
     713            0 :         probe_lsn: Lsn,
     714            0 :         ctx: &RequestContext,
     715            0 :     ) -> Result<Option<TimestampTz>, PageReconstructError> {
     716            0 :         let mut max: Option<TimestampTz> = None;
     717            0 :         self.map_all_timestamps::<()>(probe_lsn, ctx, |timestamp| {
     718            0 :             if let Some(max_prev) = max {
     719            0 :                 max = Some(max_prev.max(timestamp));
     720            0 :             } else {
     721            0 :                 max = Some(timestamp);
     722            0 :             }
     723            0 :             ControlFlow::Continue(())
     724            0 :         })
     725            0 :         .await?;
     726              : 
     727            0 :         Ok(max)
     728            0 :     }
     729              : 
     730              :     /// Runs the given function on all the timestamps for a given lsn
     731              :     ///
     732              :     /// The return value is either given by the closure, or set to the `Default`
     733              :     /// impl's output.
     734            0 :     async fn map_all_timestamps<T: Default>(
     735            0 :         &self,
     736            0 :         probe_lsn: Lsn,
     737            0 :         ctx: &RequestContext,
     738            0 :         mut f: impl FnMut(TimestampTz) -> ControlFlow<T>,
     739            0 :     ) -> Result<T, PageReconstructError> {
     740            0 :         for segno in self
     741            0 :             .list_slru_segments(SlruKind::Clog, Version::Lsn(probe_lsn), ctx)
     742            0 :             .await?
     743              :         {
     744            0 :             let nblocks = self
     745            0 :                 .get_slru_segment_size(SlruKind::Clog, segno, Version::Lsn(probe_lsn), ctx)
     746            0 :                 .await?;
     747            0 :             for blknum in (0..nblocks).rev() {
     748            0 :                 let clog_page = self
     749            0 :                     .get_slru_page_at_lsn(SlruKind::Clog, segno, blknum, probe_lsn, ctx)
     750            0 :                     .await?;
     751              : 
     752            0 :                 if clog_page.len() == BLCKSZ as usize + 8 {
     753            0 :                     let mut timestamp_bytes = [0u8; 8];
     754            0 :                     timestamp_bytes.copy_from_slice(&clog_page[BLCKSZ as usize..]);
     755            0 :                     let timestamp = TimestampTz::from_be_bytes(timestamp_bytes);
     756            0 : 
     757            0 :                     match f(timestamp) {
     758            0 :                         ControlFlow::Break(b) => return Ok(b),
     759            0 :                         ControlFlow::Continue(()) => (),
     760              :                     }
     761            0 :                 }
     762              :             }
     763              :         }
     764            0 :         Ok(Default::default())
     765            0 :     }
     766              : 
     767            0 :     pub(crate) async fn get_slru_keyspace(
     768            0 :         &self,
     769            0 :         version: Version<'_>,
     770            0 :         ctx: &RequestContext,
     771            0 :     ) -> Result<KeySpace, PageReconstructError> {
     772            0 :         let mut accum = KeySpaceAccum::new();
     773              : 
     774            0 :         for kind in SlruKind::iter() {
     775            0 :             let mut segments: Vec<u32> = self
     776            0 :                 .list_slru_segments(kind, version, ctx)
     777            0 :                 .await?
     778            0 :                 .into_iter()
     779            0 :                 .collect();
     780            0 :             segments.sort_unstable();
     781              : 
     782            0 :             for seg in segments {
     783            0 :                 let block_count = self.get_slru_segment_size(kind, seg, version, ctx).await?;
     784              : 
     785            0 :                 accum.add_range(
     786            0 :                     slru_block_to_key(kind, seg, 0)..slru_block_to_key(kind, seg, block_count),
     787            0 :                 );
     788              :             }
     789              :         }
     790              : 
     791            0 :         Ok(accum.to_keyspace())
     792            0 :     }
     793              : 
     794              :     /// Get a list of SLRU segments
     795            0 :     pub(crate) async fn list_slru_segments(
     796            0 :         &self,
     797            0 :         kind: SlruKind,
     798            0 :         version: Version<'_>,
     799            0 :         ctx: &RequestContext,
     800            0 :     ) -> Result<HashSet<u32>, PageReconstructError> {
     801            0 :         // fetch directory entry
     802            0 :         let key = slru_dir_to_key(kind);
     803              : 
     804            0 :         let buf = version.get(self, key, ctx).await?;
     805            0 :         Ok(SlruSegmentDirectory::des(&buf)?.segments)
     806            0 :     }
     807              : 
     808            0 :     pub(crate) async fn get_relmap_file(
     809            0 :         &self,
     810            0 :         spcnode: Oid,
     811            0 :         dbnode: Oid,
     812            0 :         version: Version<'_>,
     813            0 :         ctx: &RequestContext,
     814            0 :     ) -> Result<Bytes, PageReconstructError> {
     815            0 :         let key = relmap_file_key(spcnode, dbnode);
     816              : 
     817            0 :         let buf = version.get(self, key, ctx).await?;
     818            0 :         Ok(buf)
     819            0 :     }
     820              : 
     821          292 :     pub(crate) async fn list_dbdirs(
     822          292 :         &self,
     823          292 :         lsn: Lsn,
     824          292 :         ctx: &RequestContext,
     825          292 :     ) -> Result<HashMap<(Oid, Oid), bool>, PageReconstructError> {
     826              :         // fetch directory entry
     827          292 :         let buf = self.get(DBDIR_KEY, lsn, ctx).await?;
     828              : 
     829          292 :         Ok(DbDirectory::des(&buf)?.dbdirs)
     830          292 :     }
     831              : 
     832            0 :     pub(crate) async fn get_twophase_file(
     833            0 :         &self,
     834            0 :         xid: u64,
     835            0 :         lsn: Lsn,
     836            0 :         ctx: &RequestContext,
     837            0 :     ) -> Result<Bytes, PageReconstructError> {
     838            0 :         let key = twophase_file_key(xid);
     839            0 :         let buf = self.get(key, lsn, ctx).await?;
     840            0 :         Ok(buf)
     841            0 :     }
     842              : 
     843          294 :     pub(crate) async fn list_twophase_files(
     844          294 :         &self,
     845          294 :         lsn: Lsn,
     846          294 :         ctx: &RequestContext,
     847          294 :     ) -> Result<HashSet<u64>, PageReconstructError> {
     848              :         // fetch directory entry
     849          294 :         let buf = self.get(TWOPHASEDIR_KEY, lsn, ctx).await?;
     850              : 
     851          294 :         if self.pg_version >= 17 {
     852            0 :             Ok(TwoPhaseDirectoryV17::des(&buf)?.xids)
     853              :         } else {
     854          294 :             Ok(TwoPhaseDirectory::des(&buf)?
     855              :                 .xids
     856          294 :                 .iter()
     857          294 :                 .map(|x| u64::from(*x))
     858          294 :                 .collect())
     859              :         }
     860          294 :     }
     861              : 
     862            0 :     pub(crate) async fn get_control_file(
     863            0 :         &self,
     864            0 :         lsn: Lsn,
     865            0 :         ctx: &RequestContext,
     866            0 :     ) -> Result<Bytes, PageReconstructError> {
     867            0 :         self.get(CONTROLFILE_KEY, lsn, ctx).await
     868            0 :     }
     869              : 
     870           12 :     pub(crate) async fn get_checkpoint(
     871           12 :         &self,
     872           12 :         lsn: Lsn,
     873           12 :         ctx: &RequestContext,
     874           12 :     ) -> Result<Bytes, PageReconstructError> {
     875           12 :         self.get(CHECKPOINT_KEY, lsn, ctx).await
     876           12 :     }
     877              : 
     878           12 :     async fn list_aux_files_v2(
     879           12 :         &self,
     880           12 :         lsn: Lsn,
     881           12 :         ctx: &RequestContext,
     882           12 :     ) -> Result<HashMap<String, Bytes>, PageReconstructError> {
     883           12 :         let kv = self
     884           12 :             .scan(KeySpace::single(Key::metadata_aux_key_range()), lsn, ctx)
     885           12 :             .await?;
     886           12 :         let mut result = HashMap::new();
     887           12 :         let mut sz = 0;
     888           30 :         for (_, v) in kv {
     889           18 :             let v = v?;
     890           18 :             let v = aux_file::decode_file_value_bytes(&v)
     891           18 :                 .context("value decode")
     892           18 :                 .map_err(PageReconstructError::Other)?;
     893           34 :             for (fname, content) in v {
     894           16 :                 sz += fname.len();
     895           16 :                 sz += content.len();
     896           16 :                 result.insert(fname, content);
     897           16 :             }
     898              :         }
     899           12 :         self.aux_file_size_estimator.on_initial(sz);
     900           12 :         Ok(result)
     901           12 :     }
     902              : 
     903            0 :     pub(crate) async fn trigger_aux_file_size_computation(
     904            0 :         &self,
     905            0 :         lsn: Lsn,
     906            0 :         ctx: &RequestContext,
     907            0 :     ) -> Result<(), PageReconstructError> {
     908            0 :         self.list_aux_files_v2(lsn, ctx).await?;
     909            0 :         Ok(())
     910            0 :     }
     911              : 
     912           12 :     pub(crate) async fn list_aux_files(
     913           12 :         &self,
     914           12 :         lsn: Lsn,
     915           12 :         ctx: &RequestContext,
     916           12 :     ) -> Result<HashMap<String, Bytes>, PageReconstructError> {
     917           12 :         self.list_aux_files_v2(lsn, ctx).await
     918           12 :     }
     919              : 
     920            0 :     pub(crate) async fn get_replorigins(
     921            0 :         &self,
     922            0 :         lsn: Lsn,
     923            0 :         ctx: &RequestContext,
     924            0 :     ) -> Result<HashMap<RepOriginId, Lsn>, PageReconstructError> {
     925            0 :         let kv = self
     926            0 :             .scan(KeySpace::single(repl_origin_key_range()), lsn, ctx)
     927            0 :             .await?;
     928            0 :         let mut result = HashMap::new();
     929            0 :         for (k, v) in kv {
     930            0 :             let v = v?;
     931            0 :             let origin_id = k.field6 as RepOriginId;
     932            0 :             let origin_lsn = Lsn::des(&v).unwrap();
     933            0 :             if origin_lsn != Lsn::INVALID {
     934            0 :                 result.insert(origin_id, origin_lsn);
     935            0 :             }
     936              :         }
     937            0 :         Ok(result)
     938            0 :     }
     939              : 
     940              :     /// Does the same as get_current_logical_size but counted on demand.
     941              :     /// Used to initialize the logical size tracking on startup.
     942              :     ///
     943              :     /// Only relation blocks are counted currently. That excludes metadata,
     944              :     /// SLRUs, twophase files etc.
     945              :     ///
     946              :     /// # Cancel-Safety
     947              :     ///
     948              :     /// This method is cancellation-safe.
     949            0 :     pub(crate) async fn get_current_logical_size_non_incremental(
     950            0 :         &self,
     951            0 :         lsn: Lsn,
     952            0 :         ctx: &RequestContext,
     953            0 :     ) -> Result<u64, CalculateLogicalSizeError> {
     954            0 :         debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id();
     955              : 
     956              :         // Fetch list of database dirs and iterate them
     957            0 :         let buf = self.get(DBDIR_KEY, lsn, ctx).await?;
     958            0 :         let dbdir = DbDirectory::des(&buf)?;
     959              : 
     960            0 :         let mut total_size: u64 = 0;
     961            0 :         for (spcnode, dbnode) in dbdir.dbdirs.keys() {
     962            0 :             for rel in self
     963            0 :                 .list_rels(*spcnode, *dbnode, Version::Lsn(lsn), ctx)
     964            0 :                 .await?
     965              :             {
     966            0 :                 if self.cancel.is_cancelled() {
     967            0 :                     return Err(CalculateLogicalSizeError::Cancelled);
     968            0 :                 }
     969            0 :                 let relsize_key = rel_size_to_key(rel);
     970            0 :                 let mut buf = self.get(relsize_key, lsn, ctx).await?;
     971            0 :                 let relsize = buf.get_u32_le();
     972            0 : 
     973            0 :                 total_size += relsize as u64;
     974              :             }
     975              :         }
     976            0 :         Ok(total_size * BLCKSZ as u64)
     977            0 :     }
     978              : 
     979              :     /// Get a KeySpace that covers all the Keys that are in use at AND below the given LSN. This is only used
     980              :     /// for gc-compaction.
     981              :     ///
     982              :     /// gc-compaction cannot use the same `collect_keyspace` function as the legacy compaction because it
     983              :     /// processes data at multiple LSNs and needs to be aware of the fact that some key ranges might need to
     984              :     /// be kept only for a specific range of LSN.
     985              :     ///
     986              :     /// Consider the case that the user created branches at LSN 10 and 20, where the user created a table A at
     987              :     /// LSN 10 and dropped that table at LSN 20. `collect_keyspace` at LSN 10 will return the key range
     988              :     /// corresponding to that table, while LSN 20 won't. The keyspace info at a single LSN is not enough to
     989              :     /// determine which keys to retain/drop for gc-compaction.
     990              :     ///
     991              :     /// For now, it only drops AUX-v1 keys. But in the future, the function will be extended to return the keyspace
     992              :     /// to be retained for each of the branch LSN.
     993              :     ///
     994              :     /// The return value is (dense keyspace, sparse keyspace).
     995           40 :     pub(crate) async fn collect_gc_compaction_keyspace(
     996           40 :         &self,
     997           40 :     ) -> Result<(KeySpace, SparseKeySpace), CollectKeySpaceError> {
     998           40 :         let metadata_key_begin = Key::metadata_key_range().start;
     999           40 :         let aux_v1_key = AUX_FILES_KEY;
    1000           40 :         let dense_keyspace = KeySpace {
    1001           40 :             ranges: vec![Key::MIN..aux_v1_key, aux_v1_key.next()..metadata_key_begin],
    1002           40 :         };
    1003           40 :         Ok((
    1004           40 :             dense_keyspace,
    1005           40 :             SparseKeySpace(KeySpace::single(Key::metadata_key_range())),
    1006           40 :         ))
    1007           40 :     }
    1008              : 
    1009              :     ///
    1010              :     /// Get a KeySpace that covers all the Keys that are in use at the given LSN.
    1011              :     /// Anything that's not listed maybe removed from the underlying storage (from
    1012              :     /// that LSN forwards).
    1013              :     ///
    1014              :     /// The return value is (dense keyspace, sparse keyspace).
    1015          292 :     pub(crate) async fn collect_keyspace(
    1016          292 :         &self,
    1017          292 :         lsn: Lsn,
    1018          292 :         ctx: &RequestContext,
    1019          292 :     ) -> Result<(KeySpace, SparseKeySpace), CollectKeySpaceError> {
    1020          292 :         // Iterate through key ranges, greedily packing them into partitions
    1021          292 :         let mut result = KeySpaceAccum::new();
    1022          292 : 
    1023          292 :         // The dbdir metadata always exists
    1024          292 :         result.add_key(DBDIR_KEY);
    1025              : 
    1026              :         // Fetch list of database dirs and iterate them
    1027          292 :         let dbdir = self.list_dbdirs(lsn, ctx).await?;
    1028          292 :         let mut dbs: Vec<((Oid, Oid), bool)> = dbdir.into_iter().collect();
    1029          292 : 
    1030          292 :         dbs.sort_unstable_by(|(k_a, _), (k_b, _)| k_a.cmp(k_b));
    1031          292 :         for ((spcnode, dbnode), has_relmap_file) in dbs {
    1032            0 :             if has_relmap_file {
    1033            0 :                 result.add_key(relmap_file_key(spcnode, dbnode));
    1034            0 :             }
    1035            0 :             result.add_key(rel_dir_to_key(spcnode, dbnode));
    1036              : 
    1037            0 :             let mut rels: Vec<RelTag> = self
    1038            0 :                 .list_rels(spcnode, dbnode, Version::Lsn(lsn), ctx)
    1039            0 :                 .await?
    1040            0 :                 .into_iter()
    1041            0 :                 .collect();
    1042            0 :             rels.sort_unstable();
    1043            0 :             for rel in rels {
    1044            0 :                 let relsize_key = rel_size_to_key(rel);
    1045            0 :                 let mut buf = self.get(relsize_key, lsn, ctx).await?;
    1046            0 :                 let relsize = buf.get_u32_le();
    1047            0 : 
    1048            0 :                 result.add_range(rel_block_to_key(rel, 0)..rel_block_to_key(rel, relsize));
    1049            0 :                 result.add_key(relsize_key);
    1050              :             }
    1051              :         }
    1052              : 
    1053              :         // Iterate SLRUs next
    1054          292 :         if self.tenant_shard_id.is_shard_zero() {
    1055          858 :             for kind in [
    1056          286 :                 SlruKind::Clog,
    1057          286 :                 SlruKind::MultiXactMembers,
    1058          286 :                 SlruKind::MultiXactOffsets,
    1059              :             ] {
    1060          858 :                 let slrudir_key = slru_dir_to_key(kind);
    1061          858 :                 result.add_key(slrudir_key);
    1062          858 :                 let buf = self.get(slrudir_key, lsn, ctx).await?;
    1063          858 :                 let dir = SlruSegmentDirectory::des(&buf)?;
    1064          858 :                 let mut segments: Vec<u32> = dir.segments.iter().cloned().collect();
    1065          858 :                 segments.sort_unstable();
    1066          858 :                 for segno in segments {
    1067            0 :                     let segsize_key = slru_segment_size_to_key(kind, segno);
    1068            0 :                     let mut buf = self.get(segsize_key, lsn, ctx).await?;
    1069            0 :                     let segsize = buf.get_u32_le();
    1070            0 : 
    1071            0 :                     result.add_range(
    1072            0 :                         slru_block_to_key(kind, segno, 0)..slru_block_to_key(kind, segno, segsize),
    1073            0 :                     );
    1074            0 :                     result.add_key(segsize_key);
    1075              :                 }
    1076              :             }
    1077            6 :         }
    1078              : 
    1079              :         // Then pg_twophase
    1080          292 :         result.add_key(TWOPHASEDIR_KEY);
    1081              : 
    1082          292 :         let mut xids: Vec<u64> = self
    1083          292 :             .list_twophase_files(lsn, ctx)
    1084          292 :             .await?
    1085          292 :             .iter()
    1086          292 :             .cloned()
    1087          292 :             .collect();
    1088          292 :         xids.sort_unstable();
    1089          292 :         for xid in xids {
    1090            0 :             result.add_key(twophase_file_key(xid));
    1091            0 :         }
    1092              : 
    1093          292 :         result.add_key(CONTROLFILE_KEY);
    1094          292 :         result.add_key(CHECKPOINT_KEY);
    1095          292 : 
    1096          292 :         // Add extra keyspaces in the test cases. Some test cases write keys into the storage without
    1097          292 :         // creating directory keys. These test cases will add such keyspaces into `extra_test_dense_keyspace`
    1098          292 :         // and the keys will not be garbage-colllected.
    1099          292 :         #[cfg(test)]
    1100          292 :         {
    1101          292 :             let guard = self.extra_test_dense_keyspace.load();
    1102          292 :             for kr in &guard.ranges {
    1103            0 :                 result.add_range(kr.clone());
    1104            0 :             }
    1105            0 :         }
    1106            0 : 
    1107          292 :         let dense_keyspace = result.to_keyspace();
    1108          292 :         let sparse_keyspace = SparseKeySpace(KeySpace {
    1109          292 :             ranges: vec![Key::metadata_aux_key_range(), repl_origin_key_range()],
    1110          292 :         });
    1111          292 : 
    1112          292 :         if cfg!(debug_assertions) {
    1113              :             // Verify if the sparse keyspaces are ordered and non-overlapping.
    1114              : 
    1115              :             // We do not use KeySpaceAccum for sparse_keyspace because we want to ensure each
    1116              :             // category of sparse keys are split into their own image/delta files. If there
    1117              :             // are overlapping keyspaces, they will be automatically merged by keyspace accum,
    1118              :             // and we want the developer to keep the keyspaces separated.
    1119              : 
    1120          292 :             let ranges = &sparse_keyspace.0.ranges;
    1121              : 
    1122              :             // TODO: use a single overlaps_with across the codebase
    1123          292 :             fn overlaps_with<T: Ord>(a: &Range<T>, b: &Range<T>) -> bool {
    1124          292 :                 !(a.end <= b.start || b.end <= a.start)
    1125          292 :             }
    1126          584 :             for i in 0..ranges.len() {
    1127          584 :                 for j in 0..i {
    1128          292 :                     if overlaps_with(&ranges[i], &ranges[j]) {
    1129            0 :                         panic!(
    1130            0 :                             "overlapping sparse keyspace: {}..{} and {}..{}",
    1131            0 :                             ranges[i].start, ranges[i].end, ranges[j].start, ranges[j].end
    1132            0 :                         );
    1133          292 :                     }
    1134              :                 }
    1135              :             }
    1136          292 :             for i in 1..ranges.len() {
    1137          292 :                 assert!(
    1138          292 :                     ranges[i - 1].end <= ranges[i].start,
    1139            0 :                     "unordered sparse keyspace: {}..{} and {}..{}",
    1140            0 :                     ranges[i - 1].start,
    1141            0 :                     ranges[i - 1].end,
    1142            0 :                     ranges[i].start,
    1143            0 :                     ranges[i].end
    1144              :                 );
    1145              :             }
    1146            0 :         }
    1147              : 
    1148          292 :         Ok((dense_keyspace, sparse_keyspace))
    1149          292 :     }
    1150              : 
    1151              :     /// Get cached size of relation if it not updated after specified LSN
    1152       448540 :     pub fn get_cached_rel_size(&self, tag: &RelTag, lsn: Lsn) -> Option<BlockNumber> {
    1153       448540 :         let rel_size_cache = self.rel_size_cache.read().unwrap();
    1154       448540 :         if let Some((cached_lsn, nblocks)) = rel_size_cache.map.get(tag) {
    1155       448518 :             if lsn >= *cached_lsn {
    1156       443372 :                 RELSIZE_CACHE_HITS.inc();
    1157       443372 :                 return Some(*nblocks);
    1158         5146 :             }
    1159         5146 :             RELSIZE_CACHE_MISSES_OLD.inc();
    1160           22 :         }
    1161         5168 :         RELSIZE_CACHE_MISSES.inc();
    1162         5168 :         None
    1163       448540 :     }
    1164              : 
    1165              :     /// Update cached relation size if there is no more recent update
    1166         5136 :     pub fn update_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber) {
    1167         5136 :         let mut rel_size_cache = self.rel_size_cache.write().unwrap();
    1168         5136 : 
    1169         5136 :         if lsn < rel_size_cache.complete_as_of {
    1170              :             // Do not cache old values. It's safe to cache the size on read, as long as
    1171              :             // the read was at an LSN since we started the WAL ingestion. Reasoning: we
    1172              :             // never evict values from the cache, so if the relation size changed after
    1173              :             // 'lsn', the new value is already in the cache.
    1174            0 :             return;
    1175         5136 :         }
    1176         5136 : 
    1177         5136 :         match rel_size_cache.map.entry(tag) {
    1178         5136 :             hash_map::Entry::Occupied(mut entry) => {
    1179         5136 :                 let cached_lsn = entry.get_mut();
    1180         5136 :                 if lsn >= cached_lsn.0 {
    1181            0 :                     *cached_lsn = (lsn, nblocks);
    1182         5136 :                 }
    1183              :             }
    1184            0 :             hash_map::Entry::Vacant(entry) => {
    1185            0 :                 entry.insert((lsn, nblocks));
    1186            0 :                 RELSIZE_CACHE_ENTRIES.inc();
    1187            0 :             }
    1188              :         }
    1189         5136 :     }
    1190              : 
    1191              :     /// Store cached relation size
    1192       282720 :     pub fn set_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber) {
    1193       282720 :         let mut rel_size_cache = self.rel_size_cache.write().unwrap();
    1194       282720 :         if rel_size_cache.map.insert(tag, (lsn, nblocks)).is_none() {
    1195         1920 :             RELSIZE_CACHE_ENTRIES.inc();
    1196       280800 :         }
    1197       282720 :     }
    1198              : 
    1199              :     /// Remove cached relation size
    1200            2 :     pub fn remove_cached_rel_size(&self, tag: &RelTag) {
    1201            2 :         let mut rel_size_cache = self.rel_size_cache.write().unwrap();
    1202            2 :         if rel_size_cache.map.remove(tag).is_some() {
    1203            2 :             RELSIZE_CACHE_ENTRIES.dec();
    1204            2 :         }
    1205            2 :     }
    1206              : }
    1207              : 
    1208              : /// DatadirModification represents an operation to ingest an atomic set of
    1209              : /// updates to the repository.
    1210              : ///
    1211              : /// It is created by the 'begin_record' function. It is called for each WAL
    1212              : /// record, so that all the modifications by a one WAL record appear atomic.
    1213              : pub struct DatadirModification<'a> {
    1214              :     /// The timeline this modification applies to. You can access this to
    1215              :     /// read the state, but note that any pending updates are *not* reflected
    1216              :     /// in the state in 'tline' yet.
    1217              :     pub tline: &'a Timeline,
    1218              : 
    1219              :     /// Current LSN of the modification
    1220              :     lsn: Lsn,
    1221              : 
    1222              :     // The modifications are not applied directly to the underlying key-value store.
    1223              :     // The put-functions add the modifications here, and they are flushed to the
    1224              :     // underlying key-value store by the 'finish' function.
    1225              :     pending_lsns: Vec<Lsn>,
    1226              :     pending_deletions: Vec<(Range<Key>, Lsn)>,
    1227              :     pending_nblocks: i64,
    1228              : 
    1229              :     /// Metadata writes, indexed by key so that they can be read from not-yet-committed modifications
    1230              :     /// while ingesting subsequent records. See [`Self::is_data_key`] for the definition of 'metadata'.
    1231              :     pending_metadata_pages: HashMap<CompactKey, Vec<(Lsn, usize, Value)>>,
    1232              : 
    1233              :     /// Data writes, ready to be flushed into an ephemeral layer. See [`Self::is_data_key`] for
    1234              :     /// which keys are stored here.
    1235              :     pending_data_batch: Option<SerializedValueBatch>,
    1236              : 
    1237              :     /// For special "directory" keys that store key-value maps, track the size of the map
    1238              :     /// if it was updated in this modification.
    1239              :     pending_directory_entries: Vec<(DirectoryKind, usize)>,
    1240              : 
    1241              :     /// An **approximation** of how many metadata bytes will be written to the EphemeralFile.
    1242              :     pending_metadata_bytes: usize,
    1243              : }
    1244              : 
    1245              : impl<'a> DatadirModification<'a> {
    1246              :     // When a DatadirModification is committed, we do a monolithic serialization of all its contents.  WAL records can
    1247              :     // contain multiple pages, so the pageserver's record-based batch size isn't sufficient to bound this allocation: we
    1248              :     // additionally specify a limit on how much payload a DatadirModification may contain before it should be committed.
    1249              :     pub(crate) const MAX_PENDING_BYTES: usize = 8 * 1024 * 1024;
    1250              : 
    1251              :     /// Get the current lsn
    1252       418058 :     pub(crate) fn get_lsn(&self) -> Lsn {
    1253       418058 :         self.lsn
    1254       418058 :     }
    1255              : 
    1256            0 :     pub(crate) fn approx_pending_bytes(&self) -> usize {
    1257            0 :         self.pending_data_batch
    1258            0 :             .as_ref()
    1259            0 :             .map_or(0, |b| b.buffer_size())
    1260            0 :             + self.pending_metadata_bytes
    1261            0 :     }
    1262              : 
    1263            0 :     pub(crate) fn has_dirty_data(&self) -> bool {
    1264            0 :         self.pending_data_batch
    1265            0 :             .as_ref()
    1266            0 :             .map_or(false, |b| b.has_data())
    1267            0 :     }
    1268              : 
    1269              :     /// Set the current lsn
    1270       145858 :     pub(crate) fn set_lsn(&mut self, lsn: Lsn) -> anyhow::Result<()> {
    1271       145858 :         ensure!(
    1272       145858 :             lsn >= self.lsn,
    1273            0 :             "setting an older lsn {} than {} is not allowed",
    1274              :             lsn,
    1275              :             self.lsn
    1276              :         );
    1277              : 
    1278       145858 :         if lsn > self.lsn {
    1279       145858 :             self.pending_lsns.push(self.lsn);
    1280       145858 :             self.lsn = lsn;
    1281       145858 :         }
    1282       145858 :         Ok(())
    1283       145858 :     }
    1284              : 
    1285              :     /// In this context, 'metadata' means keys that are only read by the pageserver internally, and 'data' means
    1286              :     /// keys that represent literal blocks that postgres can read.  So data includes relation blocks and
    1287              :     /// SLRU blocks, which are read directly by postgres, and everything else is considered metadata.
    1288              :     ///
    1289              :     /// The distinction is important because data keys are handled on a fast path where dirty writes are
    1290              :     /// not readable until this modification is committed, whereas metadata keys are visible for read
    1291              :     /// via [`Self::get`] as soon as their record has been ingested.
    1292       850432 :     fn is_data_key(key: &Key) -> bool {
    1293       850432 :         key.is_rel_block_key() || key.is_slru_block_key()
    1294       850432 :     }
    1295              : 
    1296              :     /// Initialize a completely new repository.
    1297              :     ///
    1298              :     /// This inserts the directory metadata entries that are assumed to
    1299              :     /// always exist.
    1300          178 :     pub fn init_empty(&mut self) -> anyhow::Result<()> {
    1301          178 :         let buf = DbDirectory::ser(&DbDirectory {
    1302          178 :             dbdirs: HashMap::new(),
    1303          178 :         })?;
    1304          178 :         self.pending_directory_entries.push((DirectoryKind::Db, 0));
    1305          178 :         self.put(DBDIR_KEY, Value::Image(buf.into()));
    1306              : 
    1307          178 :         let buf = if self.tline.pg_version >= 17 {
    1308            0 :             TwoPhaseDirectoryV17::ser(&TwoPhaseDirectoryV17 {
    1309            0 :                 xids: HashSet::new(),
    1310            0 :             })
    1311              :         } else {
    1312          178 :             TwoPhaseDirectory::ser(&TwoPhaseDirectory {
    1313          178 :                 xids: HashSet::new(),
    1314          178 :             })
    1315            0 :         }?;
    1316          178 :         self.pending_directory_entries
    1317          178 :             .push((DirectoryKind::TwoPhase, 0));
    1318          178 :         self.put(TWOPHASEDIR_KEY, Value::Image(buf.into()));
    1319              : 
    1320          178 :         let buf: Bytes = SlruSegmentDirectory::ser(&SlruSegmentDirectory::default())?.into();
    1321          178 :         let empty_dir = Value::Image(buf);
    1322          178 :         self.put(slru_dir_to_key(SlruKind::Clog), empty_dir.clone());
    1323          178 :         self.pending_directory_entries
    1324          178 :             .push((DirectoryKind::SlruSegment(SlruKind::Clog), 0));
    1325          178 :         self.put(
    1326          178 :             slru_dir_to_key(SlruKind::MultiXactMembers),
    1327          178 :             empty_dir.clone(),
    1328          178 :         );
    1329          178 :         self.pending_directory_entries
    1330          178 :             .push((DirectoryKind::SlruSegment(SlruKind::Clog), 0));
    1331          178 :         self.put(slru_dir_to_key(SlruKind::MultiXactOffsets), empty_dir);
    1332          178 :         self.pending_directory_entries
    1333          178 :             .push((DirectoryKind::SlruSegment(SlruKind::MultiXactOffsets), 0));
    1334          178 : 
    1335          178 :         Ok(())
    1336          178 :     }
    1337              : 
    1338              :     #[cfg(test)]
    1339          176 :     pub fn init_empty_test_timeline(&mut self) -> anyhow::Result<()> {
    1340          176 :         self.init_empty()?;
    1341          176 :         self.put_control_file(bytes::Bytes::from_static(
    1342          176 :             b"control_file contents do not matter",
    1343          176 :         ))
    1344          176 :         .context("put_control_file")?;
    1345          176 :         self.put_checkpoint(bytes::Bytes::from_static(
    1346          176 :             b"checkpoint_file contents do not matter",
    1347          176 :         ))
    1348          176 :         .context("put_checkpoint_file")?;
    1349          176 :         Ok(())
    1350          176 :     }
    1351              : 
    1352              :     /// Creates a relation if it is not already present.
    1353              :     /// Returns the current size of the relation
    1354       418056 :     pub(crate) async fn create_relation_if_required(
    1355       418056 :         &mut self,
    1356       418056 :         rel: RelTag,
    1357       418056 :         ctx: &RequestContext,
    1358       418056 :     ) -> Result<u32, PageReconstructError> {
    1359              :         // Get current size and put rel creation if rel doesn't exist
    1360              :         //
    1361              :         // NOTE: we check the cache first even though get_rel_exists and get_rel_size would
    1362              :         //       check the cache too. This is because eagerly checking the cache results in
    1363              :         //       less work overall and 10% better performance. It's more work on cache miss
    1364              :         //       but cache miss is rare.
    1365       418056 :         if let Some(nblocks) = self.tline.get_cached_rel_size(&rel, self.get_lsn()) {
    1366       418046 :             Ok(nblocks)
    1367           10 :         } else if !self
    1368           10 :             .tline
    1369           10 :             .get_rel_exists(rel, Version::Modified(self), ctx)
    1370           10 :             .await?
    1371              :         {
    1372              :             // create it with 0 size initially, the logic below will extend it
    1373           10 :             self.put_rel_creation(rel, 0, ctx)
    1374           10 :                 .await
    1375           10 :                 .context("Relation Error")?;
    1376           10 :             Ok(0)
    1377              :         } else {
    1378            0 :             self.tline
    1379            0 :                 .get_rel_size(rel, Version::Modified(self), ctx)
    1380            0 :                 .await
    1381              :         }
    1382       418056 :     }
    1383              : 
    1384              :     /// Given a block number for a relation (which represents a newly written block),
    1385              :     /// the previous block count of the relation, and the shard info, find the gaps
    1386              :     /// that were created by the newly written block if any.
    1387       145670 :     fn find_gaps(
    1388       145670 :         rel: RelTag,
    1389       145670 :         blkno: u32,
    1390       145670 :         previous_nblocks: u32,
    1391       145670 :         shard: &ShardIdentity,
    1392       145670 :     ) -> Option<KeySpace> {
    1393       145670 :         let mut key = rel_block_to_key(rel, blkno);
    1394       145670 :         let mut gap_accum = None;
    1395              : 
    1396       145670 :         for gap_blkno in previous_nblocks..blkno {
    1397           32 :             key.field6 = gap_blkno;
    1398           32 : 
    1399           32 :             if shard.get_shard_number(&key) != shard.number {
    1400            8 :                 continue;
    1401           24 :             }
    1402           24 : 
    1403           24 :             gap_accum
    1404           24 :                 .get_or_insert_with(KeySpaceAccum::new)
    1405           24 :                 .add_key(key);
    1406              :         }
    1407              : 
    1408       145670 :         gap_accum.map(|accum| accum.to_keyspace())
    1409       145670 :     }
    1410              : 
    1411       145852 :     pub async fn ingest_batch(
    1412       145852 :         &mut self,
    1413       145852 :         mut batch: SerializedValueBatch,
    1414       145852 :         // TODO(vlad): remove this argument and replace the shard check with is_key_local
    1415       145852 :         shard: &ShardIdentity,
    1416       145852 :         ctx: &RequestContext,
    1417       145852 :     ) -> anyhow::Result<()> {
    1418       145852 :         let mut gaps_at_lsns = Vec::default();
    1419              : 
    1420       145852 :         for meta in batch.metadata.iter() {
    1421       145642 :             let (rel, blkno) = Key::from_compact(meta.key()).to_rel_block()?;
    1422       145642 :             let new_nblocks = blkno + 1;
    1423              : 
    1424       145642 :             let old_nblocks = self.create_relation_if_required(rel, ctx).await?;
    1425       145642 :             if new_nblocks > old_nblocks {
    1426         2390 :                 self.put_rel_extend(rel, new_nblocks, ctx).await?;
    1427       143252 :             }
    1428              : 
    1429       145642 :             if let Some(gaps) = Self::find_gaps(rel, blkno, old_nblocks, shard) {
    1430            0 :                 gaps_at_lsns.push((gaps, meta.lsn()));
    1431       145642 :             }
    1432              :         }
    1433              : 
    1434       145852 :         if !gaps_at_lsns.is_empty() {
    1435            0 :             batch.zero_gaps(gaps_at_lsns);
    1436       145852 :         }
    1437              : 
    1438       145852 :         match self.pending_data_batch.as_mut() {
    1439           20 :             Some(pending_batch) => {
    1440           20 :                 pending_batch.extend(batch);
    1441           20 :             }
    1442       145832 :             None if batch.has_data() => {
    1443       145630 :                 self.pending_data_batch = Some(batch);
    1444       145630 :             }
    1445          202 :             None => {
    1446          202 :                 // Nothing to initialize the batch with
    1447          202 :             }
    1448              :         }
    1449              : 
    1450       145852 :         Ok(())
    1451       145852 :     }
    1452              : 
    1453              :     /// Put a new page version that can be constructed from a WAL record
    1454              :     ///
    1455              :     /// NOTE: this will *not* implicitly extend the relation, if the page is beyond the
    1456              :     /// current end-of-file. It's up to the caller to check that the relation size
    1457              :     /// matches the blocks inserted!
    1458           12 :     pub fn put_rel_wal_record(
    1459           12 :         &mut self,
    1460           12 :         rel: RelTag,
    1461           12 :         blknum: BlockNumber,
    1462           12 :         rec: NeonWalRecord,
    1463           12 :     ) -> anyhow::Result<()> {
    1464           12 :         anyhow::ensure!(rel.relnode != 0, RelationError::InvalidRelnode);
    1465           12 :         self.put(rel_block_to_key(rel, blknum), Value::WalRecord(rec));
    1466           12 :         Ok(())
    1467           12 :     }
    1468              : 
    1469              :     // Same, but for an SLRU.
    1470            8 :     pub fn put_slru_wal_record(
    1471            8 :         &mut self,
    1472            8 :         kind: SlruKind,
    1473            8 :         segno: u32,
    1474            8 :         blknum: BlockNumber,
    1475            8 :         rec: NeonWalRecord,
    1476            8 :     ) -> anyhow::Result<()> {
    1477            8 :         if !self.tline.tenant_shard_id.is_shard_zero() {
    1478            0 :             return Ok(());
    1479            8 :         }
    1480            8 : 
    1481            8 :         self.put(
    1482            8 :             slru_block_to_key(kind, segno, blknum),
    1483            8 :             Value::WalRecord(rec),
    1484            8 :         );
    1485            8 :         Ok(())
    1486            8 :     }
    1487              : 
    1488              :     /// Like put_wal_record, but with ready-made image of the page.
    1489       277842 :     pub fn put_rel_page_image(
    1490       277842 :         &mut self,
    1491       277842 :         rel: RelTag,
    1492       277842 :         blknum: BlockNumber,
    1493       277842 :         img: Bytes,
    1494       277842 :     ) -> anyhow::Result<()> {
    1495       277842 :         anyhow::ensure!(rel.relnode != 0, RelationError::InvalidRelnode);
    1496       277842 :         let key = rel_block_to_key(rel, blknum);
    1497       277842 :         if !key.is_valid_key_on_write_path() {
    1498            0 :             anyhow::bail!(
    1499            0 :                 "the request contains data not supported by pageserver at {}",
    1500            0 :                 key
    1501            0 :             );
    1502       277842 :         }
    1503       277842 :         self.put(rel_block_to_key(rel, blknum), Value::Image(img));
    1504       277842 :         Ok(())
    1505       277842 :     }
    1506              : 
    1507            6 :     pub fn put_slru_page_image(
    1508            6 :         &mut self,
    1509            6 :         kind: SlruKind,
    1510            6 :         segno: u32,
    1511            6 :         blknum: BlockNumber,
    1512            6 :         img: Bytes,
    1513            6 :     ) -> anyhow::Result<()> {
    1514            6 :         assert!(self.tline.tenant_shard_id.is_shard_zero());
    1515              : 
    1516            6 :         let key = slru_block_to_key(kind, segno, blknum);
    1517            6 :         if !key.is_valid_key_on_write_path() {
    1518            0 :             anyhow::bail!(
    1519            0 :                 "the request contains data not supported by pageserver at {}",
    1520            0 :                 key
    1521            0 :             );
    1522            6 :         }
    1523            6 :         self.put(key, Value::Image(img));
    1524            6 :         Ok(())
    1525            6 :     }
    1526              : 
    1527         2998 :     pub(crate) fn put_rel_page_image_zero(
    1528         2998 :         &mut self,
    1529         2998 :         rel: RelTag,
    1530         2998 :         blknum: BlockNumber,
    1531         2998 :     ) -> anyhow::Result<()> {
    1532         2998 :         anyhow::ensure!(rel.relnode != 0, RelationError::InvalidRelnode);
    1533         2998 :         let key = rel_block_to_key(rel, blknum);
    1534         2998 :         if !key.is_valid_key_on_write_path() {
    1535            0 :             anyhow::bail!(
    1536            0 :                 "the request contains data not supported by pageserver: {} @ {}",
    1537            0 :                 key,
    1538            0 :                 self.lsn
    1539            0 :             );
    1540         2998 :         }
    1541         2998 : 
    1542         2998 :         let batch = self
    1543         2998 :             .pending_data_batch
    1544         2998 :             .get_or_insert_with(SerializedValueBatch::default);
    1545         2998 : 
    1546         2998 :         batch.put(key.to_compact(), Value::Image(ZERO_PAGE.clone()), self.lsn);
    1547         2998 : 
    1548         2998 :         Ok(())
    1549         2998 :     }
    1550              : 
    1551            0 :     pub(crate) fn put_slru_page_image_zero(
    1552            0 :         &mut self,
    1553            0 :         kind: SlruKind,
    1554            0 :         segno: u32,
    1555            0 :         blknum: BlockNumber,
    1556            0 :     ) -> anyhow::Result<()> {
    1557            0 :         assert!(self.tline.tenant_shard_id.is_shard_zero());
    1558            0 :         let key = slru_block_to_key(kind, segno, blknum);
    1559            0 :         if !key.is_valid_key_on_write_path() {
    1560            0 :             anyhow::bail!(
    1561            0 :                 "the request contains data not supported by pageserver: {} @ {}",
    1562            0 :                 key,
    1563            0 :                 self.lsn
    1564            0 :             );
    1565            0 :         }
    1566            0 : 
    1567            0 :         let batch = self
    1568            0 :             .pending_data_batch
    1569            0 :             .get_or_insert_with(SerializedValueBatch::default);
    1570            0 : 
    1571            0 :         batch.put(key.to_compact(), Value::Image(ZERO_PAGE.clone()), self.lsn);
    1572            0 : 
    1573            0 :         Ok(())
    1574            0 :     }
    1575              : 
    1576              :     /// Store a relmapper file (pg_filenode.map) in the repository
    1577           16 :     pub async fn put_relmap_file(
    1578           16 :         &mut self,
    1579           16 :         spcnode: Oid,
    1580           16 :         dbnode: Oid,
    1581           16 :         img: Bytes,
    1582           16 :         ctx: &RequestContext,
    1583           16 :     ) -> anyhow::Result<()> {
    1584              :         // Add it to the directory (if it doesn't exist already)
    1585           16 :         let buf = self.get(DBDIR_KEY, ctx).await?;
    1586           16 :         let mut dbdir = DbDirectory::des(&buf)?;
    1587              : 
    1588           16 :         let r = dbdir.dbdirs.insert((spcnode, dbnode), true);
    1589           16 :         if r.is_none() || r == Some(false) {
    1590              :             // The dbdir entry didn't exist, or it contained a
    1591              :             // 'false'. The 'insert' call already updated it with
    1592              :             // 'true', now write the updated 'dbdirs' map back.
    1593           16 :             let buf = DbDirectory::ser(&dbdir)?;
    1594           16 :             self.put(DBDIR_KEY, Value::Image(buf.into()));
    1595            0 :         }
    1596           16 :         if r.is_none() {
    1597              :             // Create RelDirectory
    1598            8 :             let buf = RelDirectory::ser(&RelDirectory {
    1599            8 :                 rels: HashSet::new(),
    1600            8 :             })?;
    1601            8 :             self.pending_directory_entries.push((DirectoryKind::Rel, 0));
    1602            8 :             self.put(
    1603            8 :                 rel_dir_to_key(spcnode, dbnode),
    1604            8 :                 Value::Image(Bytes::from(buf)),
    1605            8 :             );
    1606            8 :         }
    1607              : 
    1608           16 :         self.put(relmap_file_key(spcnode, dbnode), Value::Image(img));
    1609           16 :         Ok(())
    1610           16 :     }
    1611              : 
    1612            0 :     pub async fn put_twophase_file(
    1613            0 :         &mut self,
    1614            0 :         xid: u64,
    1615            0 :         img: Bytes,
    1616            0 :         ctx: &RequestContext,
    1617            0 :     ) -> anyhow::Result<()> {
    1618              :         // Add it to the directory entry
    1619            0 :         let dirbuf = self.get(TWOPHASEDIR_KEY, ctx).await?;
    1620            0 :         let newdirbuf = if self.tline.pg_version >= 17 {
    1621            0 :             let mut dir = TwoPhaseDirectoryV17::des(&dirbuf)?;
    1622            0 :             if !dir.xids.insert(xid) {
    1623            0 :                 anyhow::bail!("twophase file for xid {} already exists", xid);
    1624            0 :             }
    1625            0 :             self.pending_directory_entries
    1626            0 :                 .push((DirectoryKind::TwoPhase, dir.xids.len()));
    1627            0 :             Bytes::from(TwoPhaseDirectoryV17::ser(&dir)?)
    1628              :         } else {
    1629            0 :             let xid = xid as u32;
    1630            0 :             let mut dir = TwoPhaseDirectory::des(&dirbuf)?;
    1631            0 :             if !dir.xids.insert(xid) {
    1632            0 :                 anyhow::bail!("twophase file for xid {} already exists", xid);
    1633            0 :             }
    1634            0 :             self.pending_directory_entries
    1635            0 :                 .push((DirectoryKind::TwoPhase, dir.xids.len()));
    1636            0 :             Bytes::from(TwoPhaseDirectory::ser(&dir)?)
    1637              :         };
    1638            0 :         self.put(TWOPHASEDIR_KEY, Value::Image(newdirbuf));
    1639            0 : 
    1640            0 :         self.put(twophase_file_key(xid), Value::Image(img));
    1641            0 :         Ok(())
    1642            0 :     }
    1643              : 
    1644            0 :     pub async fn set_replorigin(
    1645            0 :         &mut self,
    1646            0 :         origin_id: RepOriginId,
    1647            0 :         origin_lsn: Lsn,
    1648            0 :     ) -> anyhow::Result<()> {
    1649            0 :         let key = repl_origin_key(origin_id);
    1650            0 :         self.put(key, Value::Image(origin_lsn.ser().unwrap().into()));
    1651            0 :         Ok(())
    1652            0 :     }
    1653              : 
    1654            0 :     pub async fn drop_replorigin(&mut self, origin_id: RepOriginId) -> anyhow::Result<()> {
    1655            0 :         self.set_replorigin(origin_id, Lsn::INVALID).await
    1656            0 :     }
    1657              : 
    1658          178 :     pub fn put_control_file(&mut self, img: Bytes) -> anyhow::Result<()> {
    1659          178 :         self.put(CONTROLFILE_KEY, Value::Image(img));
    1660          178 :         Ok(())
    1661          178 :     }
    1662              : 
    1663          192 :     pub fn put_checkpoint(&mut self, img: Bytes) -> anyhow::Result<()> {
    1664          192 :         self.put(CHECKPOINT_KEY, Value::Image(img));
    1665          192 :         Ok(())
    1666          192 :     }
    1667              : 
    1668            0 :     pub async fn drop_dbdir(
    1669            0 :         &mut self,
    1670            0 :         spcnode: Oid,
    1671            0 :         dbnode: Oid,
    1672            0 :         ctx: &RequestContext,
    1673            0 :     ) -> anyhow::Result<()> {
    1674            0 :         let total_blocks = self
    1675            0 :             .tline
    1676            0 :             .get_db_size(spcnode, dbnode, Version::Modified(self), ctx)
    1677            0 :             .await?;
    1678              : 
    1679              :         // Remove entry from dbdir
    1680            0 :         let buf = self.get(DBDIR_KEY, ctx).await?;
    1681            0 :         let mut dir = DbDirectory::des(&buf)?;
    1682            0 :         if dir.dbdirs.remove(&(spcnode, dbnode)).is_some() {
    1683            0 :             let buf = DbDirectory::ser(&dir)?;
    1684            0 :             self.pending_directory_entries
    1685            0 :                 .push((DirectoryKind::Db, dir.dbdirs.len()));
    1686            0 :             self.put(DBDIR_KEY, Value::Image(buf.into()));
    1687              :         } else {
    1688            0 :             warn!(
    1689            0 :                 "dropped dbdir for spcnode {} dbnode {} did not exist in db directory",
    1690              :                 spcnode, dbnode
    1691              :             );
    1692              :         }
    1693              : 
    1694              :         // Update logical database size.
    1695            0 :         self.pending_nblocks -= total_blocks as i64;
    1696            0 : 
    1697            0 :         // Delete all relations and metadata files for the spcnode/dnode
    1698            0 :         self.delete(dbdir_key_range(spcnode, dbnode));
    1699            0 :         Ok(())
    1700            0 :     }
    1701              : 
    1702              :     /// Create a relation fork.
    1703              :     ///
    1704              :     /// 'nblocks' is the initial size.
    1705         1920 :     pub async fn put_rel_creation(
    1706         1920 :         &mut self,
    1707         1920 :         rel: RelTag,
    1708         1920 :         nblocks: BlockNumber,
    1709         1920 :         ctx: &RequestContext,
    1710         1920 :     ) -> Result<(), RelationError> {
    1711         1920 :         if rel.relnode == 0 {
    1712            0 :             return Err(RelationError::InvalidRelnode);
    1713         1920 :         }
    1714              :         // It's possible that this is the first rel for this db in this
    1715              :         // tablespace.  Create the reldir entry for it if so.
    1716         1920 :         let mut dbdir = DbDirectory::des(&self.get(DBDIR_KEY, ctx).await.context("read db")?)
    1717         1920 :             .context("deserialize db")?;
    1718         1920 :         let rel_dir_key = rel_dir_to_key(rel.spcnode, rel.dbnode);
    1719         1920 :         let mut rel_dir =
    1720         1920 :             if let hash_map::Entry::Vacant(e) = dbdir.dbdirs.entry((rel.spcnode, rel.dbnode)) {
    1721              :                 // Didn't exist. Update dbdir
    1722            8 :                 e.insert(false);
    1723            8 :                 let buf = DbDirectory::ser(&dbdir).context("serialize db")?;
    1724            8 :                 self.pending_directory_entries
    1725            8 :                     .push((DirectoryKind::Db, dbdir.dbdirs.len()));
    1726            8 :                 self.put(DBDIR_KEY, Value::Image(buf.into()));
    1727            8 : 
    1728            8 :                 // and create the RelDirectory
    1729            8 :                 RelDirectory::default()
    1730              :             } else {
    1731              :                 // reldir already exists, fetch it
    1732         1912 :                 RelDirectory::des(&self.get(rel_dir_key, ctx).await.context("read db")?)
    1733         1912 :                     .context("deserialize db")?
    1734              :             };
    1735              : 
    1736              :         // Add the new relation to the rel directory entry, and write it back
    1737         1920 :         if !rel_dir.rels.insert((rel.relnode, rel.forknum)) {
    1738            0 :             return Err(RelationError::AlreadyExists);
    1739         1920 :         }
    1740         1920 : 
    1741         1920 :         self.pending_directory_entries
    1742         1920 :             .push((DirectoryKind::Rel, rel_dir.rels.len()));
    1743         1920 : 
    1744         1920 :         self.put(
    1745         1920 :             rel_dir_key,
    1746         1920 :             Value::Image(Bytes::from(
    1747         1920 :                 RelDirectory::ser(&rel_dir).context("serialize")?,
    1748              :             )),
    1749              :         );
    1750              : 
    1751              :         // Put size
    1752         1920 :         let size_key = rel_size_to_key(rel);
    1753         1920 :         let buf = nblocks.to_le_bytes();
    1754         1920 :         self.put(size_key, Value::Image(Bytes::from(buf.to_vec())));
    1755         1920 : 
    1756         1920 :         self.pending_nblocks += nblocks as i64;
    1757         1920 : 
    1758         1920 :         // Update relation size cache
    1759         1920 :         self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
    1760         1920 : 
    1761         1920 :         // Even if nblocks > 0, we don't insert any actual blocks here. That's up to the
    1762         1920 :         // caller.
    1763         1920 :         Ok(())
    1764         1920 :     }
    1765              : 
    1766              :     /// Truncate relation
    1767         6012 :     pub async fn put_rel_truncation(
    1768         6012 :         &mut self,
    1769         6012 :         rel: RelTag,
    1770         6012 :         nblocks: BlockNumber,
    1771         6012 :         ctx: &RequestContext,
    1772         6012 :     ) -> anyhow::Result<()> {
    1773         6012 :         anyhow::ensure!(rel.relnode != 0, RelationError::InvalidRelnode);
    1774         6012 :         if self
    1775         6012 :             .tline
    1776         6012 :             .get_rel_exists(rel, Version::Modified(self), ctx)
    1777         6012 :             .await?
    1778              :         {
    1779         6012 :             let size_key = rel_size_to_key(rel);
    1780              :             // Fetch the old size first
    1781         6012 :             let old_size = self.get(size_key, ctx).await?.get_u32_le();
    1782         6012 : 
    1783         6012 :             // Update the entry with the new size.
    1784         6012 :             let buf = nblocks.to_le_bytes();
    1785         6012 :             self.put(size_key, Value::Image(Bytes::from(buf.to_vec())));
    1786         6012 : 
    1787         6012 :             // Update relation size cache
    1788         6012 :             self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
    1789         6012 : 
    1790         6012 :             // Update logical database size.
    1791         6012 :             self.pending_nblocks -= old_size as i64 - nblocks as i64;
    1792            0 :         }
    1793         6012 :         Ok(())
    1794         6012 :     }
    1795              : 
    1796              :     /// Extend relation
    1797              :     /// If new size is smaller, do nothing.
    1798       276680 :     pub async fn put_rel_extend(
    1799       276680 :         &mut self,
    1800       276680 :         rel: RelTag,
    1801       276680 :         nblocks: BlockNumber,
    1802       276680 :         ctx: &RequestContext,
    1803       276680 :     ) -> anyhow::Result<()> {
    1804       276680 :         anyhow::ensure!(rel.relnode != 0, RelationError::InvalidRelnode);
    1805              : 
    1806              :         // Put size
    1807       276680 :         let size_key = rel_size_to_key(rel);
    1808       276680 :         let old_size = self.get(size_key, ctx).await?.get_u32_le();
    1809       276680 : 
    1810       276680 :         // only extend relation here. never decrease the size
    1811       276680 :         if nblocks > old_size {
    1812       274788 :             let buf = nblocks.to_le_bytes();
    1813       274788 :             self.put(size_key, Value::Image(Bytes::from(buf.to_vec())));
    1814       274788 : 
    1815       274788 :             // Update relation size cache
    1816       274788 :             self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
    1817       274788 : 
    1818       274788 :             self.pending_nblocks += nblocks as i64 - old_size as i64;
    1819       274788 :         }
    1820       276680 :         Ok(())
    1821       276680 :     }
    1822              : 
    1823              :     /// Drop some relations
    1824           10 :     pub(crate) async fn put_rel_drops(
    1825           10 :         &mut self,
    1826           10 :         drop_relations: HashMap<(u32, u32), Vec<RelTag>>,
    1827           10 :         ctx: &RequestContext,
    1828           10 :     ) -> anyhow::Result<()> {
    1829           12 :         for ((spc_node, db_node), rel_tags) in drop_relations {
    1830            2 :             let dir_key = rel_dir_to_key(spc_node, db_node);
    1831            2 :             let buf = self.get(dir_key, ctx).await?;
    1832            2 :             let mut dir = RelDirectory::des(&buf)?;
    1833              : 
    1834            2 :             let mut dirty = false;
    1835            4 :             for rel_tag in rel_tags {
    1836            2 :                 if dir.rels.remove(&(rel_tag.relnode, rel_tag.forknum)) {
    1837            2 :                     dirty = true;
    1838            2 : 
    1839            2 :                     // update logical size
    1840            2 :                     let size_key = rel_size_to_key(rel_tag);
    1841            2 :                     let old_size = self.get(size_key, ctx).await?.get_u32_le();
    1842            2 :                     self.pending_nblocks -= old_size as i64;
    1843            2 : 
    1844            2 :                     // Remove entry from relation size cache
    1845            2 :                     self.tline.remove_cached_rel_size(&rel_tag);
    1846            2 : 
    1847            2 :                     // Delete size entry, as well as all blocks
    1848            2 :                     self.delete(rel_key_range(rel_tag));
    1849            0 :                 }
    1850              :             }
    1851              : 
    1852            2 :             if dirty {
    1853            2 :                 self.put(dir_key, Value::Image(Bytes::from(RelDirectory::ser(&dir)?)));
    1854            2 :                 self.pending_directory_entries
    1855            2 :                     .push((DirectoryKind::Rel, dir.rels.len()));
    1856            0 :             }
    1857              :         }
    1858              : 
    1859           10 :         Ok(())
    1860           10 :     }
    1861              : 
    1862            6 :     pub async fn put_slru_segment_creation(
    1863            6 :         &mut self,
    1864            6 :         kind: SlruKind,
    1865            6 :         segno: u32,
    1866            6 :         nblocks: BlockNumber,
    1867            6 :         ctx: &RequestContext,
    1868            6 :     ) -> anyhow::Result<()> {
    1869            6 :         assert!(self.tline.tenant_shard_id.is_shard_zero());
    1870              : 
    1871              :         // Add it to the directory entry
    1872            6 :         let dir_key = slru_dir_to_key(kind);
    1873            6 :         let buf = self.get(dir_key, ctx).await?;
    1874            6 :         let mut dir = SlruSegmentDirectory::des(&buf)?;
    1875              : 
    1876            6 :         if !dir.segments.insert(segno) {
    1877            0 :             anyhow::bail!("slru segment {kind:?}/{segno} already exists");
    1878            6 :         }
    1879            6 :         self.pending_directory_entries
    1880            6 :             .push((DirectoryKind::SlruSegment(kind), dir.segments.len()));
    1881            6 :         self.put(
    1882            6 :             dir_key,
    1883            6 :             Value::Image(Bytes::from(SlruSegmentDirectory::ser(&dir)?)),
    1884              :         );
    1885              : 
    1886              :         // Put size
    1887            6 :         let size_key = slru_segment_size_to_key(kind, segno);
    1888            6 :         let buf = nblocks.to_le_bytes();
    1889            6 :         self.put(size_key, Value::Image(Bytes::from(buf.to_vec())));
    1890            6 : 
    1891            6 :         // even if nblocks > 0, we don't insert any actual blocks here
    1892            6 : 
    1893            6 :         Ok(())
    1894            6 :     }
    1895              : 
    1896              :     /// Extend SLRU segment
    1897            0 :     pub fn put_slru_extend(
    1898            0 :         &mut self,
    1899            0 :         kind: SlruKind,
    1900            0 :         segno: u32,
    1901            0 :         nblocks: BlockNumber,
    1902            0 :     ) -> anyhow::Result<()> {
    1903            0 :         assert!(self.tline.tenant_shard_id.is_shard_zero());
    1904              : 
    1905              :         // Put size
    1906            0 :         let size_key = slru_segment_size_to_key(kind, segno);
    1907            0 :         let buf = nblocks.to_le_bytes();
    1908            0 :         self.put(size_key, Value::Image(Bytes::from(buf.to_vec())));
    1909            0 :         Ok(())
    1910            0 :     }
    1911              : 
    1912              :     /// This method is used for marking truncated SLRU files
    1913            0 :     pub async fn drop_slru_segment(
    1914            0 :         &mut self,
    1915            0 :         kind: SlruKind,
    1916            0 :         segno: u32,
    1917            0 :         ctx: &RequestContext,
    1918            0 :     ) -> anyhow::Result<()> {
    1919            0 :         // Remove it from the directory entry
    1920            0 :         let dir_key = slru_dir_to_key(kind);
    1921            0 :         let buf = self.get(dir_key, ctx).await?;
    1922            0 :         let mut dir = SlruSegmentDirectory::des(&buf)?;
    1923              : 
    1924            0 :         if !dir.segments.remove(&segno) {
    1925            0 :             warn!("slru segment {:?}/{} does not exist", kind, segno);
    1926            0 :         }
    1927            0 :         self.pending_directory_entries
    1928            0 :             .push((DirectoryKind::SlruSegment(kind), dir.segments.len()));
    1929            0 :         self.put(
    1930            0 :             dir_key,
    1931            0 :             Value::Image(Bytes::from(SlruSegmentDirectory::ser(&dir)?)),
    1932              :         );
    1933              : 
    1934              :         // Delete size entry, as well as all blocks
    1935            0 :         self.delete(slru_segment_key_range(kind, segno));
    1936            0 : 
    1937            0 :         Ok(())
    1938            0 :     }
    1939              : 
    1940              :     /// Drop a relmapper file (pg_filenode.map)
    1941            0 :     pub fn drop_relmap_file(&mut self, _spcnode: Oid, _dbnode: Oid) -> anyhow::Result<()> {
    1942            0 :         // TODO
    1943            0 :         Ok(())
    1944            0 :     }
    1945              : 
    1946              :     /// This method is used for marking truncated SLRU files
    1947            0 :     pub async fn drop_twophase_file(
    1948            0 :         &mut self,
    1949            0 :         xid: u64,
    1950            0 :         ctx: &RequestContext,
    1951            0 :     ) -> anyhow::Result<()> {
    1952              :         // Remove it from the directory entry
    1953            0 :         let buf = self.get(TWOPHASEDIR_KEY, ctx).await?;
    1954            0 :         let newdirbuf = if self.tline.pg_version >= 17 {
    1955            0 :             let mut dir = TwoPhaseDirectoryV17::des(&buf)?;
    1956              : 
    1957            0 :             if !dir.xids.remove(&xid) {
    1958            0 :                 warn!("twophase file for xid {} does not exist", xid);
    1959            0 :             }
    1960            0 :             self.pending_directory_entries
    1961            0 :                 .push((DirectoryKind::TwoPhase, dir.xids.len()));
    1962            0 :             Bytes::from(TwoPhaseDirectoryV17::ser(&dir)?)
    1963              :         } else {
    1964            0 :             let xid: u32 = u32::try_from(xid)?;
    1965            0 :             let mut dir = TwoPhaseDirectory::des(&buf)?;
    1966              : 
    1967            0 :             if !dir.xids.remove(&xid) {
    1968            0 :                 warn!("twophase file for xid {} does not exist", xid);
    1969            0 :             }
    1970            0 :             self.pending_directory_entries
    1971            0 :                 .push((DirectoryKind::TwoPhase, dir.xids.len()));
    1972            0 :             Bytes::from(TwoPhaseDirectory::ser(&dir)?)
    1973              :         };
    1974            0 :         self.put(TWOPHASEDIR_KEY, Value::Image(newdirbuf));
    1975            0 : 
    1976            0 :         // Delete it
    1977            0 :         self.delete(twophase_key_range(xid));
    1978            0 : 
    1979            0 :         Ok(())
    1980            0 :     }
    1981              : 
    1982           16 :     pub async fn put_file(
    1983           16 :         &mut self,
    1984           16 :         path: &str,
    1985           16 :         content: &[u8],
    1986           16 :         ctx: &RequestContext,
    1987           16 :     ) -> anyhow::Result<()> {
    1988           16 :         let key = aux_file::encode_aux_file_key(path);
    1989              :         // retrieve the key from the engine
    1990           16 :         let old_val = match self.get(key, ctx).await {
    1991            4 :             Ok(val) => Some(val),
    1992           12 :             Err(PageReconstructError::MissingKey(_)) => None,
    1993            0 :             Err(e) => return Err(e.into()),
    1994              :         };
    1995           16 :         let files: Vec<(&str, &[u8])> = if let Some(ref old_val) = old_val {
    1996            4 :             aux_file::decode_file_value(old_val)?
    1997              :         } else {
    1998           12 :             Vec::new()
    1999              :         };
    2000           16 :         let mut other_files = Vec::with_capacity(files.len());
    2001           16 :         let mut modifying_file = None;
    2002           20 :         for file @ (p, content) in files {
    2003            4 :             if path == p {
    2004            4 :                 assert!(
    2005            4 :                     modifying_file.is_none(),
    2006            0 :                     "duplicated entries found for {}",
    2007              :                     path
    2008              :                 );
    2009            4 :                 modifying_file = Some(content);
    2010            0 :             } else {
    2011            0 :                 other_files.push(file);
    2012            0 :             }
    2013              :         }
    2014           16 :         let mut new_files = other_files;
    2015           16 :         match (modifying_file, content.is_empty()) {
    2016            2 :             (Some(old_content), false) => {
    2017            2 :                 self.tline
    2018            2 :                     .aux_file_size_estimator
    2019            2 :                     .on_update(old_content.len(), content.len());
    2020            2 :                 new_files.push((path, content));
    2021            2 :             }
    2022            2 :             (Some(old_content), true) => {
    2023            2 :                 self.tline
    2024            2 :                     .aux_file_size_estimator
    2025            2 :                     .on_remove(old_content.len());
    2026            2 :                 // not adding the file key to the final `new_files` vec.
    2027            2 :             }
    2028           12 :             (None, false) => {
    2029           12 :                 self.tline.aux_file_size_estimator.on_add(content.len());
    2030           12 :                 new_files.push((path, content));
    2031           12 :             }
    2032            0 :             (None, true) => warn!("removing non-existing aux file: {}", path),
    2033              :         }
    2034           16 :         let new_val = aux_file::encode_file_value(&new_files)?;
    2035           16 :         self.put(key, Value::Image(new_val.into()));
    2036           16 : 
    2037           16 :         Ok(())
    2038           16 :     }
    2039              : 
    2040              :     ///
    2041              :     /// Flush changes accumulated so far to the underlying repository.
    2042              :     ///
    2043              :     /// Usually, changes made in DatadirModification are atomic, but this allows
    2044              :     /// you to flush them to the underlying repository before the final `commit`.
    2045              :     /// That allows to free up the memory used to hold the pending changes.
    2046              :     ///
    2047              :     /// Currently only used during bulk import of a data directory. In that
    2048              :     /// context, breaking the atomicity is OK. If the import is interrupted, the
    2049              :     /// whole import fails and the timeline will be deleted anyway.
    2050              :     /// (Or to be precise, it will be left behind for debugging purposes and
    2051              :     /// ignored, see <https://github.com/neondatabase/neon/pull/1809>)
    2052              :     ///
    2053              :     /// Note: A consequence of flushing the pending operations is that they
    2054              :     /// won't be visible to subsequent operations until `commit`. The function
    2055              :     /// retains all the metadata, but data pages are flushed. That's again OK
    2056              :     /// for bulk import, where you are just loading data pages and won't try to
    2057              :     /// modify the same pages twice.
    2058         1930 :     pub(crate) async fn flush(&mut self, ctx: &RequestContext) -> anyhow::Result<()> {
    2059         1930 :         // Unless we have accumulated a decent amount of changes, it's not worth it
    2060         1930 :         // to scan through the pending_updates list.
    2061         1930 :         let pending_nblocks = self.pending_nblocks;
    2062         1930 :         if pending_nblocks < 10000 {
    2063         1930 :             return Ok(());
    2064            0 :         }
    2065              : 
    2066            0 :         let mut writer = self.tline.writer().await;
    2067              : 
    2068              :         // Flush relation and  SLRU data blocks, keep metadata.
    2069            0 :         if let Some(batch) = self.pending_data_batch.take() {
    2070            0 :             tracing::debug!(
    2071            0 :                 "Flushing batch with max_lsn={}. Last record LSN is {}",
    2072            0 :                 batch.max_lsn,
    2073            0 :                 self.tline.get_last_record_lsn()
    2074              :             );
    2075              : 
    2076              :             // This bails out on first error without modifying pending_updates.
    2077              :             // That's Ok, cf this function's doc comment.
    2078            0 :             writer.put_batch(batch, ctx).await?;
    2079            0 :         }
    2080              : 
    2081            0 :         if pending_nblocks != 0 {
    2082            0 :             writer.update_current_logical_size(pending_nblocks * i64::from(BLCKSZ));
    2083            0 :             self.pending_nblocks = 0;
    2084            0 :         }
    2085              : 
    2086            0 :         for (kind, count) in std::mem::take(&mut self.pending_directory_entries) {
    2087            0 :             writer.update_directory_entries_count(kind, count as u64);
    2088            0 :         }
    2089              : 
    2090            0 :         Ok(())
    2091         1930 :     }
    2092              : 
    2093              :     ///
    2094              :     /// Finish this atomic update, writing all the updated keys to the
    2095              :     /// underlying timeline.
    2096              :     /// All the modifications in this atomic update are stamped by the specified LSN.
    2097              :     ///
    2098       743064 :     pub async fn commit(&mut self, ctx: &RequestContext) -> anyhow::Result<()> {
    2099       743064 :         let mut writer = self.tline.writer().await;
    2100              : 
    2101       743064 :         let pending_nblocks = self.pending_nblocks;
    2102       743064 :         self.pending_nblocks = 0;
    2103              : 
    2104              :         // Ordering: the items in this batch do not need to be in any global order, but values for
    2105              :         // a particular Key must be in Lsn order relative to one another.  InMemoryLayer relies on
    2106              :         // this to do efficient updates to its index.  See [`wal_decoder::serialized_batch`] for
    2107              :         // more details.
    2108              : 
    2109       743064 :         let metadata_batch = {
    2110       743064 :             let pending_meta = self
    2111       743064 :                 .pending_metadata_pages
    2112       743064 :                 .drain()
    2113       743064 :                 .flat_map(|(key, values)| {
    2114       273826 :                     values
    2115       273826 :                         .into_iter()
    2116       273826 :                         .map(move |(lsn, value_size, value)| (key, lsn, value_size, value))
    2117       743064 :                 })
    2118       743064 :                 .collect::<Vec<_>>();
    2119       743064 : 
    2120       743064 :             if pending_meta.is_empty() {
    2121       472278 :                 None
    2122              :             } else {
    2123       270786 :                 Some(SerializedValueBatch::from_values(pending_meta))
    2124              :             }
    2125              :         };
    2126              : 
    2127       743064 :         let data_batch = self.pending_data_batch.take();
    2128              : 
    2129       743064 :         let maybe_batch = match (data_batch, metadata_batch) {
    2130       264556 :             (Some(mut data), Some(metadata)) => {
    2131       264556 :                 data.extend(metadata);
    2132       264556 :                 Some(data)
    2133              :             }
    2134       143262 :             (Some(data), None) => Some(data),
    2135         6230 :             (None, Some(metadata)) => Some(metadata),
    2136       329016 :             (None, None) => None,
    2137              :         };
    2138              : 
    2139       743064 :         if let Some(batch) = maybe_batch {
    2140       414048 :             tracing::debug!(
    2141            0 :                 "Flushing batch with max_lsn={}. Last record LSN is {}",
    2142            0 :                 batch.max_lsn,
    2143            0 :                 self.tline.get_last_record_lsn()
    2144              :             );
    2145              : 
    2146              :             // This bails out on first error without modifying pending_updates.
    2147              :             // That's Ok, cf this function's doc comment.
    2148       414048 :             writer.put_batch(batch, ctx).await?;
    2149       329016 :         }
    2150              : 
    2151       743064 :         if !self.pending_deletions.is_empty() {
    2152            2 :             writer.delete_batch(&self.pending_deletions, ctx).await?;
    2153            2 :             self.pending_deletions.clear();
    2154       743062 :         }
    2155              : 
    2156       743064 :         self.pending_lsns.push(self.lsn);
    2157       888922 :         for pending_lsn in self.pending_lsns.drain(..) {
    2158       888922 :             // TODO(vlad): pretty sure the comment below is not valid anymore
    2159       888922 :             // and we can call finish write with the latest LSN
    2160       888922 :             //
    2161       888922 :             // Ideally, we should be able to call writer.finish_write() only once
    2162       888922 :             // with the highest LSN. However, the last_record_lsn variable in the
    2163       888922 :             // timeline keeps track of the latest LSN and the immediate previous LSN
    2164       888922 :             // so we need to record every LSN to not leave a gap between them.
    2165       888922 :             writer.finish_write(pending_lsn);
    2166       888922 :         }
    2167              : 
    2168       743064 :         if pending_nblocks != 0 {
    2169       270570 :             writer.update_current_logical_size(pending_nblocks * i64::from(BLCKSZ));
    2170       472494 :         }
    2171              : 
    2172       743064 :         for (kind, count) in std::mem::take(&mut self.pending_directory_entries) {
    2173         2834 :             writer.update_directory_entries_count(kind, count as u64);
    2174         2834 :         }
    2175              : 
    2176       743064 :         self.pending_metadata_bytes = 0;
    2177       743064 : 
    2178       743064 :         Ok(())
    2179       743064 :     }
    2180              : 
    2181       291704 :     pub(crate) fn len(&self) -> usize {
    2182       291704 :         self.pending_metadata_pages.len()
    2183       291704 :             + self.pending_data_batch.as_ref().map_or(0, |b| b.len())
    2184       291704 :             + self.pending_deletions.len()
    2185       291704 :     }
    2186              : 
    2187              :     /// Read a page from the Timeline we are writing to.  For metadata pages, this passes through
    2188              :     /// a cache in Self, which makes writes earlier in this modification visible to WAL records later
    2189              :     /// in the modification.
    2190              :     ///
    2191              :     /// For data pages, reads pass directly to the owning Timeline: any ingest code which reads a data
    2192              :     /// page must ensure that the pages they read are already committed in Timeline, for example
    2193              :     /// DB create operations are always preceded by a call to commit().  This is special cased because
    2194              :     /// it's rare: all the 'normal' WAL operations will only read metadata pages such as relation sizes,
    2195              :     /// and not data pages.
    2196       286586 :     async fn get(&self, key: Key, ctx: &RequestContext) -> Result<Bytes, PageReconstructError> {
    2197       286586 :         if !Self::is_data_key(&key) {
    2198              :             // Have we already updated the same key? Read the latest pending updated
    2199              :             // version in that case.
    2200              :             //
    2201              :             // Note: we don't check pending_deletions. It is an error to request a
    2202              :             // value that has been removed, deletion only avoids leaking storage.
    2203       286586 :             if let Some(values) = self.pending_metadata_pages.get(&key.to_compact()) {
    2204        15928 :                 if let Some((_, _, value)) = values.last() {
    2205        15928 :                     return if let Value::Image(img) = value {
    2206        15928 :                         Ok(img.clone())
    2207              :                     } else {
    2208              :                         // Currently, we never need to read back a WAL record that we
    2209              :                         // inserted in the same "transaction". All the metadata updates
    2210              :                         // work directly with Images, and we never need to read actual
    2211              :                         // data pages. We could handle this if we had to, by calling
    2212              :                         // the walredo manager, but let's keep it simple for now.
    2213            0 :                         Err(PageReconstructError::Other(anyhow::anyhow!(
    2214            0 :                             "unexpected pending WAL record"
    2215            0 :                         )))
    2216              :                     };
    2217            0 :                 }
    2218       270658 :             }
    2219              :         } else {
    2220              :             // This is an expensive check, so we only do it in debug mode. If reading a data key,
    2221              :             // this key should never be present in pending_data_pages. We ensure this by committing
    2222              :             // modifications before ingesting DB create operations, which are the only kind that reads
    2223              :             // data pages during ingest.
    2224            0 :             if cfg!(debug_assertions) {
    2225            0 :                 assert!(!self
    2226            0 :                     .pending_data_batch
    2227            0 :                     .as_ref()
    2228            0 :                     .map_or(false, |b| b.updates_key(&key)));
    2229            0 :             }
    2230              :         }
    2231              : 
    2232              :         // Metadata page cache miss, or we're reading a data page.
    2233       270658 :         let lsn = Lsn::max(self.tline.get_last_record_lsn(), self.lsn);
    2234       270658 :         self.tline.get(key, lsn, ctx).await
    2235       286586 :     }
    2236              : 
    2237       563846 :     fn put(&mut self, key: Key, val: Value) {
    2238       563846 :         if Self::is_data_key(&key) {
    2239       277868 :             self.put_data(key.to_compact(), val)
    2240              :         } else {
    2241       285978 :             self.put_metadata(key.to_compact(), val)
    2242              :         }
    2243       563846 :     }
    2244              : 
    2245       277868 :     fn put_data(&mut self, key: CompactKey, val: Value) {
    2246       277868 :         let batch = self
    2247       277868 :             .pending_data_batch
    2248       277868 :             .get_or_insert_with(SerializedValueBatch::default);
    2249       277868 :         batch.put(key, val, self.lsn);
    2250       277868 :     }
    2251              : 
    2252       285978 :     fn put_metadata(&mut self, key: CompactKey, val: Value) {
    2253       285978 :         let values = self.pending_metadata_pages.entry(key).or_default();
    2254              :         // Replace the previous value if it exists at the same lsn
    2255       285978 :         if let Some((last_lsn, last_value_ser_size, last_value)) = values.last_mut() {
    2256        12152 :             if *last_lsn == self.lsn {
    2257              :                 // Update the pending_metadata_bytes contribution from this entry, and update the serialized size in place
    2258        12152 :                 self.pending_metadata_bytes -= *last_value_ser_size;
    2259        12152 :                 *last_value_ser_size = val.serialized_size().unwrap() as usize;
    2260        12152 :                 self.pending_metadata_bytes += *last_value_ser_size;
    2261        12152 : 
    2262        12152 :                 // Use the latest value, this replaces any earlier write to the same (key,lsn), such as much
    2263        12152 :                 // have been generated by synthesized zero page writes prior to the first real write to a page.
    2264        12152 :                 *last_value = val;
    2265        12152 :                 return;
    2266            0 :             }
    2267       273826 :         }
    2268              : 
    2269       273826 :         let val_serialized_size = val.serialized_size().unwrap() as usize;
    2270       273826 :         self.pending_metadata_bytes += val_serialized_size;
    2271       273826 :         values.push((self.lsn, val_serialized_size, val));
    2272       273826 : 
    2273       273826 :         if key == CHECKPOINT_KEY.to_compact() {
    2274          192 :             tracing::debug!("Checkpoint key added to pending with size {val_serialized_size}");
    2275       273634 :         }
    2276       285978 :     }
    2277              : 
    2278            2 :     fn delete(&mut self, key_range: Range<Key>) {
    2279            2 :         trace!("DELETE {}-{}", key_range.start, key_range.end);
    2280            2 :         self.pending_deletions.push((key_range, self.lsn));
    2281            2 :     }
    2282              : }
    2283              : 
    2284              : /// This struct facilitates accessing either a committed key from the timeline at a
    2285              : /// specific LSN, or the latest uncommitted key from a pending modification.
    2286              : ///
    2287              : /// During WAL ingestion, the records from multiple LSNs may be batched in the same
    2288              : /// modification before being flushed to the timeline. Hence, the routines in WalIngest
    2289              : /// need to look up the keys in the modification first before looking them up in the
    2290              : /// timeline to not miss the latest updates.
    2291              : #[derive(Clone, Copy)]
    2292              : pub enum Version<'a> {
    2293              :     Lsn(Lsn),
    2294              :     Modified(&'a DatadirModification<'a>),
    2295              : }
    2296              : 
    2297              : impl<'a> Version<'a> {
    2298         5176 :     async fn get(
    2299         5176 :         &self,
    2300         5176 :         timeline: &Timeline,
    2301         5176 :         key: Key,
    2302         5176 :         ctx: &RequestContext,
    2303         5176 :     ) -> Result<Bytes, PageReconstructError> {
    2304         5176 :         match self {
    2305         5156 :             Version::Lsn(lsn) => timeline.get(key, *lsn, ctx).await,
    2306           20 :             Version::Modified(modification) => modification.get(key, ctx).await,
    2307              :         }
    2308         5176 :     }
    2309              : 
    2310        35620 :     fn get_lsn(&self) -> Lsn {
    2311        35620 :         match self {
    2312        29574 :             Version::Lsn(lsn) => *lsn,
    2313         6046 :             Version::Modified(modification) => modification.lsn,
    2314              :         }
    2315        35620 :     }
    2316              : }
    2317              : 
    2318              : //--- Metadata structs stored in key-value pairs in the repository.
    2319              : 
    2320         2246 : #[derive(Debug, Serialize, Deserialize)]
    2321              : pub(crate) struct DbDirectory {
    2322              :     // (spcnode, dbnode) -> (do relmapper and PG_VERSION files exist)
    2323              :     pub(crate) dbdirs: HashMap<(Oid, Oid), bool>,
    2324              : }
    2325              : 
    2326              : // The format of TwoPhaseDirectory changed in PostgreSQL v17, because the filenames of
    2327              : // pg_twophase files was expanded from 32-bit XIDs to 64-bit XIDs.  Previously, the files
    2328              : // were named like "pg_twophase/000002E5", now they're like
    2329              : // "pg_twophsae/0000000A000002E4".
    2330              : 
    2331          294 : #[derive(Debug, Serialize, Deserialize)]
    2332              : pub(crate) struct TwoPhaseDirectory {
    2333              :     pub(crate) xids: HashSet<TransactionId>,
    2334              : }
    2335              : 
    2336            0 : #[derive(Debug, Serialize, Deserialize)]
    2337              : struct TwoPhaseDirectoryV17 {
    2338              :     xids: HashSet<u64>,
    2339              : }
    2340              : 
    2341         1932 : #[derive(Debug, Serialize, Deserialize, Default)]
    2342              : pub(crate) struct RelDirectory {
    2343              :     // Set of relations that exist. (relfilenode, forknum)
    2344              :     //
    2345              :     // TODO: Store it as a btree or radix tree or something else that spans multiple
    2346              :     // key-value pairs, if you have a lot of relations
    2347              :     pub(crate) rels: HashSet<(Oid, u8)>,
    2348              : }
    2349              : 
    2350            0 : #[derive(Debug, Serialize, Deserialize)]
    2351              : struct RelSizeEntry {
    2352              :     nblocks: u32,
    2353              : }
    2354              : 
    2355          864 : #[derive(Debug, Serialize, Deserialize, Default)]
    2356              : pub(crate) struct SlruSegmentDirectory {
    2357              :     // Set of SLRU segments that exist.
    2358              :     pub(crate) segments: HashSet<u32>,
    2359              : }
    2360              : 
    2361              : #[derive(Copy, Clone, PartialEq, Eq, Debug, enum_map::Enum)]
    2362              : #[repr(u8)]
    2363              : pub(crate) enum DirectoryKind {
    2364              :     Db,
    2365              :     TwoPhase,
    2366              :     Rel,
    2367              :     AuxFiles,
    2368              :     SlruSegment(SlruKind),
    2369              : }
    2370              : 
    2371              : impl DirectoryKind {
    2372              :     pub(crate) const KINDS_NUM: usize = <DirectoryKind as Enum>::LENGTH;
    2373         5668 :     pub(crate) fn offset(&self) -> usize {
    2374         5668 :         self.into_usize()
    2375         5668 :     }
    2376              : }
    2377              : 
    2378              : static ZERO_PAGE: Bytes = Bytes::from_static(&[0u8; BLCKSZ as usize]);
    2379              : 
    2380              : #[allow(clippy::bool_assert_comparison)]
    2381              : #[cfg(test)]
    2382              : mod tests {
    2383              :     use hex_literal::hex;
    2384              :     use pageserver_api::{models::ShardParameters, shard::ShardStripeSize};
    2385              :     use utils::{
    2386              :         id::TimelineId,
    2387              :         shard::{ShardCount, ShardNumber},
    2388              :     };
    2389              : 
    2390              :     use super::*;
    2391              : 
    2392              :     use crate::{tenant::harness::TenantHarness, DEFAULT_PG_VERSION};
    2393              : 
    2394              :     /// Test a round trip of aux file updates, from DatadirModification to reading back from the Timeline
    2395              :     #[tokio::test]
    2396            2 :     async fn aux_files_round_trip() -> anyhow::Result<()> {
    2397            2 :         let name = "aux_files_round_trip";
    2398            2 :         let harness = TenantHarness::create(name).await?;
    2399            2 : 
    2400            2 :         pub const TIMELINE_ID: TimelineId =
    2401            2 :             TimelineId::from_array(hex!("11223344556677881122334455667788"));
    2402            2 : 
    2403            2 :         let (tenant, ctx) = harness.load().await;
    2404            2 :         let tline = tenant
    2405            2 :             .create_empty_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
    2406            2 :             .await?;
    2407            2 :         let tline = tline.raw_timeline().unwrap();
    2408            2 : 
    2409            2 :         // First modification: insert two keys
    2410            2 :         let mut modification = tline.begin_modification(Lsn(0x1000));
    2411            2 :         modification.put_file("foo/bar1", b"content1", &ctx).await?;
    2412            2 :         modification.set_lsn(Lsn(0x1008))?;
    2413            2 :         modification.put_file("foo/bar2", b"content2", &ctx).await?;
    2414            2 :         modification.commit(&ctx).await?;
    2415            2 :         let expect_1008 = HashMap::from([
    2416            2 :             ("foo/bar1".to_string(), Bytes::from_static(b"content1")),
    2417            2 :             ("foo/bar2".to_string(), Bytes::from_static(b"content2")),
    2418            2 :         ]);
    2419            2 : 
    2420            2 :         let readback = tline.list_aux_files(Lsn(0x1008), &ctx).await?;
    2421            2 :         assert_eq!(readback, expect_1008);
    2422            2 : 
    2423            2 :         // Second modification: update one key, remove the other
    2424            2 :         let mut modification = tline.begin_modification(Lsn(0x2000));
    2425            2 :         modification.put_file("foo/bar1", b"content3", &ctx).await?;
    2426            2 :         modification.set_lsn(Lsn(0x2008))?;
    2427            2 :         modification.put_file("foo/bar2", b"", &ctx).await?;
    2428            2 :         modification.commit(&ctx).await?;
    2429            2 :         let expect_2008 =
    2430            2 :             HashMap::from([("foo/bar1".to_string(), Bytes::from_static(b"content3"))]);
    2431            2 : 
    2432            2 :         let readback = tline.list_aux_files(Lsn(0x2008), &ctx).await?;
    2433            2 :         assert_eq!(readback, expect_2008);
    2434            2 : 
    2435            2 :         // Reading back in time works
    2436            2 :         let readback = tline.list_aux_files(Lsn(0x1008), &ctx).await?;
    2437            2 :         assert_eq!(readback, expect_1008);
    2438            2 : 
    2439            2 :         Ok(())
    2440            2 :     }
    2441              : 
    2442              :     #[test]
    2443            2 :     fn gap_finding() {
    2444            2 :         let rel = RelTag {
    2445            2 :             spcnode: 1663,
    2446            2 :             dbnode: 208101,
    2447            2 :             relnode: 2620,
    2448            2 :             forknum: 0,
    2449            2 :         };
    2450            2 :         let base_blkno = 1;
    2451            2 : 
    2452            2 :         let base_key = rel_block_to_key(rel, base_blkno);
    2453            2 :         let before_base_key = rel_block_to_key(rel, base_blkno - 1);
    2454            2 : 
    2455            2 :         let shard = ShardIdentity::unsharded();
    2456            2 : 
    2457            2 :         let mut previous_nblocks = 0;
    2458           22 :         for i in 0..10 {
    2459           20 :             let crnt_blkno = base_blkno + i;
    2460           20 :             let gaps = DatadirModification::find_gaps(rel, crnt_blkno, previous_nblocks, &shard);
    2461           20 : 
    2462           20 :             previous_nblocks = crnt_blkno + 1;
    2463           20 : 
    2464           20 :             if i == 0 {
    2465              :                 // The first block we write is 1, so we should find the gap.
    2466            2 :                 assert_eq!(gaps.unwrap(), KeySpace::single(before_base_key..base_key));
    2467              :             } else {
    2468           18 :                 assert!(gaps.is_none());
    2469              :             }
    2470              :         }
    2471              : 
    2472              :         // This is an update to an already existing block. No gaps here.
    2473            2 :         let update_blkno = 5;
    2474            2 :         let gaps = DatadirModification::find_gaps(rel, update_blkno, previous_nblocks, &shard);
    2475            2 :         assert!(gaps.is_none());
    2476              : 
    2477              :         // This is an update past the current end block.
    2478            2 :         let after_gap_blkno = 20;
    2479            2 :         let gaps = DatadirModification::find_gaps(rel, after_gap_blkno, previous_nblocks, &shard);
    2480            2 : 
    2481            2 :         let gap_start_key = rel_block_to_key(rel, previous_nblocks);
    2482            2 :         let after_gap_key = rel_block_to_key(rel, after_gap_blkno);
    2483            2 :         assert_eq!(
    2484            2 :             gaps.unwrap(),
    2485            2 :             KeySpace::single(gap_start_key..after_gap_key)
    2486            2 :         );
    2487            2 :     }
    2488              : 
    2489              :     #[test]
    2490            2 :     fn sharded_gap_finding() {
    2491            2 :         let rel = RelTag {
    2492            2 :             spcnode: 1663,
    2493            2 :             dbnode: 208101,
    2494            2 :             relnode: 2620,
    2495            2 :             forknum: 0,
    2496            2 :         };
    2497            2 : 
    2498            2 :         let first_blkno = 6;
    2499            2 : 
    2500            2 :         // This shard will get the even blocks
    2501            2 :         let shard = ShardIdentity::from_params(
    2502            2 :             ShardNumber(0),
    2503            2 :             &ShardParameters {
    2504            2 :                 count: ShardCount(2),
    2505            2 :                 stripe_size: ShardStripeSize(1),
    2506            2 :             },
    2507            2 :         );
    2508            2 : 
    2509            2 :         // Only keys belonging to this shard are considered as gaps.
    2510            2 :         let mut previous_nblocks = 0;
    2511            2 :         let gaps =
    2512            2 :             DatadirModification::find_gaps(rel, first_blkno, previous_nblocks, &shard).unwrap();
    2513            2 :         assert!(!gaps.ranges.is_empty());
    2514            6 :         for gap_range in gaps.ranges {
    2515            4 :             let mut k = gap_range.start;
    2516            8 :             while k != gap_range.end {
    2517            4 :                 assert_eq!(shard.get_shard_number(&k), shard.number);
    2518            4 :                 k = k.next();
    2519              :             }
    2520              :         }
    2521              : 
    2522            2 :         previous_nblocks = first_blkno;
    2523            2 : 
    2524            2 :         let update_blkno = 2;
    2525            2 :         let gaps = DatadirModification::find_gaps(rel, update_blkno, previous_nblocks, &shard);
    2526            2 :         assert!(gaps.is_none());
    2527            2 :     }
    2528              : 
    2529              :     /*
    2530              :         fn assert_current_logical_size<R: Repository>(timeline: &DatadirTimeline<R>, lsn: Lsn) {
    2531              :             let incremental = timeline.get_current_logical_size();
    2532              :             let non_incremental = timeline
    2533              :                 .get_current_logical_size_non_incremental(lsn)
    2534              :                 .unwrap();
    2535              :             assert_eq!(incremental, non_incremental);
    2536              :         }
    2537              :     */
    2538              : 
    2539              :     /*
    2540              :     ///
    2541              :     /// Test list_rels() function, with branches and dropped relations
    2542              :     ///
    2543              :     #[test]
    2544              :     fn test_list_rels_drop() -> Result<()> {
    2545              :         let repo = RepoHarness::create("test_list_rels_drop")?.load();
    2546              :         let tline = create_empty_timeline(repo, TIMELINE_ID)?;
    2547              :         const TESTDB: u32 = 111;
    2548              : 
    2549              :         // Import initial dummy checkpoint record, otherwise the get_timeline() call
    2550              :         // after branching fails below
    2551              :         let mut writer = tline.begin_record(Lsn(0x10));
    2552              :         writer.put_checkpoint(ZERO_CHECKPOINT.clone())?;
    2553              :         writer.finish()?;
    2554              : 
    2555              :         // Create a relation on the timeline
    2556              :         let mut writer = tline.begin_record(Lsn(0x20));
    2557              :         writer.put_rel_page_image(TESTREL_A, 0, TEST_IMG("foo blk 0 at 2"))?;
    2558              :         writer.finish()?;
    2559              : 
    2560              :         let writer = tline.begin_record(Lsn(0x00));
    2561              :         writer.finish()?;
    2562              : 
    2563              :         // Check that list_rels() lists it after LSN 2, but no before it
    2564              :         assert!(!tline.list_rels(0, TESTDB, Lsn(0x10))?.contains(&TESTREL_A));
    2565              :         assert!(tline.list_rels(0, TESTDB, Lsn(0x20))?.contains(&TESTREL_A));
    2566              :         assert!(tline.list_rels(0, TESTDB, Lsn(0x30))?.contains(&TESTREL_A));
    2567              : 
    2568              :         // Create a branch, check that the relation is visible there
    2569              :         repo.branch_timeline(&tline, NEW_TIMELINE_ID, Lsn(0x30))?;
    2570              :         let newtline = match repo.get_timeline(NEW_TIMELINE_ID)?.local_timeline() {
    2571              :             Some(timeline) => timeline,
    2572              :             None => panic!("Should have a local timeline"),
    2573              :         };
    2574              :         let newtline = DatadirTimelineImpl::new(newtline);
    2575              :         assert!(newtline
    2576              :             .list_rels(0, TESTDB, Lsn(0x30))?
    2577              :             .contains(&TESTREL_A));
    2578              : 
    2579              :         // Drop it on the branch
    2580              :         let mut new_writer = newtline.begin_record(Lsn(0x40));
    2581              :         new_writer.drop_relation(TESTREL_A)?;
    2582              :         new_writer.finish()?;
    2583              : 
    2584              :         // Check that it's no longer listed on the branch after the point where it was dropped
    2585              :         assert!(newtline
    2586              :             .list_rels(0, TESTDB, Lsn(0x30))?
    2587              :             .contains(&TESTREL_A));
    2588              :         assert!(!newtline
    2589              :             .list_rels(0, TESTDB, Lsn(0x40))?
    2590              :             .contains(&TESTREL_A));
    2591              : 
    2592              :         // Run checkpoint and garbage collection and check that it's still not visible
    2593              :         newtline.checkpoint(CheckpointConfig::Forced)?;
    2594              :         repo.gc_iteration(Some(NEW_TIMELINE_ID), 0, true)?;
    2595              : 
    2596              :         assert!(!newtline
    2597              :             .list_rels(0, TESTDB, Lsn(0x40))?
    2598              :             .contains(&TESTREL_A));
    2599              : 
    2600              :         Ok(())
    2601              :     }
    2602              :      */
    2603              : 
    2604              :     /*
    2605              :     #[test]
    2606              :     fn test_read_beyond_eof() -> Result<()> {
    2607              :         let repo = RepoHarness::create("test_read_beyond_eof")?.load();
    2608              :         let tline = create_test_timeline(repo, TIMELINE_ID)?;
    2609              : 
    2610              :         make_some_layers(&tline, Lsn(0x20))?;
    2611              :         let mut writer = tline.begin_record(Lsn(0x60));
    2612              :         walingest.put_rel_page_image(
    2613              :             &mut writer,
    2614              :             TESTREL_A,
    2615              :             0,
    2616              :             TEST_IMG(&format!("foo blk 0 at {}", Lsn(0x60))),
    2617              :         )?;
    2618              :         writer.finish()?;
    2619              : 
    2620              :         // Test read before rel creation. Should error out.
    2621              :         assert!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x10), false).is_err());
    2622              : 
    2623              :         // Read block beyond end of relation at different points in time.
    2624              :         // These reads should fall into different delta, image, and in-memory layers.
    2625              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x20), false)?, ZERO_PAGE);
    2626              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x25), false)?, ZERO_PAGE);
    2627              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x30), false)?, ZERO_PAGE);
    2628              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x35), false)?, ZERO_PAGE);
    2629              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x40), false)?, ZERO_PAGE);
    2630              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x45), false)?, ZERO_PAGE);
    2631              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x50), false)?, ZERO_PAGE);
    2632              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x55), false)?, ZERO_PAGE);
    2633              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_A, 1, Lsn(0x60), false)?, ZERO_PAGE);
    2634              : 
    2635              :         // Test on an in-memory layer with no preceding layer
    2636              :         let mut writer = tline.begin_record(Lsn(0x70));
    2637              :         walingest.put_rel_page_image(
    2638              :             &mut writer,
    2639              :             TESTREL_B,
    2640              :             0,
    2641              :             TEST_IMG(&format!("foo blk 0 at {}", Lsn(0x70))),
    2642              :         )?;
    2643              :         writer.finish()?;
    2644              : 
    2645              :         assert_eq!(tline.get_rel_page_at_lsn(TESTREL_B, 1, Lsn(0x70), false)?6, ZERO_PAGE);
    2646              : 
    2647              :         Ok(())
    2648              :     }
    2649              :      */
    2650              : }
        

Generated by: LCOV version 2.1-beta