graft_kernel/rt/

runtime.rs

use std::{sync::Arc, time::Duration};

use bytestring::ByteString;
use culprit::ResultExt;
use graft_core::{
    PageCount, PageIdx, SegmentId, VolumeId, checksum::Checksum, commit::SegmentIdx, lsn::LSN,
    page::Page, pageset::PageSet,
};
use tracing::Instrument;
use tryiter::TryIteratorExt;

use crate::{
    KernelErr,
    graft::{Graft, GraftStatus},
    graft_reader::GraftReader,
    graft_writer::GraftWriter,
    remote::Remote,
    rt::{
        action::{Action, FetchSegment, FetchVolume, HydrateSnapshot, RemoteCommit},
        task::{autosync::AutosyncTask, supervise},
    },
    snapshot::Snapshot,
};

use crate::local::fjall_storage::FjallStorage;

type Result<T> = culprit::Result<T, KernelErr>;

#[derive(Clone, Debug)]
pub struct Runtime {
    inner: Arc<RuntimeInner>,
}

#[derive(Debug)]
struct RuntimeInner {
    tokio: tokio::runtime::Handle,
    storage: Arc<FjallStorage>,
    remote: Arc<Remote>,
}

impl Runtime {
    /// Create a Graft `Runtime` wrapping the provided Tokio runtime handle.
    pub fn new(
        tokio_rt: tokio::runtime::Handle,
        remote: Arc<Remote>,
        storage: Arc<FjallStorage>,
        autosync: Option<Duration>,
    ) -> Runtime {
        // spin up background tasks as needed
        if let Some(interval) = autosync {
            let _guard = tokio_rt.enter();
            let mut ticker = tokio::time::interval(interval);
            ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
            tokio_rt.spawn(supervise(
                storage.clone(),
                remote.clone(),
                AutosyncTask::new(ticker),
            ));
        }
        Runtime {
            inner: Arc::new(RuntimeInner { tokio: tokio_rt, storage, remote }),
        }
    }

    pub(crate) fn storage(&self) -> &FjallStorage {
        &self.inner.storage
    }

    pub(crate) fn create_staged_segment(&self) -> SegmentIdx {
        // TODO: need to keep track of staged segments in memory to prevent the GC from clearing them
        SegmentIdx::new(SegmentId::random(), PageSet::default())
    }

    pub(crate) fn read_page(&self, snapshot: &Snapshot, pageidx: PageIdx) -> Result<Page> {
        let reader = self.storage().read();
        if let Some(commit) = reader.search_page(snapshot, pageidx).or_into_ctx()? {
            let idx = commit
                .segment_idx()
                .expect("BUG: commit claims to contain pageidx");

            if let Some(page) = reader.read_page(idx.sid().clone(), pageidx).or_into_ctx()? {
                return Ok(page);
            }

            // fallthrough to loading the page from the remote
            let range = idx
                .frame_for_pageidx(pageidx)
                .expect("BUG: no frame for pageidx");

            // fetch the segment frame containing the page
            self.run_action(FetchSegment { range })?;

            // now that we've fetched the segment, read the page again using a
            // fresh storage reader
            Ok(self
                .storage()
                .read()
                .read_page(idx.sid.clone(), pageidx)
                .or_into_ctx()?
                .expect("BUG: page not found after fetching"))
        } else {
            Ok(Page::EMPTY)
        }
    }

    fn run_action<A: Action>(&self, action: A) -> Result<()> {
        let span = tracing::debug_span!("Action::run", ?action);

        self.inner.tokio.block_on(
            action
                .run(&self.inner.storage, &self.inner.remote)
                .instrument(span),
        )
    }
}

// tag methods
impl Runtime {
    pub fn tag_iter(&self) -> impl Iterator<Item = Result<(ByteString, VolumeId)>> {
        self.storage()
            .read()
            .iter_tags()
            .map_err(|err| err.map_ctx(KernelErr::from))
    }

    pub fn tag_exists(&self, name: &str) -> Result<bool> {
        self.storage().read().tag_exists(name).or_into_ctx()
    }

    pub fn tag_get(&self, tag: &str) -> Result<Option<VolumeId>> {
        self.storage().read().get_tag(tag).or_into_ctx()
    }

    /// retrieves the `VolumeId` for a tag, replacing it with the provided graft
    pub fn tag_replace(&self, tag: &str, graft: VolumeId) -> Result<Option<VolumeId>> {
        self.storage()
            .read_write()
            .tag_replace(tag, graft)
            .or_into_ctx()
    }

    pub fn tag_delete(&self, tag: &str) -> Result<()> {
        self.storage().tag_delete(tag).or_into_ctx()
    }
}

// graft methods
impl Runtime {
    pub fn graft_iter(&self) -> impl Iterator<Item = Result<Graft>> {
        self.storage()
            .read()
            .iter_grafts()
            .map_err(|err| err.map_ctx(KernelErr::from))
    }

    pub fn graft_exists(&self, graft: &VolumeId) -> Result<bool> {
        self.storage().read().graft_exists(graft).or_into_ctx()
    }

    /// opens a graft. if either the graft's `VolumeId` or the remote's `VolumeId`
    /// are missing, they will be randomly generated. If the graft already
    /// exists, this function will fail if its remote doesn't match.
    pub fn graft_open(&self, graft: Option<VolumeId>, remote: Option<VolumeId>) -> Result<Graft> {
        self.storage()
            .read_write()
            .graft_open(graft, remote)
            .or_into_ctx()
    }

    /// creates a new graft by forking an existing snapshot
    pub fn graft_from_snapshot(&self, snapshot: &Snapshot) -> Result<Graft> {
        self.storage().graft_from_snapshot(snapshot).or_into_ctx()
    }

    /// retrieves an existing graft. returns `LogicalErr::GraftNotFound` if missing
    pub fn graft_get(&self, graft: &VolumeId) -> Result<Graft> {
        self.storage().read().graft(graft).or_into_ctx()
    }

    /// removes a graft but leaves the underlying volumes in place
    pub fn graft_delete(&self, graft: &VolumeId) -> Result<()> {
        self.storage().graft_delete(graft).or_into_ctx()
    }

    /// fetches the latest changes to the remote and then pulls them into the
    /// graft
    pub fn graft_pull(&self, graft: VolumeId) -> Result<()> {
        let graft = self.inner.storage.read().graft(&graft).or_into_ctx()?;
        self.fetch_volume(graft.remote, None)?;
        self.storage()
            .read_write()
            .sync_remote_to_local(graft.local)
            .or_into_ctx()
    }

    pub fn graft_push(&self, graft: VolumeId) -> Result<()> {
        self.run_action(RemoteCommit { graft })
    }

    pub fn graft_status(&self, graft: &VolumeId) -> Result<GraftStatus> {
        let reader = self.storage().read();
        let state = reader.graft(graft).or_into_ctx()?;
        let latest_local = reader.latest_lsn(&state.local).or_into_ctx()?;
        let latest_remote = reader.latest_lsn(&state.remote).or_into_ctx()?;
        Ok(state.status(latest_local, latest_remote))
    }

    pub fn graft_snapshot(&self, graft: &VolumeId) -> Result<Snapshot> {
        self.storage().read().snapshot(graft).or_into_ctx()
    }

    pub fn graft_reader(&self, graft: VolumeId) -> Result<GraftReader> {
        let snapshot = self.graft_snapshot(&graft)?;
        Ok(GraftReader::new(self.clone(), graft, snapshot))
    }

    pub fn graft_writer(&self, graft: VolumeId) -> Result<GraftWriter> {
        let snapshot = self.graft_snapshot(&graft)?;
        let page_count = self.snapshot_pages(&snapshot)?;
        Ok(GraftWriter::new(self.clone(), graft, snapshot, page_count))
    }
}

// volume methods
impl Runtime {
    pub fn fetch_volume(&self, vid: VolumeId, max_lsn: Option<LSN>) -> Result<()> {
        self.run_action(FetchVolume { vid, max_lsn })
    }
}

// snapshot methods
impl Runtime {
    /// returns the total number of pages in the snapshot
    pub fn snapshot_pages(&self, snapshot: &Snapshot) -> Result<PageCount> {
        if let Some((vid, lsn)) = snapshot.head() {
            Ok(self
                .storage()
                .read()
                .page_count(vid, lsn)
                .or_into_ctx()?
                .expect("BUG: missing head commit for snapshot"))
        } else {
            Ok(PageCount::ZERO)
        }
    }

    pub fn snapshot_is_latest(&self, graft: &VolumeId, snapshot: &Snapshot) -> Result<bool> {
        self.storage()
            .read()
            .is_latest_snapshot(graft, snapshot)
            .or_into_ctx()
    }

    /// returns the checksum of the snapshot
    pub fn snapshot_checksum(&self, snapshot: &Snapshot) -> Result<Checksum> {
        self.storage().read().checksum(snapshot).or_into_ctx()
    }

    pub fn snapshot_missing_pages(&self, snapshot: &Snapshot) -> Result<PageSet> {
        let missing_frames = self
            .storage()
            .read()
            .find_missing_frames(snapshot)
            .or_into_ctx()?;
        // merge missing_frames into a single PageSet
        Ok(missing_frames
            .into_iter()
            .fold(PageSet::EMPTY, |mut pageset, frame| {
                pageset |= frame.pageset;
                pageset
            }))
    }

    pub fn snapshot_hydrate(&self, snapshot: Snapshot) -> Result<()> {
        self.run_action(HydrateSnapshot { snapshot })
    }
}

#[cfg(test)]
mod tests {
    use std::{sync::Arc, time::Duration};

    use graft_core::{PageIdx, VolumeId, page::Page};
    use tokio::time::sleep;

    use crate::{
        graft_reader::GraftRead, graft_writer::GraftWrite, local::fjall_storage::FjallStorage,
        remote::RemoteConfig, rt::runtime::Runtime,
    };

    #[graft_test::test]
    fn runtime_sanity() {
        let tokio_rt = tokio::runtime::Builder::new_current_thread()
            .start_paused(true)
            .enable_all()
            .build()
            .unwrap();

        let remote = Arc::new(RemoteConfig::Memory.build().unwrap());
        let storage = Arc::new(FjallStorage::open_temporary().unwrap());
        let runtime = Runtime::new(
            tokio_rt.handle().clone(),
            remote.clone(),
            storage,
            Some(Duration::from_secs(1)),
        );

        let remote_vid = VolumeId::random();
        let graft = runtime
            .graft_open(None, Some(remote_vid.clone()))
            .unwrap()
            .local;

        assert_eq!(runtime.graft_status(&graft).unwrap().to_string(), "_ r_",);

        // sanity check volume writer semantics
        let mut writer = runtime.graft_writer(graft.clone()).unwrap();
        for i in [1u8, 2, 5, 9] {
            let pageidx = PageIdx::must_new(i as u32);
            let page = Page::test_filled(i);
            writer.write_page(pageidx, page.clone()).unwrap();
            assert_eq!(writer.read_page(pageidx).unwrap(), page);
        }
        writer.commit().unwrap();

        assert_eq!(runtime.graft_status(&graft).unwrap().to_string(), "1 r_",);

        // sanity check volume reader semantics
        let reader = runtime.graft_reader(graft.clone()).unwrap();
        tracing::debug!("got snapshot {:?}", reader.snapshot());
        for i in [1u8, 2, 5, 9] {
            let pageidx = PageIdx::must_new(i as u32);
            let page = Page::test_filled(i);
            assert_eq!(
                reader.read_page(pageidx).unwrap().into_bytes(),
                page.into_bytes()
            );
        }

        // create a second runtime connected to the same remote
        let storage = Arc::new(FjallStorage::open_temporary().unwrap());
        let runtime_2 = Runtime::new(
            tokio_rt.handle().clone(),
            remote.clone(),
            storage,
            Some(Duration::from_secs(1)),
        );

        // open the same remote volume in the second runtime
        let graft_2 = runtime_2.graft_open(None, Some(remote_vid)).unwrap().local;

        // let both runtimes run for a little while
        tokio_rt.block_on(async {
            // this sleep lets tokio advance time, allowing the runtime to flush all its jobs
            sleep(Duration::from_secs(5)).await;
            let tree = remote.testonly_format_tree().await;
            tracing::info!("remote tree\n{tree}")
        });

        assert_eq!(runtime.graft_status(&graft).unwrap().to_string(), "1 r1",);
        assert_eq!(
            runtime_2.graft_status(&graft_2).unwrap().to_string(),
            "_ r1",
        );

        // sanity check volume reader semantics in the second runtime
        let reader_2 = runtime_2.graft_reader(graft_2.clone()).unwrap();
        let task = tokio_rt.spawn_blocking(move || {
            for i in [1u8, 2, 5, 9] {
                let pageidx = PageIdx::must_new(i as u32);
                tracing::info!("checking page {pageidx}");
                let expected = Page::test_filled(i);
                let actual = reader_2.read_page(pageidx).unwrap();
                assert_eq!(expected, actual, "read unexpected page contents");
            }
        });
        tokio_rt.block_on(task).unwrap();

        // now write to the second volume, and sync back to the first
        let mut writer_2 = runtime_2.graft_writer(graft_2.clone()).unwrap();
        for i in [3u8, 4, 5, 7] {
            let pageidx = PageIdx::must_new(i as u32);
            let page = Page::test_filled(i + 10);
            writer_2.write_page(pageidx, page.clone()).unwrap();
            assert_eq!(writer_2.read_page(pageidx).unwrap(), page);
        }
        writer_2.commit().unwrap();

        // let both runtimes run for a little while
        tokio_rt.block_on(async {
            // this sleep lets tokio advance time, allowing the runtime to flush all its jobs
            sleep(Duration::from_secs(5)).await;
            let tree = remote.testonly_format_tree().await;
            tracing::info!("remote tree\n{tree}")
        });

        assert_eq!(runtime.graft_status(&graft).unwrap().to_string(), "1 r2",);
        assert_eq!(
            runtime_2.graft_status(&graft_2).unwrap().to_string(),
            "1 r2",
        );

        // sanity check volume reader semantics in the first runtime
        let reader = runtime.graft_reader(graft.clone()).unwrap();
        let task = tokio_rt.spawn_blocking(move || {
            for i in [3u8, 4, 5, 7] {
                let pageidx = PageIdx::must_new(i as u32);
                tracing::info!("checking page {pageidx}");
                let expected = Page::test_filled(i + 10);
                let actual = reader.read_page(pageidx).unwrap();
                assert_eq!(expected, actual, "read unexpected page contents");
            }
        });
        tokio_rt.block_on(task).unwrap();
    }
}