LCOV - code coverage report
Current view: top level - pageserver/src/tenant/timeline/import_pgdata - flow.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 0.0 % 399 0
Test Date: 2025-02-20 13:11:02 Functions: 0.0 % 50 0

            Line data    Source code
       1              : //! Import a PGDATA directory into an empty root timeline.
       2              : //!
       3              : //! This module is adapted hackathon code by Heikki and Stas.
       4              : //! Other code in the parent module was written by Christian as part of a customer PoC.
       5              : //!
       6              : //! The hackathon code was producing image layer files as a free-standing program.
       7              : //!
       8              : //! It has been modified to
       9              : //! - run inside a running Pageserver, within the proper lifecycles of Timeline -> Tenant(Shard)
      10              : //! - => sharding-awareness: produce image layers with only the data relevant for this shard
      11              : //! - => S3 as the source for the PGDATA instead of local filesystem
      12              : //!
      13              : //! TODOs before productionization:
      14              : //! - ChunkProcessingJob size / ImportJob::total_size does not account for sharding.
      15              : //!   => produced image layers likely too small.
      16              : //! - ChunkProcessingJob should cut up an ImportJob to hit exactly target image layer size.
      17              : //! - asserts / unwraps need to be replaced with errors
      18              : //! - don't trust remote objects will be small (=prevent OOMs in those cases)
      19              : //!     - limit all in-memory buffers in size, or download to disk and read from there
      20              : //! - limit task concurrency
      21              : //! - generally play nice with other tenants in the system
      22              : //!   - importbucket is different bucket than main pageserver storage, so, should be fine wrt S3 rate limits
      23              : //!   - but concerns like network bandwidth, local disk write bandwidth, local disk capacity, etc
      24              : //! - integrate with layer eviction system
      25              : //! - audit for Tenant::cancel nor Timeline::cancel responsivity
      26              : //! - audit for Tenant/Timeline gate holding (we spawn tokio tasks during this flow!)
      27              : //!
      28              : //! An incomplete set of TODOs from the Hackathon:
      29              : //! - version-specific CheckPointData (=> pgv abstraction, already exists for regular walingest)
      30              : 
      31              : use std::sync::Arc;
      32              : 
      33              : use anyhow::{bail, ensure};
      34              : use bytes::Bytes;
      35              : 
      36              : use itertools::Itertools;
      37              : use pageserver_api::{
      38              :     key::{rel_block_to_key, rel_dir_to_key, rel_size_to_key, relmap_file_key, DBDIR_KEY},
      39              :     reltag::RelTag,
      40              :     shard::ShardIdentity,
      41              : };
      42              : use postgres_ffi::{pg_constants, relfile_utils::parse_relfilename, BLCKSZ};
      43              : use tokio::task::JoinSet;
      44              : use tracing::{debug, info_span, instrument, Instrument};
      45              : 
      46              : use crate::{
      47              :     assert_u64_eq_usize::UsizeIsU64,
      48              :     pgdatadir_mapping::{SlruSegmentDirectory, TwoPhaseDirectory},
      49              : };
      50              : use crate::{
      51              :     context::{DownloadBehavior, RequestContext},
      52              :     pgdatadir_mapping::{DbDirectory, RelDirectory},
      53              :     task_mgr::TaskKind,
      54              :     tenant::storage_layer::{ImageLayerWriter, Layer},
      55              : };
      56              : 
      57              : use pageserver_api::key::Key;
      58              : use pageserver_api::key::{
      59              :     slru_block_to_key, slru_dir_to_key, slru_segment_size_to_key, CHECKPOINT_KEY, CONTROLFILE_KEY,
      60              :     TWOPHASEDIR_KEY,
      61              : };
      62              : use pageserver_api::keyspace::singleton_range;
      63              : use pageserver_api::keyspace::{contiguous_range_len, is_contiguous_range};
      64              : use pageserver_api::reltag::SlruKind;
      65              : use utils::bin_ser::BeSer;
      66              : use utils::lsn::Lsn;
      67              : 
      68              : use std::collections::HashSet;
      69              : use std::ops::Range;
      70              : 
      71              : use super::{
      72              :     importbucket_client::{ControlFile, RemoteStorageWrapper},
      73              :     Timeline,
      74              : };
      75              : 
      76              : use remote_storage::RemotePath;
      77              : 
      78            0 : pub async fn run(
      79            0 :     timeline: Arc<Timeline>,
      80            0 :     pgdata_lsn: Lsn,
      81            0 :     control_file: ControlFile,
      82            0 :     storage: RemoteStorageWrapper,
      83            0 :     ctx: &RequestContext,
      84            0 : ) -> anyhow::Result<()> {
      85            0 :     Flow {
      86            0 :         timeline,
      87            0 :         pgdata_lsn,
      88            0 :         control_file,
      89            0 :         tasks: Vec::new(),
      90            0 :         storage,
      91            0 :     }
      92            0 :     .run(ctx)
      93            0 :     .await
      94            0 : }
      95              : 
      96              : struct Flow {
      97              :     timeline: Arc<Timeline>,
      98              :     pgdata_lsn: Lsn,
      99              :     control_file: ControlFile,
     100              :     tasks: Vec<AnyImportTask>,
     101              :     storage: RemoteStorageWrapper,
     102              : }
     103              : 
     104              : impl Flow {
     105              :     /// Perform the ingestion into [`Self::timeline`].
     106              :     /// Assumes the timeline is empty (= no layers).
     107            0 :     pub async fn run(mut self, ctx: &RequestContext) -> anyhow::Result<()> {
     108            0 :         let pgdata_lsn = Lsn(self.control_file.control_file_data().checkPoint).align();
     109            0 : 
     110            0 :         self.pgdata_lsn = pgdata_lsn;
     111              : 
     112            0 :         let datadir = PgDataDir::new(&self.storage).await?;
     113              : 
     114              :         // Import dbdir (00:00:00 keyspace)
     115              :         // This is just constructed here, but will be written to the image layer in the first call to import_db()
     116            0 :         let dbdir_buf = Bytes::from(DbDirectory::ser(&DbDirectory {
     117            0 :             dbdirs: datadir
     118            0 :                 .dbs
     119            0 :                 .iter()
     120            0 :                 .map(|db| ((db.spcnode, db.dboid), true))
     121            0 :                 .collect(),
     122            0 :         })?);
     123            0 :         self.tasks
     124            0 :             .push(ImportSingleKeyTask::new(DBDIR_KEY, dbdir_buf).into());
     125              : 
     126              :         // Import databases (00:spcnode:dbnode keyspace for each db)
     127            0 :         for db in datadir.dbs {
     128            0 :             self.import_db(&db).await?;
     129              :         }
     130              : 
     131              :         // Import SLRUs
     132            0 :         if self.timeline.tenant_shard_id.is_shard_zero() {
     133              :             // pg_xact (01:00 keyspace)
     134            0 :             self.import_slru(SlruKind::Clog, &self.storage.pgdata().join("pg_xact"))
     135            0 :                 .await?;
     136              :             // pg_multixact/members (01:01 keyspace)
     137            0 :             self.import_slru(
     138            0 :                 SlruKind::MultiXactMembers,
     139            0 :                 &self.storage.pgdata().join("pg_multixact/members"),
     140            0 :             )
     141            0 :             .await?;
     142              :             // pg_multixact/offsets (01:02 keyspace)
     143            0 :             self.import_slru(
     144            0 :                 SlruKind::MultiXactOffsets,
     145            0 :                 &self.storage.pgdata().join("pg_multixact/offsets"),
     146            0 :             )
     147            0 :             .await?;
     148            0 :         }
     149              : 
     150              :         // Import pg_twophase.
     151              :         // TODO: as empty
     152            0 :         let twophasedir_buf = TwoPhaseDirectory::ser(&TwoPhaseDirectory {
     153            0 :             xids: HashSet::new(),
     154            0 :         })?;
     155            0 :         self.tasks
     156            0 :             .push(AnyImportTask::SingleKey(ImportSingleKeyTask::new(
     157            0 :                 TWOPHASEDIR_KEY,
     158            0 :                 Bytes::from(twophasedir_buf),
     159            0 :             )));
     160            0 : 
     161            0 :         // Controlfile, checkpoint
     162            0 :         self.tasks
     163            0 :             .push(AnyImportTask::SingleKey(ImportSingleKeyTask::new(
     164            0 :                 CONTROLFILE_KEY,
     165            0 :                 self.control_file.control_file_buf().clone(),
     166            0 :             )));
     167              : 
     168            0 :         let checkpoint_buf = self
     169            0 :             .control_file
     170            0 :             .control_file_data()
     171            0 :             .checkPointCopy
     172            0 :             .encode()?;
     173            0 :         self.tasks
     174            0 :             .push(AnyImportTask::SingleKey(ImportSingleKeyTask::new(
     175            0 :                 CHECKPOINT_KEY,
     176            0 :                 checkpoint_buf,
     177            0 :             )));
     178            0 : 
     179            0 :         // Assigns parts of key space to later parallel jobs
     180            0 :         let mut last_end_key = Key::MIN;
     181            0 :         let mut current_chunk = Vec::new();
     182            0 :         let mut current_chunk_size: usize = 0;
     183            0 :         let mut parallel_jobs = Vec::new();
     184            0 :         for task in std::mem::take(&mut self.tasks).into_iter() {
     185            0 :             if current_chunk_size + task.total_size() > 1024 * 1024 * 1024 {
     186            0 :                 let key_range = last_end_key..task.key_range().start;
     187            0 :                 parallel_jobs.push(ChunkProcessingJob::new(
     188            0 :                     key_range.clone(),
     189            0 :                     std::mem::take(&mut current_chunk),
     190            0 :                     &self,
     191            0 :                 ));
     192            0 :                 last_end_key = key_range.end;
     193            0 :                 current_chunk_size = 0;
     194            0 :             }
     195            0 :             current_chunk_size += task.total_size();
     196            0 :             current_chunk.push(task);
     197              :         }
     198            0 :         parallel_jobs.push(ChunkProcessingJob::new(
     199            0 :             last_end_key..Key::MAX,
     200            0 :             current_chunk,
     201            0 :             &self,
     202            0 :         ));
     203            0 : 
     204            0 :         // Start all jobs simultaneosly
     205            0 :         let mut work = JoinSet::new();
     206              :         // TODO: semaphore?
     207            0 :         for job in parallel_jobs {
     208            0 :             let ctx: RequestContext =
     209            0 :                 ctx.detached_child(TaskKind::ImportPgdata, DownloadBehavior::Error);
     210            0 :             work.spawn(async move { job.run(&ctx).await }.instrument(info_span!("parallel_job")));
     211              :         }
     212            0 :         let mut results = Vec::new();
     213            0 :         while let Some(result) = work.join_next().await {
     214            0 :             match result {
     215            0 :                 Ok(res) => {
     216            0 :                     results.push(res);
     217            0 :                 }
     218            0 :                 Err(_joinset_err) => {
     219            0 :                     results.push(Err(anyhow::anyhow!(
     220            0 :                         "parallel job panicked or cancelled, check pageserver logs"
     221            0 :                     )));
     222            0 :                 }
     223              :             }
     224              :         }
     225              : 
     226            0 :         if results.iter().all(|r| r.is_ok()) {
     227            0 :             Ok(())
     228              :         } else {
     229            0 :             let mut msg = String::new();
     230            0 :             for result in results {
     231            0 :                 if let Err(err) = result {
     232            0 :                     msg.push_str(&format!("{err:?}\n\n"));
     233            0 :                 }
     234              :             }
     235            0 :             bail!("Some parallel jobs failed:\n\n{msg}");
     236              :         }
     237            0 :     }
     238              : 
     239            0 :     #[instrument(level = tracing::Level::DEBUG, skip_all, fields(dboid=%db.dboid, tablespace=%db.spcnode, path=%db.path))]
     240              :     async fn import_db(&mut self, db: &PgDataDirDb) -> anyhow::Result<()> {
     241              :         debug!("start");
     242              :         scopeguard::defer! {
     243              :             debug!("return");
     244              :         }
     245              : 
     246              :         // Import relmap (00:spcnode:dbnode:00:*:00)
     247              :         let relmap_key = relmap_file_key(db.spcnode, db.dboid);
     248              :         debug!("Constructing relmap entry, key {relmap_key}");
     249              :         let relmap_path = db.path.join("pg_filenode.map");
     250              :         let relmap_buf = self.storage.get(&relmap_path).await?;
     251              :         self.tasks
     252              :             .push(AnyImportTask::SingleKey(ImportSingleKeyTask::new(
     253              :                 relmap_key, relmap_buf,
     254              :             )));
     255              : 
     256              :         // Import reldir (00:spcnode:dbnode:00:*:01)
     257              :         let reldir_key = rel_dir_to_key(db.spcnode, db.dboid);
     258              :         debug!("Constructing reldirs entry, key {reldir_key}");
     259              :         let reldir_buf = RelDirectory::ser(&RelDirectory {
     260              :             rels: db
     261              :                 .files
     262              :                 .iter()
     263            0 :                 .map(|f| (f.rel_tag.relnode, f.rel_tag.forknum))
     264              :                 .collect(),
     265              :         })?;
     266              :         self.tasks
     267              :             .push(AnyImportTask::SingleKey(ImportSingleKeyTask::new(
     268              :                 reldir_key,
     269              :                 Bytes::from(reldir_buf),
     270              :             )));
     271              : 
     272              :         // Import data (00:spcnode:dbnode:reloid:fork:blk) and set sizes for each last
     273              :         // segment in a given relation (00:spcnode:dbnode:reloid:fork:ff)
     274              :         for file in &db.files {
     275              :             debug!(%file.path, %file.filesize, "importing file");
     276              :             let len = file.filesize;
     277              :             ensure!(len % 8192 == 0);
     278              :             let start_blk: u32 = file.segno * (1024 * 1024 * 1024 / 8192);
     279              :             let start_key = rel_block_to_key(file.rel_tag, start_blk);
     280              :             let end_key = rel_block_to_key(file.rel_tag, start_blk + (len / 8192) as u32);
     281              :             self.tasks
     282              :                 .push(AnyImportTask::RelBlocks(ImportRelBlocksTask::new(
     283              :                     *self.timeline.get_shard_identity(),
     284              :                     start_key..end_key,
     285              :                     &file.path,
     286              :                     self.storage.clone(),
     287              :                 )));
     288              : 
     289              :             // Set relsize for the last segment (00:spcnode:dbnode:reloid:fork:ff)
     290              :             if let Some(nblocks) = file.nblocks {
     291              :                 let size_key = rel_size_to_key(file.rel_tag);
     292              :                 //debug!("Setting relation size (path={path}, rel_tag={rel_tag}, segno={segno}) to {nblocks}, key {size_key}");
     293              :                 let buf = nblocks.to_le_bytes();
     294              :                 self.tasks
     295              :                     .push(AnyImportTask::SingleKey(ImportSingleKeyTask::new(
     296              :                         size_key,
     297              :                         Bytes::from(buf.to_vec()),
     298              :                     )));
     299              :             }
     300              :         }
     301              : 
     302              :         Ok(())
     303              :     }
     304              : 
     305            0 :     async fn import_slru(&mut self, kind: SlruKind, path: &RemotePath) -> anyhow::Result<()> {
     306            0 :         assert!(self.timeline.tenant_shard_id.is_shard_zero());
     307              : 
     308            0 :         let segments = self.storage.listfilesindir(path).await?;
     309            0 :         let segments: Vec<(String, u32, usize)> = segments
     310            0 :             .into_iter()
     311            0 :             .filter_map(|(path, size)| {
     312            0 :                 let filename = path.object_name()?;
     313            0 :                 let segno = u32::from_str_radix(filename, 16).ok()?;
     314            0 :                 Some((filename.to_string(), segno, size))
     315            0 :             })
     316            0 :             .collect();
     317            0 : 
     318            0 :         // Write SlruDir
     319            0 :         let slrudir_key = slru_dir_to_key(kind);
     320            0 :         let segnos: HashSet<u32> = segments
     321            0 :             .iter()
     322            0 :             .map(|(_path, segno, _size)| *segno)
     323            0 :             .collect();
     324            0 :         let slrudir = SlruSegmentDirectory { segments: segnos };
     325            0 :         let slrudir_buf = SlruSegmentDirectory::ser(&slrudir)?;
     326            0 :         self.tasks
     327            0 :             .push(AnyImportTask::SingleKey(ImportSingleKeyTask::new(
     328            0 :                 slrudir_key,
     329            0 :                 Bytes::from(slrudir_buf),
     330            0 :             )));
     331              : 
     332            0 :         for (segpath, segno, size) in segments {
     333              :             // SlruSegBlocks for each segment
     334            0 :             let p = path.join(&segpath);
     335            0 :             let file_size = size;
     336            0 :             ensure!(file_size % 8192 == 0);
     337            0 :             let nblocks = u32::try_from(file_size / 8192)?;
     338            0 :             let start_key = slru_block_to_key(kind, segno, 0);
     339            0 :             let end_key = slru_block_to_key(kind, segno, nblocks);
     340            0 :             debug!(%p, segno=%segno, %size, %start_key, %end_key, "scheduling SLRU segment");
     341            0 :             self.tasks
     342            0 :                 .push(AnyImportTask::SlruBlocks(ImportSlruBlocksTask::new(
     343            0 :                     start_key..end_key,
     344            0 :                     &p,
     345            0 :                     self.storage.clone(),
     346            0 :                 )));
     347            0 : 
     348            0 :             // Followed by SlruSegSize
     349            0 :             let segsize_key = slru_segment_size_to_key(kind, segno);
     350            0 :             let segsize_buf = nblocks.to_le_bytes();
     351            0 :             self.tasks
     352            0 :                 .push(AnyImportTask::SingleKey(ImportSingleKeyTask::new(
     353            0 :                     segsize_key,
     354            0 :                     Bytes::copy_from_slice(&segsize_buf),
     355            0 :                 )));
     356              :         }
     357            0 :         Ok(())
     358            0 :     }
     359              : }
     360              : 
     361              : //
     362              : // dbdir iteration tools
     363              : //
     364              : 
     365              : struct PgDataDir {
     366              :     pub dbs: Vec<PgDataDirDb>, // spcnode, dboid, path
     367              : }
     368              : 
     369              : struct PgDataDirDb {
     370              :     pub spcnode: u32,
     371              :     pub dboid: u32,
     372              :     pub path: RemotePath,
     373              :     pub files: Vec<PgDataDirDbFile>,
     374              : }
     375              : 
     376              : struct PgDataDirDbFile {
     377              :     pub path: RemotePath,
     378              :     pub rel_tag: RelTag,
     379              :     pub segno: u32,
     380              :     pub filesize: usize,
     381              :     // Cummulative size of the given fork, set only for the last segment of that fork
     382              :     pub nblocks: Option<usize>,
     383              : }
     384              : 
     385              : impl PgDataDir {
     386            0 :     async fn new(storage: &RemoteStorageWrapper) -> anyhow::Result<Self> {
     387            0 :         let datadir_path = storage.pgdata();
     388            0 :         // Import ordinary databases, DEFAULTTABLESPACE_OID is smaller than GLOBALTABLESPACE_OID, so import them first
     389            0 :         // Traverse database in increasing oid order
     390            0 : 
     391            0 :         let basedir = &datadir_path.join("base");
     392            0 :         let db_oids: Vec<_> = storage
     393            0 :             .listdir(basedir)
     394            0 :             .await?
     395            0 :             .into_iter()
     396            0 :             .filter_map(|path| path.object_name().and_then(|name| name.parse::<u32>().ok()))
     397            0 :             .sorted()
     398            0 :             .collect();
     399            0 :         debug!(?db_oids, "found databases");
     400            0 :         let mut databases = Vec::new();
     401            0 :         for dboid in db_oids {
     402            0 :             databases.push(
     403            0 :                 PgDataDirDb::new(
     404            0 :                     storage,
     405            0 :                     &basedir.join(dboid.to_string()),
     406            0 :                     pg_constants::DEFAULTTABLESPACE_OID,
     407            0 :                     dboid,
     408            0 :                     &datadir_path,
     409            0 :                 )
     410            0 :                 .await?,
     411              :             );
     412              :         }
     413              : 
     414              :         // special case for global catalogs
     415            0 :         databases.push(
     416            0 :             PgDataDirDb::new(
     417            0 :                 storage,
     418            0 :                 &datadir_path.join("global"),
     419            0 :                 postgres_ffi::pg_constants::GLOBALTABLESPACE_OID,
     420            0 :                 0,
     421            0 :                 &datadir_path,
     422            0 :             )
     423            0 :             .await?,
     424              :         );
     425              : 
     426            0 :         databases.sort_by_key(|db| (db.spcnode, db.dboid));
     427            0 : 
     428            0 :         Ok(Self { dbs: databases })
     429            0 :     }
     430              : }
     431              : 
     432              : impl PgDataDirDb {
     433            0 :     #[instrument(level = tracing::Level::DEBUG, skip_all, fields(%dboid, %db_path))]
     434              :     async fn new(
     435              :         storage: &RemoteStorageWrapper,
     436              :         db_path: &RemotePath,
     437              :         spcnode: u32,
     438              :         dboid: u32,
     439              :         datadir_path: &RemotePath,
     440              :     ) -> anyhow::Result<Self> {
     441              :         let mut files: Vec<PgDataDirDbFile> = storage
     442              :             .listfilesindir(db_path)
     443              :             .await?
     444              :             .into_iter()
     445            0 :             .filter_map(|(path, size)| {
     446            0 :                 debug!(%path, %size, "found file in dbdir");
     447            0 :                 path.object_name().and_then(|name| {
     448            0 :                     // returns (relnode, forknum, segno)
     449            0 :                     parse_relfilename(name).ok().map(|x| (size, x))
     450            0 :                 })
     451            0 :             })
     452            0 :             .sorted_by_key(|(_, relfilename)| *relfilename)
     453            0 :             .map(|(filesize, (relnode, forknum, segno))| {
     454            0 :                 let rel_tag = RelTag {
     455            0 :                     spcnode,
     456            0 :                     dbnode: dboid,
     457            0 :                     relnode,
     458            0 :                     forknum,
     459            0 :                 };
     460            0 : 
     461            0 :                 let path = datadir_path.join(rel_tag.to_segfile_name(segno));
     462            0 :                 assert!(filesize % BLCKSZ as usize == 0); // TODO: this should result in an error
     463            0 :                 let nblocks = filesize / BLCKSZ as usize;
     464            0 : 
     465            0 :                 PgDataDirDbFile {
     466            0 :                     path,
     467            0 :                     filesize,
     468            0 :                     rel_tag,
     469            0 :                     segno,
     470            0 :                     nblocks: Some(nblocks), // first non-cummulative sizes
     471            0 :                 }
     472            0 :             })
     473              :             .collect();
     474              : 
     475              :         // Set cummulative sizes. Do all of that math here, so that later we could easier
     476              :         // parallelize over segments and know with which segments we need to write relsize
     477              :         // entry.
     478              :         let mut cumulative_nblocks: usize = 0;
     479              :         let mut prev_rel_tag: Option<RelTag> = None;
     480              :         for i in 0..files.len() {
     481              :             if prev_rel_tag == Some(files[i].rel_tag) {
     482              :                 cumulative_nblocks += files[i].nblocks.unwrap();
     483              :             } else {
     484              :                 cumulative_nblocks = files[i].nblocks.unwrap();
     485              :             }
     486              : 
     487              :             files[i].nblocks = if i == files.len() - 1 || files[i + 1].rel_tag != files[i].rel_tag {
     488              :                 Some(cumulative_nblocks)
     489              :             } else {
     490              :                 None
     491              :             };
     492              : 
     493              :             prev_rel_tag = Some(files[i].rel_tag);
     494              :         }
     495              : 
     496              :         Ok(PgDataDirDb {
     497              :             files,
     498              :             path: db_path.clone(),
     499              :             spcnode,
     500              :             dboid,
     501              :         })
     502              :     }
     503              : }
     504              : 
     505              : trait ImportTask {
     506              :     fn key_range(&self) -> Range<Key>;
     507              : 
     508            0 :     fn total_size(&self) -> usize {
     509            0 :         // TODO: revisit this
     510            0 :         if is_contiguous_range(&self.key_range()) {
     511            0 :             contiguous_range_len(&self.key_range()) as usize * 8192
     512              :         } else {
     513            0 :             u32::MAX as usize
     514              :         }
     515            0 :     }
     516              : 
     517              :     async fn doit(
     518              :         self,
     519              :         layer_writer: &mut ImageLayerWriter,
     520              :         ctx: &RequestContext,
     521              :     ) -> anyhow::Result<usize>;
     522              : }
     523              : 
     524              : struct ImportSingleKeyTask {
     525              :     key: Key,
     526              :     buf: Bytes,
     527              : }
     528              : 
     529              : impl ImportSingleKeyTask {
     530            0 :     fn new(key: Key, buf: Bytes) -> Self {
     531            0 :         ImportSingleKeyTask { key, buf }
     532            0 :     }
     533              : }
     534              : 
     535              : impl ImportTask for ImportSingleKeyTask {
     536            0 :     fn key_range(&self) -> Range<Key> {
     537            0 :         singleton_range(self.key)
     538            0 :     }
     539              : 
     540            0 :     async fn doit(
     541            0 :         self,
     542            0 :         layer_writer: &mut ImageLayerWriter,
     543            0 :         ctx: &RequestContext,
     544            0 :     ) -> anyhow::Result<usize> {
     545            0 :         layer_writer.put_image(self.key, self.buf, ctx).await?;
     546            0 :         Ok(1)
     547            0 :     }
     548              : }
     549              : 
     550              : struct ImportRelBlocksTask {
     551              :     shard_identity: ShardIdentity,
     552              :     key_range: Range<Key>,
     553              :     path: RemotePath,
     554              :     storage: RemoteStorageWrapper,
     555              : }
     556              : 
     557              : impl ImportRelBlocksTask {
     558            0 :     fn new(
     559            0 :         shard_identity: ShardIdentity,
     560            0 :         key_range: Range<Key>,
     561            0 :         path: &RemotePath,
     562            0 :         storage: RemoteStorageWrapper,
     563            0 :     ) -> Self {
     564            0 :         ImportRelBlocksTask {
     565            0 :             shard_identity,
     566            0 :             key_range,
     567            0 :             path: path.clone(),
     568            0 :             storage,
     569            0 :         }
     570            0 :     }
     571              : }
     572              : 
     573              : impl ImportTask for ImportRelBlocksTask {
     574            0 :     fn key_range(&self) -> Range<Key> {
     575            0 :         self.key_range.clone()
     576            0 :     }
     577              : 
     578            0 :     #[instrument(level = tracing::Level::DEBUG, skip_all, fields(%self.path))]
     579              :     async fn doit(
     580              :         self,
     581              :         layer_writer: &mut ImageLayerWriter,
     582              :         ctx: &RequestContext,
     583              :     ) -> anyhow::Result<usize> {
     584              :         debug!("Importing relation file");
     585              : 
     586              :         let (rel_tag, start_blk) = self.key_range.start.to_rel_block()?;
     587              :         let (rel_tag_end, end_blk) = self.key_range.end.to_rel_block()?;
     588              :         assert_eq!(rel_tag, rel_tag_end);
     589              : 
     590              :         let ranges = (start_blk..end_blk)
     591              :             .enumerate()
     592            0 :             .filter_map(|(i, blknum)| {
     593            0 :                 let key = rel_block_to_key(rel_tag, blknum);
     594            0 :                 if self.shard_identity.is_key_disposable(&key) {
     595            0 :                     return None;
     596            0 :                 }
     597            0 :                 let file_offset = i.checked_mul(8192).unwrap();
     598            0 :                 Some((
     599            0 :                     vec![key],
     600            0 :                     file_offset,
     601            0 :                     file_offset.checked_add(8192).unwrap(),
     602            0 :                 ))
     603            0 :             })
     604            0 :             .coalesce(|(mut acc, acc_start, acc_end), (mut key, start, end)| {
     605            0 :                 assert_eq!(key.len(), 1);
     606            0 :                 assert!(!acc.is_empty());
     607            0 :                 assert!(acc_end > acc_start);
     608            0 :                 if acc_end == start /* TODO additional max range check here, to limit memory consumption per task to X */ {
     609            0 :                     acc.push(key.pop().unwrap());
     610            0 :                     Ok((acc, acc_start, end))
     611              :                 } else {
     612            0 :                     Err(((acc, acc_start, acc_end), (key, start, end)))
     613              :                 }
     614            0 :             });
     615              : 
     616              :         let mut nimages = 0;
     617              :         for (keys, range_start, range_end) in ranges {
     618              :             let range_buf = self
     619              :                 .storage
     620              :                 .get_range(&self.path, range_start.into_u64(), range_end.into_u64())
     621              :                 .await?;
     622              :             let mut buf = Bytes::from(range_buf);
     623              :             // TODO: batched writes
     624              :             for key in keys {
     625              :                 let image = buf.split_to(8192);
     626              :                 layer_writer.put_image(key, image, ctx).await?;
     627              :                 nimages += 1;
     628              :             }
     629              :         }
     630              : 
     631              :         Ok(nimages)
     632              :     }
     633              : }
     634              : 
     635              : struct ImportSlruBlocksTask {
     636              :     key_range: Range<Key>,
     637              :     path: RemotePath,
     638              :     storage: RemoteStorageWrapper,
     639              : }
     640              : 
     641              : impl ImportSlruBlocksTask {
     642            0 :     fn new(key_range: Range<Key>, path: &RemotePath, storage: RemoteStorageWrapper) -> Self {
     643            0 :         ImportSlruBlocksTask {
     644            0 :             key_range,
     645            0 :             path: path.clone(),
     646            0 :             storage,
     647            0 :         }
     648            0 :     }
     649              : }
     650              : 
     651              : impl ImportTask for ImportSlruBlocksTask {
     652            0 :     fn key_range(&self) -> Range<Key> {
     653            0 :         self.key_range.clone()
     654            0 :     }
     655              : 
     656            0 :     async fn doit(
     657            0 :         self,
     658            0 :         layer_writer: &mut ImageLayerWriter,
     659            0 :         ctx: &RequestContext,
     660            0 :     ) -> anyhow::Result<usize> {
     661            0 :         debug!("Importing SLRU segment file {}", self.path);
     662            0 :         let buf = self.storage.get(&self.path).await?;
     663              : 
     664            0 :         let (kind, segno, start_blk) = self.key_range.start.to_slru_block()?;
     665            0 :         let (_kind, _segno, end_blk) = self.key_range.end.to_slru_block()?;
     666            0 :         let mut blknum = start_blk;
     667            0 :         let mut nimages = 0;
     668            0 :         let mut file_offset = 0;
     669            0 :         while blknum < end_blk {
     670            0 :             let key = slru_block_to_key(kind, segno, blknum);
     671            0 :             let buf = &buf[file_offset..(file_offset + 8192)];
     672            0 :             file_offset += 8192;
     673            0 :             layer_writer
     674            0 :                 .put_image(key, Bytes::copy_from_slice(buf), ctx)
     675            0 :                 .await?;
     676            0 :             nimages += 1;
     677            0 :             blknum += 1;
     678              :         }
     679            0 :         Ok(nimages)
     680            0 :     }
     681              : }
     682              : 
     683              : enum AnyImportTask {
     684              :     SingleKey(ImportSingleKeyTask),
     685              :     RelBlocks(ImportRelBlocksTask),
     686              :     SlruBlocks(ImportSlruBlocksTask),
     687              : }
     688              : 
     689              : impl ImportTask for AnyImportTask {
     690            0 :     fn key_range(&self) -> Range<Key> {
     691            0 :         match self {
     692            0 :             Self::SingleKey(t) => t.key_range(),
     693            0 :             Self::RelBlocks(t) => t.key_range(),
     694            0 :             Self::SlruBlocks(t) => t.key_range(),
     695              :         }
     696            0 :     }
     697              :     /// returns the number of images put into the `layer_writer`
     698            0 :     async fn doit(
     699            0 :         self,
     700            0 :         layer_writer: &mut ImageLayerWriter,
     701            0 :         ctx: &RequestContext,
     702            0 :     ) -> anyhow::Result<usize> {
     703            0 :         match self {
     704            0 :             Self::SingleKey(t) => t.doit(layer_writer, ctx).await,
     705            0 :             Self::RelBlocks(t) => t.doit(layer_writer, ctx).await,
     706            0 :             Self::SlruBlocks(t) => t.doit(layer_writer, ctx).await,
     707              :         }
     708            0 :     }
     709              : }
     710              : 
     711              : impl From<ImportSingleKeyTask> for AnyImportTask {
     712            0 :     fn from(t: ImportSingleKeyTask) -> Self {
     713            0 :         Self::SingleKey(t)
     714            0 :     }
     715              : }
     716              : 
     717              : impl From<ImportRelBlocksTask> for AnyImportTask {
     718            0 :     fn from(t: ImportRelBlocksTask) -> Self {
     719            0 :         Self::RelBlocks(t)
     720            0 :     }
     721              : }
     722              : 
     723              : impl From<ImportSlruBlocksTask> for AnyImportTask {
     724            0 :     fn from(t: ImportSlruBlocksTask) -> Self {
     725            0 :         Self::SlruBlocks(t)
     726            0 :     }
     727              : }
     728              : 
     729              : struct ChunkProcessingJob {
     730              :     timeline: Arc<Timeline>,
     731              :     range: Range<Key>,
     732              :     tasks: Vec<AnyImportTask>,
     733              : 
     734              :     pgdata_lsn: Lsn,
     735              : }
     736              : 
     737              : impl ChunkProcessingJob {
     738            0 :     fn new(range: Range<Key>, tasks: Vec<AnyImportTask>, env: &Flow) -> Self {
     739            0 :         assert!(env.pgdata_lsn.is_valid());
     740            0 :         Self {
     741            0 :             timeline: env.timeline.clone(),
     742            0 :             range,
     743            0 :             tasks,
     744            0 :             pgdata_lsn: env.pgdata_lsn,
     745            0 :         }
     746            0 :     }
     747              : 
     748            0 :     async fn run(self, ctx: &RequestContext) -> anyhow::Result<()> {
     749            0 :         let mut writer = ImageLayerWriter::new(
     750            0 :             self.timeline.conf,
     751            0 :             self.timeline.timeline_id,
     752            0 :             self.timeline.tenant_shard_id,
     753            0 :             &self.range,
     754            0 :             self.pgdata_lsn,
     755            0 :             ctx,
     756            0 :         )
     757            0 :         .await?;
     758              : 
     759            0 :         let mut nimages = 0;
     760            0 :         for task in self.tasks {
     761            0 :             nimages += task.doit(&mut writer, ctx).await?;
     762              :         }
     763              : 
     764            0 :         let resident_layer = if nimages > 0 {
     765            0 :             let (desc, path) = writer.finish(ctx).await?;
     766            0 :             Layer::finish_creating(self.timeline.conf, &self.timeline, desc, &path)?
     767              :         } else {
     768              :             // dropping the writer cleans up
     769            0 :             return Ok(());
     770              :         };
     771              : 
     772              :         // this is sharing the same code as create_image_layers
     773            0 :         let mut guard = self.timeline.layers.write().await;
     774            0 :         guard
     775            0 :             .open_mut()?
     776            0 :             .track_new_image_layers(&[resident_layer.clone()], &self.timeline.metrics);
     777            0 :         crate::tenant::timeline::drop_wlock(guard);
     778            0 : 
     779            0 :         // Schedule the layer for upload but don't add barriers such as
     780            0 :         // wait for completion or index upload, so we don't inhibit upload parallelism.
     781            0 :         // TODO: limit upload parallelism somehow (e.g. by limiting concurrency of jobs?)
     782            0 :         // TODO: or regulate parallelism by upload queue depth? Prob should happen at a higher level.
     783            0 :         self.timeline
     784            0 :             .remote_client
     785            0 :             .schedule_layer_file_upload(resident_layer)?;
     786              : 
     787            0 :         Ok(())
     788            0 :     }
     789              : }
        

Generated by: LCOV version 2.1-beta