LCOV - code coverage report
Current view: top level - safekeeper/src - control_file.rs (source / functions) Coverage Total Hit
Test: 691a4c28fe7169edd60b367c52d448a0a6605f1f.info Lines: 83.7 % 208 174
Test Date: 2024-05-10 13:18:37 Functions: 69.6 % 23 16

            Line data    Source code
       1              : //! Control file serialization, deserialization and persistence.
       2              : 
       3              : use anyhow::{bail, ensure, Context, Result};
       4              : use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
       5              : use camino::Utf8PathBuf;
       6              : use tokio::fs::File;
       7              : use tokio::io::AsyncWriteExt;
       8              : use utils::crashsafe::durable_rename;
       9              : 
      10              : use std::io::Read;
      11              : use std::ops::Deref;
      12              : use std::path::Path;
      13              : use std::time::Instant;
      14              : 
      15              : use crate::control_file_upgrade::upgrade_control_file;
      16              : use crate::metrics::PERSIST_CONTROL_FILE_SECONDS;
      17              : use crate::state::TimelinePersistentState;
      18              : use utils::{bin_ser::LeSer, id::TenantTimelineId};
      19              : 
      20              : use crate::SafeKeeperConf;
      21              : 
      22              : pub const SK_MAGIC: u32 = 0xcafeceefu32;
      23              : pub const SK_FORMAT_VERSION: u32 = 8;
      24              : 
      25              : // contains persistent metadata for safekeeper
      26              : const CONTROL_FILE_NAME: &str = "safekeeper.control";
      27              : // needed to atomically update the state using `rename`
      28              : const CONTROL_FILE_NAME_PARTIAL: &str = "safekeeper.control.partial";
      29              : pub const CHECKSUM_SIZE: usize = std::mem::size_of::<u32>();
      30              : 
      31              : /// Storage should keep actual state inside of it. It should implement Deref
      32              : /// trait to access state fields and have persist method for updating that state.
      33              : #[async_trait::async_trait]
      34              : pub trait Storage: Deref<Target = TimelinePersistentState> {
      35              :     /// Persist safekeeper state on disk and update internal state.
      36              :     async fn persist(&mut self, s: &TimelinePersistentState) -> Result<()>;
      37              : 
      38              :     /// Timestamp of last persist.
      39              :     fn last_persist_at(&self) -> Instant;
      40              : }
      41              : 
      42              : #[derive(Debug)]
      43              : pub struct FileStorage {
      44              :     // save timeline dir to avoid reconstructing it every time
      45              :     timeline_dir: Utf8PathBuf,
      46              :     conf: SafeKeeperConf,
      47              : 
      48              :     /// Last state persisted to disk.
      49              :     state: TimelinePersistentState,
      50              :     /// Not preserved across restarts.
      51              :     last_persist_at: Instant,
      52              : }
      53              : 
      54              : impl FileStorage {
      55              :     /// Initialize storage by loading state from disk.
      56            4 :     pub fn restore_new(ttid: &TenantTimelineId, conf: &SafeKeeperConf) -> Result<FileStorage> {
      57            4 :         let timeline_dir = conf.timeline_dir(ttid);
      58              : 
      59            4 :         let state = Self::load_control_file_conf(conf, ttid)?;
      60              : 
      61            2 :         Ok(FileStorage {
      62            2 :             timeline_dir,
      63            2 :             conf: conf.clone(),
      64            2 :             state,
      65            2 :             last_persist_at: Instant::now(),
      66            2 :         })
      67            4 :     }
      68              : 
      69              :     /// Create file storage for a new timeline, but don't persist it yet.
      70            4 :     pub fn create_new(
      71            4 :         timeline_dir: Utf8PathBuf,
      72            4 :         conf: &SafeKeeperConf,
      73            4 :         state: TimelinePersistentState,
      74            4 :     ) -> Result<FileStorage> {
      75            4 :         let store = FileStorage {
      76            4 :             timeline_dir,
      77            4 :             conf: conf.clone(),
      78            4 :             state,
      79            4 :             last_persist_at: Instant::now(),
      80            4 :         };
      81            4 : 
      82            4 :         Ok(store)
      83            4 :     }
      84              : 
      85              :     /// Check the magic/version in the on-disk data and deserialize it, if possible.
      86            4 :     fn deser_sk_state(buf: &mut &[u8]) -> Result<TimelinePersistentState> {
      87              :         // Read the version independent part
      88            4 :         let magic = ReadBytesExt::read_u32::<LittleEndian>(buf)?;
      89            4 :         if magic != SK_MAGIC {
      90            0 :             bail!(
      91            0 :                 "bad control file magic: {:X}, expected {:X}",
      92            0 :                 magic,
      93            0 :                 SK_MAGIC
      94            0 :             );
      95            4 :         }
      96            4 :         let version = ReadBytesExt::read_u32::<LittleEndian>(buf)?;
      97            4 :         if version == SK_FORMAT_VERSION {
      98            4 :             let res = TimelinePersistentState::des(buf)?;
      99            4 :             return Ok(res);
     100            0 :         }
     101            0 :         // try to upgrade
     102            0 :         upgrade_control_file(buf, version)
     103            4 :     }
     104              : 
     105              :     /// Load control file for given ttid at path specified by conf.
     106            6 :     pub fn load_control_file_conf(
     107            6 :         conf: &SafeKeeperConf,
     108            6 :         ttid: &TenantTimelineId,
     109            6 :     ) -> Result<TimelinePersistentState> {
     110            6 :         let path = conf.timeline_dir(ttid).join(CONTROL_FILE_NAME);
     111            6 :         Self::load_control_file(path)
     112            6 :     }
     113              : 
     114              :     /// Read in the control file.
     115            6 :     pub fn load_control_file<P: AsRef<Path>>(
     116            6 :         control_file_path: P,
     117            6 :     ) -> Result<TimelinePersistentState> {
     118            6 :         let mut control_file = std::fs::OpenOptions::new()
     119            6 :             .read(true)
     120            6 :             .write(true)
     121            6 :             .open(&control_file_path)
     122            6 :             .with_context(|| {
     123            0 :                 format!(
     124            0 :                     "failed to open control file at {}",
     125            0 :                     control_file_path.as_ref().display(),
     126            0 :                 )
     127            6 :             })?;
     128              : 
     129            6 :         let mut buf = Vec::new();
     130            6 :         control_file
     131            6 :             .read_to_end(&mut buf)
     132            6 :             .context("failed to read control file")?;
     133              : 
     134            6 :         let calculated_checksum = crc32c::crc32c(&buf[..buf.len() - CHECKSUM_SIZE]);
     135              : 
     136            6 :         let expected_checksum_bytes: &[u8; CHECKSUM_SIZE] =
     137            6 :             buf[buf.len() - CHECKSUM_SIZE..].try_into()?;
     138            6 :         let expected_checksum = u32::from_le_bytes(*expected_checksum_bytes);
     139            6 : 
     140            6 :         ensure!(
     141            6 :             calculated_checksum == expected_checksum,
     142            2 :             format!(
     143            2 :                 "safekeeper control file checksum mismatch: expected {} got {}",
     144            2 :                 expected_checksum, calculated_checksum
     145            2 :             )
     146              :         );
     147              : 
     148            4 :         let state = FileStorage::deser_sk_state(&mut &buf[..buf.len() - CHECKSUM_SIZE])
     149            4 :             .with_context(|| {
     150            0 :                 format!(
     151            0 :                     "while reading control file {}",
     152            0 :                     control_file_path.as_ref().display(),
     153            0 :                 )
     154            4 :             })?;
     155            4 :         Ok(state)
     156            6 :     }
     157              : }
     158              : 
     159              : impl Deref for FileStorage {
     160              :     type Target = TimelinePersistentState;
     161              : 
     162            0 :     fn deref(&self) -> &Self::Target {
     163            0 :         &self.state
     164            0 :     }
     165              : }
     166              : 
     167              : #[async_trait::async_trait]
     168              : impl Storage for FileStorage {
     169              :     /// Persists state durably to the underlying storage.
     170              :     ///
     171              :     /// For a description, see <https://lwn.net/Articles/457667/>.
     172            4 :     async fn persist(&mut self, s: &TimelinePersistentState) -> Result<()> {
     173            4 :         let _timer = PERSIST_CONTROL_FILE_SECONDS.start_timer();
     174            4 : 
     175            4 :         // write data to safekeeper.control.partial
     176            4 :         let control_partial_path = self.timeline_dir.join(CONTROL_FILE_NAME_PARTIAL);
     177            4 :         let mut control_partial = File::create(&control_partial_path).await.with_context(|| {
     178            0 :             format!(
     179            0 :                 "failed to create partial control file at: {}",
     180            0 :                 &control_partial_path
     181            0 :             )
     182            4 :         })?;
     183            4 :         let mut buf: Vec<u8> = Vec::new();
     184            4 :         WriteBytesExt::write_u32::<LittleEndian>(&mut buf, SK_MAGIC)?;
     185            4 :         WriteBytesExt::write_u32::<LittleEndian>(&mut buf, SK_FORMAT_VERSION)?;
     186            4 :         s.ser_into(&mut buf)?;
     187            4 : 
     188            4 :         // calculate checksum before resize
     189            4 :         let checksum = crc32c::crc32c(&buf);
     190            4 :         buf.extend_from_slice(&checksum.to_le_bytes());
     191            4 : 
     192            4 :         control_partial.write_all(&buf).await.with_context(|| {
     193            0 :             format!(
     194            0 :                 "failed to write safekeeper state into control file at: {}",
     195            0 :                 control_partial_path
     196            0 :             )
     197            4 :         })?;
     198            4 :         control_partial.flush().await.with_context(|| {
     199            0 :             format!(
     200            0 :                 "failed to flush safekeeper state into control file at: {}",
     201            0 :                 control_partial_path
     202            0 :             )
     203            4 :         })?;
     204            4 : 
     205            4 :         let control_path = self.timeline_dir.join(CONTROL_FILE_NAME);
     206           28 :         durable_rename(&control_partial_path, &control_path, !self.conf.no_sync).await?;
     207            4 : 
     208            4 :         // update internal state
     209            4 :         self.state = s.clone();
     210            4 :         Ok(())
     211            4 :     }
     212              : 
     213            0 :     fn last_persist_at(&self) -> Instant {
     214            0 :         self.last_persist_at
     215            0 :     }
     216              : }
     217              : 
     218              : #[cfg(test)]
     219              : mod test {
     220              :     use super::*;
     221              :     use tokio::fs;
     222              :     use utils::lsn::Lsn;
     223              : 
     224            4 :     fn stub_conf() -> SafeKeeperConf {
     225            4 :         let workdir = camino_tempfile::tempdir().unwrap().into_path();
     226            4 :         SafeKeeperConf {
     227            4 :             workdir,
     228            4 :             ..SafeKeeperConf::dummy()
     229            4 :         }
     230            4 :     }
     231              : 
     232            4 :     async fn load_from_control_file(
     233            4 :         conf: &SafeKeeperConf,
     234            4 :         ttid: &TenantTimelineId,
     235            4 :     ) -> Result<(FileStorage, TimelinePersistentState)> {
     236            4 :         fs::create_dir_all(conf.timeline_dir(ttid))
     237            4 :             .await
     238            4 :             .expect("failed to create timeline dir");
     239            4 :         Ok((
     240            4 :             FileStorage::restore_new(ttid, conf)?,
     241            2 :             FileStorage::load_control_file_conf(conf, ttid)?,
     242              :         ))
     243            4 :     }
     244              : 
     245            4 :     async fn create(
     246            4 :         conf: &SafeKeeperConf,
     247            4 :         ttid: &TenantTimelineId,
     248            4 :     ) -> Result<(FileStorage, TimelinePersistentState)> {
     249            4 :         fs::create_dir_all(conf.timeline_dir(ttid))
     250            4 :             .await
     251            4 :             .expect("failed to create timeline dir");
     252            4 :         let state = TimelinePersistentState::empty();
     253            4 :         let timeline_dir = conf.timeline_dir(ttid);
     254            4 :         let storage = FileStorage::create_new(timeline_dir, conf, state.clone())?;
     255            4 :         Ok((storage, state))
     256            4 :     }
     257              : 
     258              :     #[tokio::test]
     259            2 :     async fn test_read_write_safekeeper_state() {
     260            2 :         let conf = stub_conf();
     261            2 :         let ttid = TenantTimelineId::generate();
     262            2 :         {
     263            2 :             let (mut storage, mut state) =
     264            2 :                 create(&conf, &ttid).await.expect("failed to create state");
     265            2 :             // change something
     266            2 :             state.commit_lsn = Lsn(42);
     267            2 :             storage
     268            2 :                 .persist(&state)
     269           18 :                 .await
     270            2 :                 .expect("failed to persist state");
     271            2 :         }
     272            2 : 
     273            2 :         let (_, state) = load_from_control_file(&conf, &ttid)
     274            2 :             .await
     275            2 :             .expect("failed to read state");
     276            2 :         assert_eq!(state.commit_lsn, Lsn(42));
     277            2 :     }
     278              : 
     279              :     #[tokio::test]
     280            2 :     async fn test_safekeeper_state_checksum_mismatch() {
     281            2 :         let conf = stub_conf();
     282            2 :         let ttid = TenantTimelineId::generate();
     283            2 :         {
     284            2 :             let (mut storage, mut state) =
     285            2 :                 create(&conf, &ttid).await.expect("failed to read state");
     286            2 : 
     287            2 :             // change something
     288            2 :             state.commit_lsn = Lsn(42);
     289            2 :             storage
     290            2 :                 .persist(&state)
     291           18 :                 .await
     292            2 :                 .expect("failed to persist state");
     293            2 :         }
     294            2 :         let control_path = conf.timeline_dir(&ttid).join(CONTROL_FILE_NAME);
     295            2 :         let mut data = fs::read(&control_path).await.unwrap();
     296            2 :         data[0] += 1; // change the first byte of the file to fail checksum validation
     297            2 :         fs::write(&control_path, &data)
     298            2 :             .await
     299            2 :             .expect("failed to write control file");
     300            2 : 
     301            2 :         match load_from_control_file(&conf, &ttid).await {
     302            2 :             Err(err) => assert!(err
     303            2 :                 .to_string()
     304            2 :                 .contains("safekeeper control file checksum mismatch")),
     305            2 :             Ok(_) => panic!("expected error"),
     306            2 :         }
     307            2 :     }
     308              : }
        

Generated by: LCOV version 2.1-beta