LCOV - code coverage report
Current view: top level - safekeeper/src - control_file.rs (source / functions) Coverage Total Hit
Test: ec3669e3e2c475fa8f629fda7ec6b3e52960c08d.info Lines: 83.5 % 206 172
Test Date: 2024-06-19 16:48:34 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::{Utf8Path, 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::metrics::PERSIST_CONTROL_FILE_SECONDS;
      16              : use crate::state::TimelinePersistentState;
      17              : use crate::{control_file_upgrade::upgrade_control_file, timeline::get_timeline_dir};
      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              : pub 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              :     no_sync: bool,
      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 = get_timeline_dir(conf, ttid);
      58            4 :         let state = Self::load_control_file_from_dir(&timeline_dir)?;
      59              : 
      60            2 :         Ok(FileStorage {
      61            2 :             timeline_dir,
      62            2 :             no_sync: conf.no_sync,
      63            2 :             state,
      64            2 :             last_persist_at: Instant::now(),
      65            2 :         })
      66            4 :     }
      67              : 
      68              :     /// Create file storage for a new timeline, but don't persist it yet.
      69            4 :     pub fn create_new(
      70            4 :         timeline_dir: Utf8PathBuf,
      71            4 :         conf: &SafeKeeperConf,
      72            4 :         state: TimelinePersistentState,
      73            4 :     ) -> Result<FileStorage> {
      74            4 :         let store = FileStorage {
      75            4 :             timeline_dir,
      76            4 :             no_sync: conf.no_sync,
      77            4 :             state,
      78            4 :             last_persist_at: Instant::now(),
      79            4 :         };
      80            4 : 
      81            4 :         Ok(store)
      82            4 :     }
      83              : 
      84              :     /// Check the magic/version in the on-disk data and deserialize it, if possible.
      85            4 :     fn deser_sk_state(buf: &mut &[u8]) -> Result<TimelinePersistentState> {
      86              :         // Read the version independent part
      87            4 :         let magic = ReadBytesExt::read_u32::<LittleEndian>(buf)?;
      88            4 :         if magic != SK_MAGIC {
      89            0 :             bail!(
      90            0 :                 "bad control file magic: {:X}, expected {:X}",
      91            0 :                 magic,
      92            0 :                 SK_MAGIC
      93            0 :             );
      94            4 :         }
      95            4 :         let version = ReadBytesExt::read_u32::<LittleEndian>(buf)?;
      96            4 :         if version == SK_FORMAT_VERSION {
      97            4 :             let res = TimelinePersistentState::des(buf)?;
      98            4 :             return Ok(res);
      99            0 :         }
     100            0 :         // try to upgrade
     101            0 :         upgrade_control_file(buf, version)
     102            4 :     }
     103              : 
     104              :     /// Load control file from given directory.
     105            6 :     pub fn load_control_file_from_dir(timeline_dir: &Utf8Path) -> Result<TimelinePersistentState> {
     106            6 :         let path = timeline_dir.join(CONTROL_FILE_NAME);
     107            6 :         Self::load_control_file(path)
     108            6 :     }
     109              : 
     110              :     /// Read in the control file.
     111            6 :     pub fn load_control_file<P: AsRef<Path>>(
     112            6 :         control_file_path: P,
     113            6 :     ) -> Result<TimelinePersistentState> {
     114            6 :         let mut control_file = std::fs::OpenOptions::new()
     115            6 :             .read(true)
     116            6 :             .write(true)
     117            6 :             .open(&control_file_path)
     118            6 :             .with_context(|| {
     119            0 :                 format!(
     120            0 :                     "failed to open control file at {}",
     121            0 :                     control_file_path.as_ref().display(),
     122            0 :                 )
     123            6 :             })?;
     124              : 
     125            6 :         let mut buf = Vec::new();
     126            6 :         control_file
     127            6 :             .read_to_end(&mut buf)
     128            6 :             .context("failed to read control file")?;
     129              : 
     130            6 :         let calculated_checksum = crc32c::crc32c(&buf[..buf.len() - CHECKSUM_SIZE]);
     131              : 
     132            6 :         let expected_checksum_bytes: &[u8; CHECKSUM_SIZE] =
     133            6 :             buf[buf.len() - CHECKSUM_SIZE..].try_into()?;
     134            6 :         let expected_checksum = u32::from_le_bytes(*expected_checksum_bytes);
     135            6 : 
     136            6 :         ensure!(
     137            6 :             calculated_checksum == expected_checksum,
     138            2 :             format!(
     139            2 :                 "safekeeper control file checksum mismatch: expected {} got {}",
     140            2 :                 expected_checksum, calculated_checksum
     141            2 :             )
     142              :         );
     143              : 
     144            4 :         let state = FileStorage::deser_sk_state(&mut &buf[..buf.len() - CHECKSUM_SIZE])
     145            4 :             .with_context(|| {
     146            0 :                 format!(
     147            0 :                     "while reading control file {}",
     148            0 :                     control_file_path.as_ref().display(),
     149            0 :                 )
     150            4 :             })?;
     151            4 :         Ok(state)
     152            6 :     }
     153              : }
     154              : 
     155              : impl Deref for FileStorage {
     156              :     type Target = TimelinePersistentState;
     157              : 
     158            0 :     fn deref(&self) -> &Self::Target {
     159            0 :         &self.state
     160            0 :     }
     161              : }
     162              : 
     163              : #[async_trait::async_trait]
     164              : impl Storage for FileStorage {
     165              :     /// Persists state durably to the underlying storage.
     166              :     ///
     167              :     /// For a description, see <https://lwn.net/Articles/457667/>.
     168            4 :     async fn persist(&mut self, s: &TimelinePersistentState) -> Result<()> {
     169            4 :         let _timer = PERSIST_CONTROL_FILE_SECONDS.start_timer();
     170            4 : 
     171            4 :         // write data to safekeeper.control.partial
     172            4 :         let control_partial_path = self.timeline_dir.join(CONTROL_FILE_NAME_PARTIAL);
     173            4 :         let mut control_partial = File::create(&control_partial_path).await.with_context(|| {
     174            0 :             format!(
     175            0 :                 "failed to create partial control file at: {}",
     176            0 :                 &control_partial_path
     177            0 :             )
     178            4 :         })?;
     179            4 :         let mut buf: Vec<u8> = Vec::new();
     180            4 :         WriteBytesExt::write_u32::<LittleEndian>(&mut buf, SK_MAGIC)?;
     181            4 :         WriteBytesExt::write_u32::<LittleEndian>(&mut buf, SK_FORMAT_VERSION)?;
     182            4 :         s.ser_into(&mut buf)?;
     183            4 : 
     184            4 :         // calculate checksum before resize
     185            4 :         let checksum = crc32c::crc32c(&buf);
     186            4 :         buf.extend_from_slice(&checksum.to_le_bytes());
     187            4 : 
     188            4 :         control_partial.write_all(&buf).await.with_context(|| {
     189            0 :             format!(
     190            0 :                 "failed to write safekeeper state into control file at: {}",
     191            0 :                 control_partial_path
     192            0 :             )
     193            4 :         })?;
     194            4 :         control_partial.flush().await.with_context(|| {
     195            0 :             format!(
     196            0 :                 "failed to flush safekeeper state into control file at: {}",
     197            0 :                 control_partial_path
     198            0 :             )
     199            4 :         })?;
     200            4 : 
     201            4 :         let control_path = self.timeline_dir.join(CONTROL_FILE_NAME);
     202           28 :         durable_rename(&control_partial_path, &control_path, !self.no_sync).await?;
     203            4 : 
     204            4 :         // update internal state
     205            4 :         self.state = s.clone();
     206            4 :         Ok(())
     207            4 :     }
     208              : 
     209            0 :     fn last_persist_at(&self) -> Instant {
     210            0 :         self.last_persist_at
     211            0 :     }
     212              : }
     213              : 
     214              : #[cfg(test)]
     215              : mod test {
     216              :     use super::*;
     217              :     use tokio::fs;
     218              :     use utils::lsn::Lsn;
     219              : 
     220            4 :     fn stub_conf() -> SafeKeeperConf {
     221            4 :         let workdir = camino_tempfile::tempdir().unwrap().into_path();
     222            4 :         SafeKeeperConf {
     223            4 :             workdir,
     224            4 :             ..SafeKeeperConf::dummy()
     225            4 :         }
     226            4 :     }
     227              : 
     228            4 :     async fn load_from_control_file(
     229            4 :         conf: &SafeKeeperConf,
     230            4 :         ttid: &TenantTimelineId,
     231            4 :     ) -> Result<(FileStorage, TimelinePersistentState)> {
     232            4 :         let timeline_dir = get_timeline_dir(conf, ttid);
     233            4 :         fs::create_dir_all(&timeline_dir)
     234            4 :             .await
     235            4 :             .expect("failed to create timeline dir");
     236            4 :         Ok((
     237            4 :             FileStorage::restore_new(ttid, conf)?,
     238            2 :             FileStorage::load_control_file_from_dir(&timeline_dir)?,
     239              :         ))
     240            4 :     }
     241              : 
     242            4 :     async fn create(
     243            4 :         conf: &SafeKeeperConf,
     244            4 :         ttid: &TenantTimelineId,
     245            4 :     ) -> Result<(FileStorage, TimelinePersistentState)> {
     246            4 :         let timeline_dir = get_timeline_dir(conf, ttid);
     247            4 :         fs::create_dir_all(&timeline_dir)
     248            4 :             .await
     249            4 :             .expect("failed to create timeline dir");
     250            4 :         let state = TimelinePersistentState::empty();
     251            4 :         let storage = FileStorage::create_new(timeline_dir, conf, state.clone())?;
     252            4 :         Ok((storage, state))
     253            4 :     }
     254              : 
     255              :     #[tokio::test]
     256            2 :     async fn test_read_write_safekeeper_state() {
     257            2 :         let conf = stub_conf();
     258            2 :         let ttid = TenantTimelineId::generate();
     259            2 :         {
     260            2 :             let (mut storage, mut state) =
     261            2 :                 create(&conf, &ttid).await.expect("failed to create state");
     262            2 :             // change something
     263            2 :             state.commit_lsn = Lsn(42);
     264            2 :             storage
     265            2 :                 .persist(&state)
     266           18 :                 .await
     267            2 :                 .expect("failed to persist state");
     268            2 :         }
     269            2 : 
     270            2 :         let (_, state) = load_from_control_file(&conf, &ttid)
     271            2 :             .await
     272            2 :             .expect("failed to read state");
     273            2 :         assert_eq!(state.commit_lsn, Lsn(42));
     274            2 :     }
     275              : 
     276              :     #[tokio::test]
     277            2 :     async fn test_safekeeper_state_checksum_mismatch() {
     278            2 :         let conf = stub_conf();
     279            2 :         let ttid = TenantTimelineId::generate();
     280            2 :         {
     281            2 :             let (mut storage, mut state) =
     282            2 :                 create(&conf, &ttid).await.expect("failed to read state");
     283            2 : 
     284            2 :             // change something
     285            2 :             state.commit_lsn = Lsn(42);
     286            2 :             storage
     287            2 :                 .persist(&state)
     288           18 :                 .await
     289            2 :                 .expect("failed to persist state");
     290            2 :         }
     291            2 :         let control_path = get_timeline_dir(&conf, &ttid).join(CONTROL_FILE_NAME);
     292            2 :         let mut data = fs::read(&control_path).await.unwrap();
     293            2 :         data[0] += 1; // change the first byte of the file to fail checksum validation
     294            2 :         fs::write(&control_path, &data)
     295            2 :             .await
     296            2 :             .expect("failed to write control file");
     297            2 : 
     298            2 :         match load_from_control_file(&conf, &ttid).await {
     299            2 :             Err(err) => assert!(err
     300            2 :                 .to_string()
     301            2 :                 .contains("safekeeper control file checksum mismatch")),
     302            2 :             Ok(_) => panic!("expected error"),
     303            2 :         }
     304            2 :     }
     305              : }
        

Generated by: LCOV version 2.1-beta