LCOV - code coverage report
Current view: top level - pageserver/src - import_datadir.rs (source / functions) Coverage Total Hit
Test: c639aa5f7ab62b43d647b10f40d15a15686ce8a9.info Lines: 84.6 % 468 396
Test Date: 2024-02-12 20:26:03 Functions: 34.6 % 104 36

            Line data    Source code
       1              : //!
       2              : //! Import data and WAL from a PostgreSQL data directory and WAL segments into
       3              : //! a neon Timeline.
       4              : //!
       5              : use std::io::SeekFrom;
       6              : use std::path::{Path, PathBuf};
       7              : 
       8              : use anyhow::{bail, ensure, Context, Result};
       9              : use async_compression::tokio::bufread::ZstdDecoder;
      10              : use async_compression::{tokio::write::ZstdEncoder, zstd::CParameter, Level};
      11              : use bytes::Bytes;
      12              : use camino::Utf8Path;
      13              : use futures::StreamExt;
      14              : use nix::NixPath;
      15              : use tokio::fs::{File, OpenOptions};
      16              : use tokio::io::{AsyncBufRead, AsyncRead, AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
      17              : use tokio_tar::Archive;
      18              : use tokio_tar::Builder;
      19              : use tokio_tar::HeaderMode;
      20              : use tracing::*;
      21              : use walkdir::WalkDir;
      22              : 
      23              : use crate::context::RequestContext;
      24              : use crate::metrics::WAL_INGEST;
      25              : use crate::pgdatadir_mapping::*;
      26              : use crate::tenant::remote_timeline_client::INITDB_PATH;
      27              : use crate::tenant::Timeline;
      28              : use crate::walingest::WalIngest;
      29              : use crate::walrecord::DecodedWALRecord;
      30              : use pageserver_api::reltag::{RelTag, SlruKind};
      31              : use postgres_ffi::pg_constants;
      32              : use postgres_ffi::relfile_utils::*;
      33              : use postgres_ffi::waldecoder::WalStreamDecoder;
      34              : use postgres_ffi::ControlFileData;
      35              : use postgres_ffi::DBState_DB_SHUTDOWNED;
      36              : use postgres_ffi::Oid;
      37              : use postgres_ffi::XLogFileName;
      38              : use postgres_ffi::{BLCKSZ, WAL_SEGMENT_SIZE};
      39              : use utils::lsn::Lsn;
      40              : 
      41              : // Returns checkpoint LSN from controlfile
      42          597 : pub fn get_lsn_from_controlfile(path: &Utf8Path) -> Result<Lsn> {
      43          597 :     // Read control file to extract the LSN
      44          597 :     let controlfile_path = path.join("global").join("pg_control");
      45          597 :     let controlfile_buf = std::fs::read(&controlfile_path)
      46          597 :         .with_context(|| format!("reading controlfile: {controlfile_path}"))?;
      47          597 :     let controlfile = ControlFileData::decode(&controlfile_buf)?;
      48          597 :     let lsn = controlfile.checkPoint;
      49          597 : 
      50          597 :     Ok(Lsn(lsn))
      51          597 : }
      52              : 
      53              : ///
      54              : /// Import all relation data pages from local disk into the repository.
      55              : ///
      56              : /// This is currently only used to import a cluster freshly created by initdb.
      57              : /// The code that deals with the checkpoint would not work right if the
      58              : /// cluster was not shut down cleanly.
      59          596 : pub async fn import_timeline_from_postgres_datadir(
      60          596 :     tline: &Timeline,
      61          596 :     pgdata_path: &Utf8Path,
      62          596 :     pgdata_lsn: Lsn,
      63          596 :     ctx: &RequestContext,
      64          596 : ) -> Result<()> {
      65          596 :     let mut pg_control: Option<ControlFileData> = None;
      66          596 : 
      67          596 :     // TODO this shoud be start_lsn, which is not necessarily equal to end_lsn (aka lsn)
      68          596 :     // Then fishing out pg_control would be unnecessary
      69          596 :     let mut modification = tline.begin_modification(pgdata_lsn);
      70          596 :     modification.init_empty()?;
      71              : 
      72              :     // Import all but pg_wal
      73          596 :     let all_but_wal = WalkDir::new(pgdata_path)
      74          596 :         .into_iter()
      75       580600 :         .filter_entry(|entry| !entry.path().ends_with("pg_wal"));
      76       580600 :     for entry in all_but_wal {
      77       580004 :         let entry = entry?;
      78       580004 :         let metadata = entry.metadata().expect("error getting dir entry metadata");
      79       580004 :         if metadata.is_file() {
      80       565700 :             let absolute_path = entry.path();
      81       565700 :             let relative_path = absolute_path.strip_prefix(pgdata_path)?;
      82              : 
      83       565700 :             let mut file = tokio::fs::File::open(absolute_path).await?;
      84       565700 :             let len = metadata.len() as usize;
      85          596 :             if let Some(control_file) =
      86      2386427 :                 import_file(&mut modification, relative_path, &mut file, len, ctx).await?
      87          596 :             {
      88          596 :                 pg_control = Some(control_file);
      89       565104 :             }
      90       565700 :             modification.flush(ctx).await?;
      91        14304 :         }
      92              :     }
      93              : 
      94              :     // We're done importing all the data files.
      95        42862 :     modification.commit(ctx).await?;
      96              : 
      97              :     // We expect the Postgres server to be shut down cleanly.
      98          596 :     let pg_control = pg_control.context("pg_control file not found")?;
      99          596 :     ensure!(
     100          596 :         pg_control.state == DBState_DB_SHUTDOWNED,
     101            0 :         "Postgres cluster was not shut down cleanly"
     102              :     );
     103          596 :     ensure!(
     104          596 :         pg_control.checkPointCopy.redo == pgdata_lsn.0,
     105            0 :         "unexpected checkpoint REDO pointer"
     106              :     );
     107              : 
     108              :     // Import WAL. This is needed even when starting from a shutdown checkpoint, because
     109              :     // this reads the checkpoint record itself, advancing the tip of the timeline to
     110              :     // *after* the checkpoint record. And crucially, it initializes the 'prev_lsn'.
     111          596 :     import_wal(
     112          596 :         &pgdata_path.join("pg_wal"),
     113          596 :         tline,
     114          596 :         Lsn(pg_control.checkPointCopy.redo),
     115          596 :         pgdata_lsn,
     116          596 :         ctx,
     117          596 :     )
     118          589 :     .await?;
     119              : 
     120          596 :     Ok(())
     121          596 : }
     122              : 
     123              : // subroutine of import_timeline_from_postgres_datadir(), to load one relation file.
     124       565245 : async fn import_rel(
     125       565245 :     modification: &mut DatadirModification<'_>,
     126       565245 :     path: &Path,
     127       565245 :     spcoid: Oid,
     128       565245 :     dboid: Oid,
     129       565245 :     reader: &mut (impl AsyncRead + Unpin),
     130       565245 :     len: usize,
     131       565245 :     ctx: &RequestContext,
     132       565245 : ) -> anyhow::Result<()> {
     133              :     // Does it look like a relation file?
     134            0 :     trace!("importing rel file {}", path.display());
     135              : 
     136       565245 :     let filename = &path
     137       565245 :         .file_name()
     138       565245 :         .expect("missing rel filename")
     139       565245 :         .to_string_lossy();
     140       565245 :     let (relnode, forknum, segno) = parse_relfilename(filename).map_err(|e| {
     141            0 :         warn!("unrecognized file in postgres datadir: {:?} ({})", path, e);
     142            0 :         e
     143       565245 :     })?;
     144              : 
     145       565245 :     let mut buf: [u8; 8192] = [0u8; 8192];
     146       565245 : 
     147       565245 :     ensure!(len % BLCKSZ as usize == 0);
     148       565245 :     let nblocks = len / BLCKSZ as usize;
     149       565245 : 
     150       565245 :     let rel = RelTag {
     151       565245 :         spcnode: spcoid,
     152       565245 :         dbnode: dboid,
     153       565245 :         relnode,
     154       565245 :         forknum,
     155       565245 :     };
     156       565245 : 
     157       565245 :     let mut blknum: u32 = segno * (1024 * 1024 * 1024 / BLCKSZ as u32);
     158              : 
     159              :     // Call put_rel_creation for every segment of the relation,
     160              :     // because there is no guarantee about the order in which we are processing segments.
     161              :     // ignore "relation already exists" error
     162              :     //
     163              :     // FIXME: Keep track of which relations we've already created?
     164              :     // https://github.com/neondatabase/neon/issues/3309
     165       565245 :     if let Err(e) = modification
     166       565245 :         .put_rel_creation(rel, nblocks as u32, ctx)
     167            0 :         .await
     168              :     {
     169            0 :         match e {
     170              :             RelationError::AlreadyExists => {
     171            0 :                 debug!("Relation {} already exist. We must be extending it.", rel)
     172              :             }
     173            0 :             _ => return Err(e.into()),
     174              :         }
     175       565245 :     }
     176              : 
     177              :     loop {
     178      2454141 :         let r = reader.read_exact(&mut buf).await;
     179      2454141 :         match r {
     180              :             Ok(_) => {
     181      1888896 :                 modification.put_rel_page_image(rel, blknum, Bytes::copy_from_slice(&buf))?;
     182              :             }
     183              : 
     184              :             // TODO: UnexpectedEof is expected
     185       565245 :             Err(err) => match err.kind() {
     186              :                 std::io::ErrorKind::UnexpectedEof => {
     187              :                     // reached EOF. That's expected.
     188       565245 :                     let relative_blknum = blknum - segno * (1024 * 1024 * 1024 / BLCKSZ as u32);
     189       565245 :                     ensure!(relative_blknum == nblocks as u32, "unexpected EOF");
     190       565245 :                     break;
     191              :                 }
     192              :                 _ => {
     193            0 :                     bail!("error reading file {}: {:#}", path.display(), err);
     194              :                 }
     195              :             },
     196              :         };
     197      1888896 :         blknum += 1;
     198              :     }
     199              : 
     200              :     // Update relation size
     201              :     //
     202              :     // If we process rel segments out of order,
     203              :     // put_rel_extend will skip the update.
     204       565245 :     modification.put_rel_extend(rel, blknum, ctx).await?;
     205              : 
     206       565245 :     Ok(())
     207       565245 : }
     208              : 
     209              : /// Import an SLRU segment file
     210              : ///
     211         1822 : async fn import_slru(
     212         1822 :     modification: &mut DatadirModification<'_>,
     213         1822 :     slru: SlruKind,
     214         1822 :     path: &Path,
     215         1822 :     reader: &mut (impl AsyncRead + Unpin),
     216         1822 :     len: usize,
     217         1822 :     ctx: &RequestContext,
     218         1822 : ) -> anyhow::Result<()> {
     219         1822 :     info!("importing slru file {path:?}");
     220              : 
     221         1822 :     let mut buf: [u8; 8192] = [0u8; 8192];
     222         1822 :     let filename = &path
     223         1822 :         .file_name()
     224         1822 :         .with_context(|| format!("missing slru filename for path {path:?}"))?
     225         1822 :         .to_string_lossy();
     226         1822 :     let segno = u32::from_str_radix(filename, 16)?;
     227              : 
     228         1822 :     ensure!(len % BLCKSZ as usize == 0); // we assume SLRU block size is the same as BLCKSZ
     229         1822 :     let nblocks = len / BLCKSZ as usize;
     230         1822 : 
     231         1822 :     ensure!(nblocks <= pg_constants::SLRU_PAGES_PER_SEGMENT as usize);
     232              : 
     233         1822 :     modification
     234         1822 :         .put_slru_segment_creation(slru, segno, nblocks as u32, ctx)
     235            0 :         .await?;
     236              : 
     237         1822 :     let mut rpageno = 0;
     238              :     loop {
     239         3652 :         let r = reader.read_exact(&mut buf).await;
     240         3652 :         match r {
     241              :             Ok(_) => {
     242         1830 :                 modification.put_slru_page_image(
     243         1830 :                     slru,
     244         1830 :                     segno,
     245         1830 :                     rpageno,
     246         1830 :                     Bytes::copy_from_slice(&buf),
     247         1830 :                 )?;
     248              :             }
     249              : 
     250              :             // TODO: UnexpectedEof is expected
     251         1822 :             Err(err) => match err.kind() {
     252              :                 std::io::ErrorKind::UnexpectedEof => {
     253              :                     // reached EOF. That's expected.
     254         1822 :                     ensure!(rpageno == nblocks as u32, "unexpected EOF");
     255         1822 :                     break;
     256              :                 }
     257              :                 _ => {
     258            0 :                     bail!("error reading file {}: {:#}", path.display(), err);
     259              :                 }
     260              :             },
     261              :         };
     262         1830 :         rpageno += 1;
     263              :     }
     264              : 
     265         1822 :     Ok(())
     266         1822 : }
     267              : 
     268              : /// Scan PostgreSQL WAL files in given directory and load all records between
     269              : /// 'startpoint' and 'endpoint' into the repository.
     270          596 : async fn import_wal(
     271          596 :     walpath: &Utf8Path,
     272          596 :     tline: &Timeline,
     273          596 :     startpoint: Lsn,
     274          596 :     endpoint: Lsn,
     275          596 :     ctx: &RequestContext,
     276          596 : ) -> anyhow::Result<()> {
     277          596 :     let mut waldecoder = WalStreamDecoder::new(startpoint, tline.pg_version);
     278          596 : 
     279          596 :     let mut segno = startpoint.segment_number(WAL_SEGMENT_SIZE);
     280          596 :     let mut offset = startpoint.segment_offset(WAL_SEGMENT_SIZE);
     281          596 :     let mut last_lsn = startpoint;
     282              : 
     283          596 :     let mut walingest = WalIngest::new(tline, startpoint, ctx).await?;
     284              : 
     285         1192 :     while last_lsn <= endpoint {
     286              :         // FIXME: assume postgresql tli 1 for now
     287          596 :         let filename = XLogFileName(1, segno, WAL_SEGMENT_SIZE);
     288          596 :         let mut buf = Vec::new();
     289          596 : 
     290          596 :         // Read local file
     291          596 :         let mut path = walpath.join(&filename);
     292          596 : 
     293          596 :         // It could be as .partial
     294          596 :         if !PathBuf::from(&path).exists() {
     295            0 :             path = walpath.join(filename + ".partial");
     296          596 :         }
     297              : 
     298              :         // Slurp the WAL file
     299          596 :         let mut file = std::fs::File::open(&path)?;
     300              : 
     301          596 :         if offset > 0 {
     302              :             use std::io::Seek;
     303          596 :             file.seek(std::io::SeekFrom::Start(offset as u64))?;
     304            0 :         }
     305              : 
     306              :         use std::io::Read;
     307          596 :         let nread = file.read_to_end(&mut buf)?;
     308          596 :         if nread != WAL_SEGMENT_SIZE - offset {
     309              :             // Maybe allow this for .partial files?
     310            0 :             error!("read only {} bytes from WAL file", nread);
     311          596 :         }
     312              : 
     313          596 :         waldecoder.feed_bytes(&buf);
     314          596 : 
     315          596 :         let mut nrecords = 0;
     316          596 :         let mut modification = tline.begin_modification(last_lsn);
     317          596 :         let mut decoded = DecodedWALRecord::default();
     318         1192 :         while last_lsn <= endpoint {
     319          596 :             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
     320          596 :                 walingest
     321          596 :                     .ingest_record(recdata, lsn, &mut modification, &mut decoded, ctx)
     322            0 :                     .await?;
     323          596 :                 WAL_INGEST.records_committed.inc();
     324          596 : 
     325          596 :                 modification.commit(ctx).await?;
     326          596 :                 last_lsn = lsn;
     327          596 : 
     328          596 :                 nrecords += 1;
     329              : 
     330            0 :                 trace!("imported record at {} (end {})", lsn, endpoint);
     331            0 :             }
     332              :         }
     333              : 
     334            0 :         debug!("imported {} records up to {}", nrecords, last_lsn);
     335              : 
     336          596 :         segno += 1;
     337          596 :         offset = 0;
     338              :     }
     339              : 
     340          596 :     if last_lsn != startpoint {
     341          596 :         info!("reached end of WAL at {}", last_lsn);
     342              :     } else {
     343            0 :         info!("no WAL to import at {}", last_lsn);
     344              :     }
     345              : 
     346          596 :     Ok(())
     347          596 : }
     348              : 
     349           12 : pub async fn import_basebackup_from_tar(
     350           12 :     tline: &Timeline,
     351           12 :     reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     352           12 :     base_lsn: Lsn,
     353           12 :     ctx: &RequestContext,
     354           12 : ) -> Result<()> {
     355           12 :     info!("importing base at {base_lsn}");
     356           12 :     let mut modification = tline.begin_modification(base_lsn);
     357           12 :     modification.init_empty()?;
     358              : 
     359           12 :     let mut pg_control: Option<ControlFileData> = None;
     360              : 
     361              :     // Import base
     362           12 :     let mut entries = Archive::new(reader).entries()?;
     363        10765 :     while let Some(base_tar_entry) = entries.next().await {
     364        10753 :         let mut entry = base_tar_entry?;
     365        10753 :         let header = entry.header();
     366        10753 :         let len = header.entry_size()? as usize;
     367        10753 :         let file_path = header.path()?.into_owned();
     368        10753 : 
     369        10753 :         match header.entry_type() {
     370              :             tokio_tar::EntryType::Regular => {
     371           10 :                 if let Some(res) =
     372        10478 :                     import_file(&mut modification, file_path.as_ref(), &mut entry, len, ctx).await?
     373           10 :                 {
     374           10 :                     // We found the pg_control file.
     375           10 :                     pg_control = Some(res);
     376        10468 :                 }
     377        10478 :                 modification.flush(ctx).await?;
     378              :             }
     379              :             tokio_tar::EntryType::Directory => {
     380            0 :                 debug!("directory {:?}", file_path);
     381              :             }
     382              :             _ => {
     383            0 :                 bail!(
     384            0 :                     "entry {} in backup tar archive is of unexpected type: {:?}",
     385            0 :                     file_path.display(),
     386            0 :                     header.entry_type()
     387            0 :                 );
     388              :             }
     389              :         }
     390              :     }
     391              : 
     392              :     // sanity check: ensure that pg_control is loaded
     393           12 :     let _pg_control = pg_control.context("pg_control file not found")?;
     394              : 
     395          907 :     modification.commit(ctx).await?;
     396           10 :     Ok(())
     397           12 : }
     398              : 
     399            2 : pub async fn import_wal_from_tar(
     400            2 :     tline: &Timeline,
     401            2 :     reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     402            2 :     start_lsn: Lsn,
     403            2 :     end_lsn: Lsn,
     404            2 :     ctx: &RequestContext,
     405            2 : ) -> Result<()> {
     406            2 :     // Set up walingest mutable state
     407            2 :     let mut waldecoder = WalStreamDecoder::new(start_lsn, tline.pg_version);
     408            2 :     let mut segno = start_lsn.segment_number(WAL_SEGMENT_SIZE);
     409            2 :     let mut offset = start_lsn.segment_offset(WAL_SEGMENT_SIZE);
     410            2 :     let mut last_lsn = start_lsn;
     411            2 :     let mut walingest = WalIngest::new(tline, start_lsn, ctx).await?;
     412              : 
     413              :     // Ingest wal until end_lsn
     414            2 :     info!("importing wal until {}", end_lsn);
     415            2 :     let mut pg_wal_tar = Archive::new(reader);
     416            2 :     let mut pg_wal_entries = pg_wal_tar.entries()?;
     417            4 :     while last_lsn <= end_lsn {
     418            2 :         let bytes = {
     419            2 :             let mut entry = pg_wal_entries
     420            2 :                 .next()
     421            2 :                 .await
     422            2 :                 .ok_or_else(|| anyhow::anyhow!("expected more wal"))??;
     423            2 :             let header = entry.header();
     424            2 :             let file_path = header.path()?.into_owned();
     425            2 : 
     426            2 :             match header.entry_type() {
     427              :                 tokio_tar::EntryType::Regular => {
     428              :                     // FIXME: assume postgresql tli 1 for now
     429            2 :                     let expected_filename = XLogFileName(1, segno, WAL_SEGMENT_SIZE);
     430            2 :                     let file_name = file_path
     431            2 :                         .file_name()
     432            2 :                         .expect("missing wal filename")
     433            2 :                         .to_string_lossy();
     434            2 :                     ensure!(expected_filename == file_name);
     435              : 
     436            0 :                     debug!("processing wal file {:?}", file_path);
     437         2257 :                     read_all_bytes(&mut entry).await?
     438              :                 }
     439              :                 tokio_tar::EntryType::Directory => {
     440            0 :                     debug!("directory {:?}", file_path);
     441            0 :                     continue;
     442              :                 }
     443              :                 _ => {
     444            0 :                     bail!(
     445            0 :                         "entry {} in WAL tar archive is of unexpected type: {:?}",
     446            0 :                         file_path.display(),
     447            0 :                         header.entry_type()
     448            0 :                     );
     449              :                 }
     450              :             }
     451              :         };
     452              : 
     453            2 :         waldecoder.feed_bytes(&bytes[offset..]);
     454            2 : 
     455            2 :         let mut modification = tline.begin_modification(last_lsn);
     456            2 :         let mut decoded = DecodedWALRecord::default();
     457           10 :         while last_lsn <= end_lsn {
     458            8 :             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
     459            8 :                 walingest
     460            8 :                     .ingest_record(recdata, lsn, &mut modification, &mut decoded, ctx)
     461            0 :                     .await?;
     462            8 :                 modification.commit(ctx).await?;
     463            8 :                 last_lsn = lsn;
     464              : 
     465            0 :                 debug!("imported record at {} (end {})", lsn, end_lsn);
     466            0 :             }
     467              :         }
     468              : 
     469            0 :         debug!("imported records up to {}", last_lsn);
     470            2 :         segno += 1;
     471            2 :         offset = 0;
     472              :     }
     473              : 
     474            2 :     if last_lsn != start_lsn {
     475            2 :         info!("reached end of WAL at {}", last_lsn);
     476              :     } else {
     477            0 :         info!("there was no WAL to import at {}", last_lsn);
     478              :     }
     479              : 
     480              :     // Log any extra unused files
     481            2 :     while let Some(e) = pg_wal_entries.next().await {
     482            0 :         let entry = e?;
     483            0 :         let header = entry.header();
     484            0 :         let file_path = header.path()?.into_owned();
     485            0 :         info!("skipping {:?}", file_path);
     486              :     }
     487              : 
     488            2 :     Ok(())
     489            2 : }
     490              : 
     491       576178 : async fn import_file(
     492       576178 :     modification: &mut DatadirModification<'_>,
     493       576178 :     file_path: &Path,
     494       576178 :     reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     495       576178 :     len: usize,
     496       576178 :     ctx: &RequestContext,
     497       576178 : ) -> Result<Option<ControlFileData>> {
     498       576178 :     let file_name = match file_path.file_name() {
     499       576178 :         Some(name) => name.to_string_lossy(),
     500            0 :         None => return Ok(None),
     501              :     };
     502              : 
     503       576178 :     if file_name.starts_with('.') {
     504              :         // tar archives on macOs, created without COPYFILE_DISABLE=1 env var
     505              :         // will contain "fork files", skip them.
     506            0 :         return Ok(None);
     507       576178 :     }
     508       576178 : 
     509       576178 :     if file_path.starts_with("global") {
     510        34616 :         let spcnode = postgres_ffi::pg_constants::GLOBALTABLESPACE_OID;
     511        34616 :         let dbnode = 0;
     512        34616 : 
     513        34616 :         match file_name.as_ref() {
     514        34616 :             "pg_control" => {
     515         5840 :                 let bytes = read_all_bytes(reader).await?;
     516              : 
     517              :                 // Extract the checkpoint record and import it separately.
     518          606 :                 let pg_control = ControlFileData::decode(&bytes[..])?;
     519          606 :                 let checkpoint_bytes = pg_control.checkPointCopy.encode()?;
     520          606 :                 modification.put_checkpoint(checkpoint_bytes)?;
     521            0 :                 debug!("imported control file");
     522              : 
     523              :                 // Import it as ControlFile
     524          606 :                 modification.put_control_file(bytes)?;
     525          606 :                 return Ok(Some(pg_control));
     526              :             }
     527        34010 :             "pg_filenode.map" => {
     528         3496 :                 let bytes = read_all_bytes(reader).await?;
     529          607 :                 modification
     530          607 :                     .put_relmap_file(spcnode, dbnode, bytes, ctx)
     531            0 :                     .await?;
     532            0 :                 debug!("imported relmap file")
     533              :             }
     534        33403 :             "PG_VERSION" => {
     535            0 :                 debug!("ignored PG_VERSION file");
     536              :             }
     537              :             _ => {
     538        70474 :                 import_rel(modification, file_path, spcnode, dbnode, reader, len, ctx).await?;
     539            0 :                 debug!("imported rel creation");
     540              :             }
     541              :         }
     542       541562 :     } else if file_path.starts_with("base") {
     543       535484 :         let spcnode = pg_constants::DEFAULTTABLESPACE_OID;
     544       535484 :         let dbnode: u32 = file_path
     545       535484 :             .iter()
     546       535484 :             .nth(1)
     547       535484 :             .expect("invalid file path, expected dbnode")
     548       535484 :             .to_string_lossy()
     549       535484 :             .parse()?;
     550              : 
     551       535484 :         match file_name.as_ref() {
     552       535484 :             "pg_filenode.map" => {
     553        10580 :                 let bytes = read_all_bytes(reader).await?;
     554         1821 :                 modification
     555         1821 :                     .put_relmap_file(spcnode, dbnode, bytes, ctx)
     556            0 :                     .await?;
     557            0 :                 debug!("imported relmap file")
     558              :             }
     559       533663 :             "PG_VERSION" => {
     560            0 :                 debug!("ignored PG_VERSION file");
     561              :             }
     562              :             _ => {
     563      2296577 :                 import_rel(modification, file_path, spcnode, dbnode, reader, len, ctx).await?;
     564            0 :                 debug!("imported rel creation");
     565              :             }
     566              :         }
     567         6078 :     } else if file_path.starts_with("pg_xact") {
     568          608 :         let slru = SlruKind::Clog;
     569          608 : 
     570         1181 :         import_slru(modification, slru, file_path, reader, len, ctx).await?;
     571            0 :         debug!("imported clog slru");
     572         5470 :     } else if file_path.starts_with("pg_multixact/offsets") {
     573          607 :         let slru = SlruKind::MultiXactOffsets;
     574          607 : 
     575         1177 :         import_slru(modification, slru, file_path, reader, len, ctx).await?;
     576            0 :         debug!("imported multixact offsets slru");
     577         4863 :     } else if file_path.starts_with("pg_multixact/members") {
     578          607 :         let slru = SlruKind::MultiXactMembers;
     579          607 : 
     580         1184 :         import_slru(modification, slru, file_path, reader, len, ctx).await?;
     581            0 :         debug!("imported multixact members slru");
     582         4256 :     } else if file_path.starts_with("pg_twophase") {
     583            0 :         let xid = u32::from_str_radix(file_name.as_ref(), 16)?;
     584              : 
     585            0 :         let bytes = read_all_bytes(reader).await?;
     586            0 :         modification
     587            0 :             .put_twophase_file(xid, Bytes::copy_from_slice(&bytes[..]), ctx)
     588            0 :             .await?;
     589            0 :         debug!("imported twophase file");
     590         4256 :     } else if file_path.starts_with("pg_wal") {
     591            0 :         debug!("found wal file in base section. ignore it");
     592         4249 :     } else if file_path.starts_with("zenith.signal") {
     593              :         // Parse zenith signal file to set correct previous LSN
     594            7 :         let bytes = read_all_bytes(reader).await?;
     595              :         // zenith.signal format is "PREV LSN: prev_lsn"
     596              :         // TODO write serialization and deserialization in the same place.
     597            7 :         let zenith_signal = std::str::from_utf8(&bytes)?.trim();
     598            7 :         let prev_lsn = match zenith_signal {
     599            7 :             "PREV LSN: none" => Lsn(0),
     600            7 :             "PREV LSN: invalid" => Lsn(0),
     601            7 :             other => {
     602            7 :                 let split = other.split(':').collect::<Vec<_>>();
     603            7 :                 split[1]
     604            7 :                     .trim()
     605            7 :                     .parse::<Lsn>()
     606            7 :                     .context("can't parse zenith.signal")?
     607              :             }
     608              :         };
     609              : 
     610              :         // zenith.signal is not necessarily the last file, that we handle
     611              :         // but it is ok to call `finish_write()`, because final `modification.commit()`
     612              :         // will update lsn once more to the final one.
     613            7 :         let writer = modification.tline.writer().await;
     614            7 :         writer.finish_write(prev_lsn);
     615              : 
     616            0 :         debug!("imported zenith signal {}", prev_lsn);
     617         4242 :     } else if file_path.starts_with("pg_tblspc") {
     618              :         // TODO Backups exported from neon won't have pg_tblspc, but we will need
     619              :         // this to import arbitrary postgres databases.
     620            0 :         bail!("Importing pg_tblspc is not implemented");
     621              :     } else {
     622            0 :         debug!(
     623            0 :             "ignoring unrecognized file \"{}\" in tar archive",
     624            0 :             file_path.display()
     625            0 :         );
     626              :     }
     627              : 
     628       575572 :     Ok(None)
     629       576178 : }
     630              : 
     631         3043 : async fn read_all_bytes(reader: &mut (impl AsyncRead + Unpin)) -> Result<Bytes> {
     632         3043 :     let mut buf: Vec<u8> = vec![];
     633        22173 :     reader.read_to_end(&mut buf).await?;
     634         3043 :     Ok(Bytes::from(buf))
     635         3043 : }
     636              : 
     637          567 : pub async fn create_tar_zst(pgdata_path: &Utf8Path, tmp_path: &Utf8Path) -> Result<(File, u64)> {
     638          567 :     let file = OpenOptions::new()
     639          567 :         .create(true)
     640          567 :         .truncate(true)
     641          567 :         .read(true)
     642          567 :         .write(true)
     643          567 :         .open(&tmp_path)
     644          877 :         .await
     645          567 :         .with_context(|| format!("tempfile creation {tmp_path}"))?;
     646              : 
     647          567 :     let mut paths = Vec::new();
     648       553456 :     for entry in WalkDir::new(pgdata_path) {
     649       553456 :         let entry = entry?;
     650       553456 :         let metadata = entry.metadata().expect("error getting dir entry metadata");
     651       553456 :         // Also allow directories so that we also get empty directories
     652       553456 :         if !(metadata.is_file() || metadata.is_dir()) {
     653            0 :             continue;
     654       553456 :         }
     655       553456 :         let path = entry.into_path();
     656       553456 :         paths.push(path);
     657              :     }
     658              :     // Do a sort to get a more consistent listing
     659          567 :     paths.sort_unstable();
     660          567 :     let zstd = ZstdEncoder::with_quality_and_params(
     661          567 :         file,
     662          567 :         Level::Default,
     663          567 :         &[CParameter::enable_long_distance_matching(true)],
     664          567 :     );
     665          567 :     let mut builder = Builder::new(zstd);
     666          567 :     // Use reproducible header mode
     667          567 :     builder.mode(HeaderMode::Deterministic);
     668       553056 :     for path in paths {
     669       552490 :         let rel_path = path.strip_prefix(pgdata_path)?;
     670       552490 :         if rel_path.is_empty() {
     671              :             // The top directory should not be compressed,
     672              :             // the tar crate doesn't like that
     673          567 :             continue;
     674       551923 :         }
     675      4660913 :         builder.append_path_with_name(&path, rel_path).await?;
     676              :     }
     677          566 :     let mut zstd = builder.into_inner().await?;
     678          568 :     zstd.shutdown().await?;
     679          566 :     let mut compressed = zstd.into_inner();
     680          566 :     let compressed_len = compressed.metadata().await?.len();
     681          566 :     const INITDB_TAR_ZST_WARN_LIMIT: u64 = 2 * 1024 * 1024;
     682          566 :     if compressed_len > INITDB_TAR_ZST_WARN_LIMIT {
     683            0 :         warn!("compressed {INITDB_PATH} size of {compressed_len} is above limit {INITDB_TAR_ZST_WARN_LIMIT}.");
     684          566 :     }
     685          566 :     compressed.seek(SeekFrom::Start(0)).await?;
     686          566 :     Ok((compressed, compressed_len))
     687          566 : }
     688              : 
     689            4 : pub async fn extract_tar_zst(
     690            4 :     pgdata_path: &Utf8Path,
     691            4 :     tar_zst: impl AsyncBufRead + Unpin,
     692            4 : ) -> Result<()> {
     693            4 :     let tar = Box::pin(ZstdDecoder::new(tar_zst));
     694            4 :     let mut archive = Archive::new(tar);
     695        20343 :     archive.unpack(pgdata_path).await?;
     696            4 :     Ok(())
     697            4 : }
        

Generated by: LCOV version 2.1-beta