mito2/manifest/storage/

checkpoint.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::sync::Arc;

use common_datasource::compression::CompressionType;
use common_telemetry::{debug, warn};
use object_store::{ErrorKind, ObjectStore};
use serde::{Deserialize, Serialize};
use snafu::{OptionExt, ResultExt};
use store_api::ManifestVersion;

use crate::cache::manifest_cache::ManifestCache;
use crate::error::{
    CompressObjectSnafu, DecompressObjectSnafu, ManifestCheckpointVanishedSnafu, OpenDalSnafu,
    Result, SerdeJsonSnafu, Utf8Snafu,
};
use crate::manifest::storage::size_tracker::Tracker;
use crate::manifest::storage::utils::{get_from_cache_decoded, write_and_put_cache};
use crate::manifest::storage::{
    FALL_BACK_COMPRESS_TYPE, LAST_CHECKPOINT_FILE, checkpoint_checksum, checkpoint_file, gen_path,
    verify_checksum,
};

#[derive(Serialize, Deserialize, Debug)]
pub(crate) struct CheckpointMetadata {
    pub size: usize,
    /// The latest version this checkpoint contains.
    pub version: ManifestVersion,
    pub checksum: Option<u32>,
    pub extend_metadata: HashMap<String, String>,
}

impl CheckpointMetadata {
    fn encode(&self) -> Result<Vec<u8>> {
        Ok(serde_json::to_string(self)
            .context(SerdeJsonSnafu)?
            .into_bytes())
    }

    fn decode(bs: &[u8]) -> Result<Self> {
        let data = std::str::from_utf8(bs).context(Utf8Snafu)?;

        serde_json::from_str(data).context(SerdeJsonSnafu)
    }
}

/// Handle checkpoint storage operations.
#[derive(Debug, Clone)]
pub(crate) struct CheckpointStorage<T: Tracker> {
    object_store: ObjectStore,
    compress_type: CompressionType,
    path: String,
    manifest_cache: Option<ManifestCache>,
    size_tracker: Arc<T>,
}

impl<T: Tracker> CheckpointStorage<T> {
    pub fn new(
        path: String,
        object_store: ObjectStore,
        compress_type: CompressionType,
        manifest_cache: Option<ManifestCache>,
        size_tracker: Arc<T>,
    ) -> Self {
        Self {
            object_store,
            compress_type,
            path,
            manifest_cache,
            size_tracker,
        }
    }

    /// Returns the last checkpoint path, because the last checkpoint is not compressed,
    /// so its path name has nothing to do with the compression algorithm used by `ManifestObjectStore`
    pub(crate) fn last_checkpoint_path(&self) -> String {
        format!("{}{}", self.path, LAST_CHECKPOINT_FILE)
    }

    /// Returns the checkpoint file path under the **current** compression algorithm
    fn checkpoint_file_path(&self, version: ManifestVersion) -> String {
        gen_path(&self.path, &checkpoint_file(version), self.compress_type)
    }

    pub(crate) async fn load_checkpoint(
        &mut self,
        metadata: CheckpointMetadata,
    ) -> Result<Option<(ManifestVersion, Vec<u8>)>> {
        let version = metadata.version;
        let path = self.checkpoint_file_path(version);

        // Try to get from cache first
        if let Some(data) =
            get_from_cache_decoded(self.manifest_cache.as_ref(), &path, self.compress_type).await
        {
            verify_checksum(&data, metadata.checksum)?;
            self.record_cached_checkpoint_size(version, &path).await;
            return Ok(Some((version, data)));
        }

        // Due to backward compatibility, it is possible that the user's checkpoint not compressed,
        // so if we don't find file by compressed type. fall back to checkpoint not compressed find again.
        let checkpoint_data = match self.object_store.read(&path).await {
            Ok(checkpoint) => {
                let checkpoint_size = checkpoint.len();
                let decompress_data = self
                    .compress_type
                    .decode(checkpoint.clone())
                    .await
                    .with_context(|_| DecompressObjectSnafu {
                        compress_type: self.compress_type,
                        path: path.clone(),
                    })?;
                verify_checksum(&decompress_data, metadata.checksum)?;
                // set the checkpoint size
                self.size_tracker.record(version, checkpoint_size as u64);
                // Add to cache
                if let Some(cache) = &self.manifest_cache {
                    cache.put_file(path, checkpoint.to_vec()).await;
                }
                Ok(Some(decompress_data))
            }
            Err(e) => {
                if e.kind() == ErrorKind::NotFound {
                    if self.compress_type != FALL_BACK_COMPRESS_TYPE {
                        let fall_back_path = gen_path(
                            &self.path,
                            &checkpoint_file(version),
                            FALL_BACK_COMPRESS_TYPE,
                        );
                        debug!(
                            "Failed to load checkpoint from path: {}, fall back to path: {}",
                            path, fall_back_path
                        );

                        // Try to get fallback from cache first
                        if let Some(data) = get_from_cache_decoded(
                            self.manifest_cache.as_ref(),
                            &fall_back_path,
                            FALL_BACK_COMPRESS_TYPE,
                        )
                        .await
                        {
                            verify_checksum(&data, metadata.checksum)?;
                            self.record_cached_checkpoint_size(version, &fall_back_path)
                                .await;
                            return Ok(Some((version, data)));
                        }

                        match self.object_store.read(&fall_back_path).await {
                            Ok(checkpoint) => {
                                let checkpoint_size = checkpoint.len();
                                let decompress_data = FALL_BACK_COMPRESS_TYPE
                                    .decode(checkpoint.clone())
                                    .await
                                    .with_context(|_| DecompressObjectSnafu {
                                        compress_type: FALL_BACK_COMPRESS_TYPE,
                                        path: fall_back_path.clone(),
                                    })?;
                                verify_checksum(&decompress_data, metadata.checksum)?;
                                self.size_tracker.record(version, checkpoint_size as u64);
                                // Add fallback to cache
                                if let Some(cache) = &self.manifest_cache {
                                    cache.put_file(fall_back_path, checkpoint.to_vec()).await;
                                }
                                Ok(Some(decompress_data))
                            }
                            Err(e) if e.kind() == ErrorKind::NotFound => Ok(None),
                            Err(e) => return Err(e).context(OpenDalSnafu),
                        }
                    } else {
                        Ok(None)
                    }
                } else {
                    Err(e).context(OpenDalSnafu)
                }
            }
        }?;
        Ok(checkpoint_data.map(|data| (version, data)))
    }

    /// Load the latest checkpoint.
    /// Return manifest version and the raw [RegionCheckpoint](crate::manifest::action::RegionCheckpoint) content if any
    pub async fn load_last_checkpoint(&mut self) -> Result<Option<(ManifestVersion, Vec<u8>)>> {
        let Some(metadata) = self.read_last_checkpoint_metadata().await? else {
            return Ok(None);
        };
        self.load_checkpoint(metadata).await
    }

    /// Loads the checkpoint that LIST found at `version` with `compress_type`.
    ///
    /// The checkpoint pointer only provides the checksum, when it points at the same version: the
    /// checkpoint task of a stale writer can leave it at an older one.
    ///
    /// Fails if cleanup removed the checkpoint after the LIST, because the deltas it covered can
    /// be gone too.
    pub(crate) async fn load_listed_checkpoint(
        &self,
        version: ManifestVersion,
        compress_type: CompressionType,
    ) -> Result<(ManifestVersion, Vec<u8>)> {
        let checksum = self
            .read_last_checkpoint_metadata()
            .await?
            .filter(|metadata| metadata.version == version)
            .and_then(|metadata| metadata.checksum);
        let mut storage = Self::new(
            self.path.clone(),
            self.object_store.clone(),
            compress_type,
            self.manifest_cache.clone(),
            self.size_tracker.clone(),
        );
        storage
            .load_checkpoint(CheckpointMetadata {
                size: 0,
                version,
                checksum,
                extend_metadata: HashMap::new(),
            })
            .await?
            .context(ManifestCheckpointVanishedSnafu {
                version,
                manifest_dir: &self.path,
            })
    }

    /// Reads the small checkpoint pointer without loading its payload.
    pub(crate) async fn read_last_checkpoint_metadata(&self) -> Result<Option<CheckpointMetadata>> {
        let last_checkpoint_path = self.last_checkpoint_path();

        // Fetch from remote object store without cache
        let last_checkpoint_data = match self.object_store.read(&last_checkpoint_path).await {
            Ok(data) => data.to_vec(),
            Err(e) if e.kind() == ErrorKind::NotFound => {
                return Ok(None);
            }
            Err(e) => {
                return Err(e).context(OpenDalSnafu)?;
            }
        };

        let checkpoint_metadata = CheckpointMetadata::decode(&last_checkpoint_data)?;

        debug!(
            "Load checkpoint in path: {}, metadata: {:?}",
            last_checkpoint_path, checkpoint_metadata
        );

        Ok(Some(checkpoint_metadata))
    }

    /// Records the remote object size; accounting failures must not prevent recovery.
    async fn record_cached_checkpoint_size(&self, version: ManifestVersion, path: &str) {
        match self.object_store.stat(path).await {
            Ok(metadata) => self.size_tracker.record(version, metadata.content_length()),
            Err(error) => warn!(error; "Failed to account for cached checkpoint {}", path),
        }
    }

    /// Records stored checkpoint bytes, including the legacy uncompressed path.
    pub(crate) async fn record_checkpoint_size(&self, version: ManifestVersion) -> Result<()> {
        let path = self.checkpoint_file_path(version);
        let metadata = match self.object_store.stat(&path).await {
            Ok(metadata) => metadata,
            Err(error)
                if error.kind() == ErrorKind::NotFound
                    && self.compress_type != FALL_BACK_COMPRESS_TYPE =>
            {
                let path = gen_path(
                    &self.path,
                    &checkpoint_file(version),
                    FALL_BACK_COMPRESS_TYPE,
                );
                self.object_store.stat(&path).await.context(OpenDalSnafu)?
            }
            Err(error) => return Err(error).context(OpenDalSnafu),
        };
        self.size_tracker.record(version, metadata.content_length());
        Ok(())
    }

    /// Save the checkpoint manifest file.
    pub(crate) async fn save_checkpoint(
        &self,
        version: ManifestVersion,
        bytes: &[u8],
    ) -> Result<()> {
        let path = self.checkpoint_file_path(version);
        let data = self
            .compress_type
            .encode(bytes)
            .await
            .context(CompressObjectSnafu {
                compress_type: self.compress_type,
                path: &path,
            })?;
        let checkpoint_size = data.len();
        let checksum = checkpoint_checksum(bytes);

        write_and_put_cache(
            &self.object_store,
            self.manifest_cache.as_ref(),
            &path,
            data,
        )
        .await?;

        // Because last checkpoint file only contain size and version, which is tiny, so we don't compress it.
        let last_checkpoint_path = self.last_checkpoint_path();

        let checkpoint_metadata = CheckpointMetadata {
            size: bytes.len(),
            version,
            checksum: Some(checksum),
            extend_metadata: HashMap::new(),
        };

        debug!(
            "Save checkpoint in path: {},  metadata: {:?}",
            last_checkpoint_path, checkpoint_metadata
        );

        let bytes = checkpoint_metadata.encode()?;
        self.object_store
            .write(&last_checkpoint_path, bytes)
            .await
            .context(OpenDalSnafu)?;

        self.size_tracker.record(version, checkpoint_size as u64);
        Ok(())
    }
}

#[cfg(test)]
impl<T: Tracker> CheckpointStorage<T> {
    pub async fn write_last_checkpoint(
        &self,
        version: ManifestVersion,
        bytes: &[u8],
    ) -> Result<()> {
        let path = self.checkpoint_file_path(version);
        let data = self
            .compress_type
            .encode(bytes)
            .await
            .context(CompressObjectSnafu {
                compress_type: self.compress_type,
                path: &path,
            })?;

        let checkpoint_size = data.len();

        self.object_store
            .write(&path, data)
            .await
            .context(OpenDalSnafu)?;

        self.size_tracker.record(version, checkpoint_size as u64);

        let last_checkpoint_path = self.last_checkpoint_path();
        let checkpoint_metadata = CheckpointMetadata {
            size: bytes.len(),
            version,
            checksum: Some(1218259706),
            extend_metadata: HashMap::new(),
        };

        debug!(
            "Rewrite checkpoint in path: {},  metadata: {:?}",
            last_checkpoint_path, checkpoint_metadata
        );

        let bytes = checkpoint_metadata.encode()?;

        // Overwrite the last checkpoint with the modified content
        self.object_store
            .write(&last_checkpoint_path, bytes.clone())
            .await
            .context(OpenDalSnafu)?;
        Ok(())
    }

    pub fn set_compress_type(&mut self, compress_type: CompressionType) {
        self.compress_type = compress_type;
    }
}