mito2/sst/

version.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.

//! SST version.
use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;

use common_time::{TimeToLive, Timestamp};
use store_api::metadata::RegionMetadataRef;
use store_api::storage::{FileId, RegionId};

use crate::sst::file::{FileHandle, FileMeta, FileTimeRange, Level, MAX_LEVEL};
use crate::sst::file_purger::FilePurgerRef;
use crate::sst::primary_key::PrimaryKeyRangeMapper;

/// A version of all SSTs in a region.
#[derive(Debug, Clone)]
pub(crate) struct SstVersion {
    /// SST metadata organized by levels.
    levels: Arc<LevelMetaArray>,
    primary_key_mapper: Arc<PrimaryKeyRangeMapper>,
}

pub(crate) type SstVersionRef = Arc<SstVersion>;

impl SstVersion {
    /// Returns a new [SstVersion].
    pub(crate) fn new(metadata: RegionMetadataRef) -> SstVersion {
        SstVersion {
            levels: Arc::new(new_level_meta_vec()),
            primary_key_mapper: Arc::new(PrimaryKeyRangeMapper::new(metadata)),
        }
    }

    /// Changes the target schema without copying the SST list or rebinding handles.
    pub(crate) fn set_metadata(&mut self, metadata: RegionMetadataRef) {
        self.primary_key_mapper = Arc::new(self.primary_key_mapper.with_metadata(metadata));
    }

    /// Shares the target schema and encoded defaults with comparisons.
    pub(crate) fn primary_key_mapper(&self) -> Arc<PrimaryKeyRangeMapper> {
        self.primary_key_mapper.clone()
    }

    /// Returns a slice to metadatas of all levels.
    pub(crate) fn levels(&self) -> &[LevelMeta] {
        self.levels.as_ref()
    }

    /// Returns the current handle matching the selected file's identity in its immutable level.
    pub(crate) fn file_for_compaction(&self, selected: &FileHandle) -> Option<&FileHandle> {
        let current = self
            .levels
            .get(selected.level() as usize)?
            .files
            .get(&selected.file_id().file_id())?;
        (current.file_id() == selected.file_id()).then_some(current)
    }

    /// Add files to the version. If a file with the same `file_id` already exists,
    /// it will be overwritten with the new file.
    ///
    /// # Panics
    /// Panics if level of [FileMeta] is greater than [MAX_LEVEL].
    pub(crate) fn add_files(
        &mut self,
        file_purger: FilePurgerRef,
        files_to_add: impl Iterator<Item = FileMeta>,
    ) {
        let levels = Arc::make_mut(&mut self.levels);
        for file in files_to_add {
            let level = file.level;
            let new_index_version = file.index_version;
            // If the file already exists, then we should only replace the handle when the index is outdated.
            levels[level as usize]
                .files
                .entry(file.file_id)
                .and_modify(|f| {
                    if *f.meta_ref() == file || f.meta_ref().is_index_up_to_date(&file) {
                        // same file meta or current file handle's index is up-to-date, skip adding
                        if f.index_id().version > new_index_version {
                            // what does it mean for us to see older index version?
                            common_telemetry::warn!(
                                "Adding file with older index version, existing: {:?}, new: {:?}, ignoring new file",
                                f.meta_ref(),
                                file
                            );
                        }
                    } else {
                        // include case like old file have no index or index is outdated
                        *f = f.with_updated_index(file.clone());
                    }
                })
                .or_insert_with(|| {
                    FileHandle::new(file.clone(), file_purger.clone())
                });
        }
    }

    /// Remove files from the version.
    ///
    /// # Panics
    /// Panics if level of [FileMeta] is greater than [MAX_LEVEL].
    pub(crate) fn remove_files(&mut self, files_to_remove: impl Iterator<Item = FileMeta>) {
        let levels = Arc::make_mut(&mut self.levels);
        for file in files_to_remove {
            let level = file.level;
            if let Some(handle) = levels[level as usize].files.remove(&file.file_id) {
                handle.mark_deleted();
            }
        }
    }

    /// Marks all SSTs in this version as deleted.
    pub(crate) fn mark_all_deleted(&self) {
        for level_meta in self.levels.iter() {
            for file_handle in level_meta.files.values() {
                file_handle.mark_deleted();
            }
        }
    }

    /// Returns the number of rows in SST files owned by `region_id`.
    ///
    /// Rows from SST files referenced from other regions, for example after
    /// repartition, are not counted.
    /// For historical reasons, the result is not precise for old SST files.
    pub(crate) fn owned_num_rows(&self, region_id: RegionId) -> u64 {
        self.levels
            .iter()
            .map(|level_meta| {
                level_meta
                    .files
                    .values()
                    .filter(|file_handle| file_handle.region_id() == region_id)
                    .map(|file_handle| {
                        let meta = file_handle.meta_ref();
                        meta.num_rows
                    })
                    .sum::<u64>()
            })
            .sum()
    }

    /// Returns the number of SST files owned by `region_id`.
    pub(crate) fn owned_num_files(&self, region_id: RegionId) -> u64 {
        self.levels
            .iter()
            .map(|level_meta| {
                level_meta
                    .files
                    .values()
                    .filter(|file_handle| file_handle.region_id() == region_id)
                    .count() as u64
            })
            .sum()
    }

    /// Returns the time range covered by every file in this version, including
    /// files referenced from other regions after a repartition.
    pub(crate) fn time_range(&self) -> Option<FileTimeRange> {
        self.levels
            .iter()
            .flat_map(|level_meta| level_meta.files.values())
            .map(|file_handle| file_handle.time_range())
            .reduce(|(min_a, max_a), (min_b, max_b)| (min_a.min(min_b), max_a.max(max_b)))
    }

    /// Returns the space occupied by SST data files owned by `region_id`.
    pub(crate) fn owned_sst_usage(&self, region_id: RegionId) -> u64 {
        self.levels
            .iter()
            .map(|level_meta| {
                level_meta
                    .files
                    .values()
                    .filter(|file_handle| file_handle.region_id() == region_id)
                    .map(|file_handle| {
                        let meta = file_handle.meta_ref();
                        meta.file_size
                    })
                    .sum::<u64>()
            })
            .sum()
    }

    /// Returns the space occupied by SST index files owned by `region_id`.
    pub(crate) fn owned_index_usage(&self, region_id: RegionId) -> u64 {
        self.levels
            .iter()
            .map(|level_meta| {
                level_meta
                    .files
                    .values()
                    .filter(|file_handle| file_handle.region_id() == region_id)
                    .map(|file_handle| {
                        let meta = file_handle.meta_ref();
                        meta.index_file_size
                    })
                    .sum::<u64>()
            })
            .sum()
    }
}

// We only has fixed number of level, so we use array to hold elements. This implementation
// detail of LevelMetaArray should not be exposed to users of [LevelMetas].
type LevelMetaArray = [LevelMeta; MAX_LEVEL as usize];

/// Metadata of files in the same SST level.
#[derive(Clone)]
pub struct LevelMeta {
    /// Level number.
    pub level: Level,
    /// Handles of SSTs in this level.
    pub files: HashMap<FileId, FileHandle>,
}

impl LevelMeta {
    /// Returns an empty meta of specific `level`.
    pub(crate) fn new(level: Level) -> LevelMeta {
        LevelMeta {
            level,
            files: HashMap::new(),
        }
    }

    /// Returns expired SSTs from current level.
    pub fn get_expired_files(&self, now: &Timestamp, ttl: &TimeToLive) -> Vec<FileHandle> {
        self.files
            .values()
            .filter(|v| {
                let (_, end) = v.time_range();

                match ttl.is_expired(&end, now) {
                    Ok(expired) => expired,
                    Err(e) => {
                        common_telemetry::error!(e; "Failed to calculate region TTL expire time");
                        false
                    }
                }
            })
            .cloned()
            .collect()
    }

    pub fn files(&self) -> impl Iterator<Item = &FileHandle> {
        self.files.values()
    }
}

impl fmt::Debug for LevelMeta {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("LevelMeta")
            .field("level", &self.level)
            .field("files", &self.files.keys())
            .finish()
    }
}

fn new_level_meta_vec() -> LevelMetaArray {
    (0u8..MAX_LEVEL)
        .map(LevelMeta::new)
        .collect::<Vec<_>>()
        .try_into()
        .unwrap() // safety: LevelMetaArray is a fixed length array with length MAX_LEVEL
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_util::new_noop_file_purger;

    #[test]
    fn time_range_spans_files_referenced_from_other_regions() {
        let purger = new_noop_file_purger();
        let owned = FileMeta {
            file_id: FileId::random(),
            region_id: RegionId::new(1, 1),
            time_range: (
                Timestamp::new_millisecond(200),
                Timestamp::new_millisecond(300),
            ),
            ..Default::default()
        };
        let referenced = FileMeta {
            file_id: FileId::random(),
            region_id: RegionId::new(2, 1),
            time_range: (
                Timestamp::new_millisecond(50),
                Timestamp::new_millisecond(100),
            ),
            ..Default::default()
        };

        let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test());
        version.add_files(purger, [owned, referenced].into_iter());

        assert_eq!(
            version.time_range(),
            Some((
                Timestamp::new_millisecond(50),
                Timestamp::new_millisecond(300)
            ))
        );
    }

    #[test]
    fn time_range_compares_across_units() {
        let purger = new_noop_file_purger();
        let seconds = FileMeta {
            file_id: FileId::random(),
            time_range: (Timestamp::new_second(1), Timestamp::new_second(2)),
            ..Default::default()
        };
        let millis = FileMeta {
            file_id: FileId::random(),
            time_range: (
                Timestamp::new_millisecond(500),
                Timestamp::new_millisecond(2500),
            ),
            ..Default::default()
        };

        let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test());
        version.add_files(purger, [seconds, millis].into_iter());

        assert_eq!(
            version.time_range(),
            Some((
                Timestamp::new_millisecond(500),
                Timestamp::new_millisecond(2500)
            ))
        );
    }

    #[test]
    fn time_range_is_none_without_files() {
        let version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test());
        assert_eq!(version.time_range(), None);
    }

    #[rstest::rstest]
    fn test_index_replacement_shares_retirement_with_old_readers(
        #[values(false, true)] retire_before_update: bool,
    ) {
        let purger = new_noop_file_purger();
        let meta = FileMeta {
            file_id: FileId::random(),
            ..Default::default()
        };
        let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test());
        version.add_files(purger.clone(), [meta.clone()].into_iter());
        let old_snapshot = version.clone();
        let old_reader = old_snapshot.levels()[0].files[&meta.file_id].clone();
        let independent = FileHandle::new(meta.clone(), purger.clone());
        if retire_before_update {
            old_reader.mark_deleted();
        }
        let mut updated = meta.clone();
        updated.index_version = 1;
        updated
            .available_indexes
            .push(crate::sst::file::IndexType::InvertedIndex);
        version.add_files(purger, [updated.clone()].into_iter());
        let new_reader = version.levels()[0].files[&meta.file_id].clone();

        assert_eq!(old_reader.meta_ref(), &meta);
        assert_eq!(new_reader.meta_ref(), &updated);
        assert_eq!(new_reader.is_deleted(), retire_before_update);
        version.remove_files([updated].into_iter());
        assert!(new_reader.is_deleted());
        assert!(old_reader.is_deleted());
        assert!(!independent.is_deleted());
        drop(version);
        assert!(old_snapshot.levels()[0].files[&meta.file_id].is_deleted());
    }

    #[test]
    fn test_add_files() {
        let purger = new_noop_file_purger();

        let files = (1..=3)
            .map(|_| FileMeta {
                file_id: FileId::random(),
                ..Default::default()
            })
            .collect::<Vec<_>>();

        let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test());
        // files[1] is added multiple times, and that's ok.
        version.add_files(purger.clone(), files[..=1].iter().cloned());
        version.add_files(purger, files[1..].iter().cloned());

        let added_files = &version.levels()[0].files;
        assert_eq!(added_files.len(), 3);
        files.iter().for_each(|f| {
            assert!(added_files.contains_key(&f.file_id));
        });
    }

    #[test]
    fn test_file_for_compaction_uses_selected_level() {
        let purger = new_noop_file_purger();
        let file_id = FileId::random();
        let selected = FileHandle::new(
            FileMeta {
                file_id,
                level: 1,
                ..Default::default()
            },
            purger.clone(),
        );
        let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test());
        version.add_files(
            purger,
            [
                FileMeta {
                    file_id,
                    level: 0,
                    ..Default::default()
                },
                selected.meta_ref().clone(),
            ]
            .into_iter(),
        );

        let current = version.file_for_compaction(&selected).unwrap();
        assert_eq!(selected.file_id(), current.file_id());
        assert_eq!(selected.level(), current.level());
    }

    #[test]
    fn test_usage_only_counts_owned_files() {
        let purger = new_noop_file_purger();
        let region_id = RegionId::new(1, 1);
        let other_region_id = RegionId::new(1, 2);

        let files = [
            FileMeta {
                region_id,
                file_id: FileId::random(),
                file_size: 100,
                index_file_size: 10,
                num_rows: 1,
                ..Default::default()
            },
            FileMeta {
                region_id,
                file_id: FileId::random(),
                file_size: 200,
                index_file_size: 20,
                num_rows: 2,
                ..Default::default()
            },
            FileMeta {
                region_id: other_region_id,
                file_id: FileId::random(),
                file_size: 300,
                index_file_size: 30,
                num_rows: 3,
                ..Default::default()
            },
        ];

        let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test());
        version.add_files(purger, files.iter().cloned());

        assert_eq!(3, version.owned_num_rows(region_id));
        assert_eq!(2, version.owned_num_files(region_id));
        assert_eq!(300, version.owned_sst_usage(region_id));
        assert_eq!(30, version.owned_index_usage(region_id));
        assert_eq!(3, version.owned_num_rows(other_region_id));
        assert_eq!(1, version.owned_num_files(other_region_id));
        assert_eq!(300, version.owned_sst_usage(other_region_id));
        assert_eq!(30, version.owned_index_usage(other_region_id));
    }
}