feat(mito): list SSTs from manifest and storage (#6766)

* feat(engine): list SSTs from manifest and storage

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* mito engine only

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* try best to remove allocation

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* address comments

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

---------

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

src/metric-engine/src/engine.rs

@@ -487,7 +487,7 @@ mod test {
     use mito2::sst::location::region_dir_from_table_dir;
     use store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY;
     use store_api::region_request::{
-        PathType, RegionCloseRequest, RegionFlushRequest, RegionOpenRequest,
+        PathType, RegionCloseRequest, RegionFlushRequest, RegionOpenRequest, RegionRequest,
     };
 
     use super::*;

src/metric-engine/src/engine/flush.rs

@@ -50,3 +50,117 @@ impl MetricEngineInner {
             .map(|response| response.affected_rows)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use api::v1::Rows;
+    use futures_util::TryStreamExt;
+    use itertools::Itertools as _;
+    use store_api::region_request::RegionPutRequest;
+
+    use super::*;
+    use crate::test_util::{build_rows, row_schema_with_tags, TestEnv};
+
+    #[tokio::test]
+    async fn test_list_ssts_after_write_and_flush_metric() {
+        let env = TestEnv::new().await;
+        let engine = env.metric();
+
+        let phy_to_logi = vec![
+            (
+                RegionId::new(11, 1),
+                vec![RegionId::new(1111, 11), RegionId::new(1111, 12)],
+            ),
+            (RegionId::new(11, 2), vec![RegionId::new(1111, 2)]),
+            (RegionId::new(22, 42), vec![RegionId::new(2222, 42)]),
+        ];
+
+        for (phy_region_id, logi_region_ids) in &phy_to_logi {
+            env.create_physical_region(*phy_region_id, &TestEnv::default_table_dir())
+                .await;
+            for logi_region_id in logi_region_ids {
+                env.create_logical_region(*phy_region_id, *logi_region_id)
+                    .await;
+            }
+        }
+
+        // write to logical regions
+        for (_, logi_region_ids) in &phy_to_logi {
+            for logi_region_id in logi_region_ids {
+                let schema = row_schema_with_tags(&["job"]);
+                let rows = build_rows(1, 10);
+                let request = RegionRequest::Put(RegionPutRequest {
+                    rows: Rows { schema, rows },
+                    hint: None,
+                });
+                engine
+                    .handle_request(*logi_region_id, request)
+                    .await
+                    .unwrap();
+            }
+        }
+
+        for (phy_region_id, _) in &phy_to_logi {
+            engine
+                .handle_request(*phy_region_id, RegionRequest::Flush(Default::default()))
+                .await
+                .unwrap();
+        }
+
+        // list from manifest
+        let mito = env.mito();
+        let debug_format = mito
+            .all_ssts_from_manifest()
+            .map(|mut e| {
+                e.file_path = e.file_path.replace(&e.file_id, "<file_id>");
+                e.index_file_path = e
+                    .index_file_path
+                    .map(|path| path.replace(&e.file_id, "<file_id>"));
+                e.file_id = "<file_id>".to_string();
+                format!("\n{:?}", e)
+            })
+            .sorted()
+            .collect::<Vec<_>>()
+            .join("");
+        assert_eq!(
+            debug_format,
+            r#"
+ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", level: 0, file_path: "test_metric_region/11_0000000001/data/<file_id>.parquet", file_size: 3157, index_file_path: Some("test_metric_region/11_0000000001/data/index/<file_id>.puffin"), index_file_size: Some(235), num_rows: 10, num_row_groups: 1, min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20) }
+ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", level: 0, file_path: "test_metric_region/11_0000000002/data/<file_id>.parquet", file_size: 3157, index_file_path: Some("test_metric_region/11_0000000002/data/index/<file_id>.puffin"), index_file_size: Some(235), num_rows: 10, num_row_groups: 1, min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10) }
+ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "<file_id>", level: 0, file_path: "test_metric_region/11_0000000001/metadata/<file_id>.parquet", file_size: 3201, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8) }
+ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "<file_id>", level: 0, file_path: "test_metric_region/11_0000000002/metadata/<file_id>.parquet", file_size: 3185, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4) }
+ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", level: 0, file_path: "test_metric_region/22_0000000042/data/<file_id>.parquet", file_size: 3157, index_file_path: Some("test_metric_region/22_0000000042/data/index/<file_id>.puffin"), index_file_size: Some(235), num_rows: 10, num_row_groups: 1, min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10) }
+ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "<file_id>", level: 0, file_path: "test_metric_region/22_0000000042/metadata/<file_id>.parquet", file_size: 3185, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4) }"#
+        );
+
+        // list from storage
+        let storage_entries = mito
+            .all_ssts_from_storage()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let debug_format = storage_entries
+            .into_iter()
+            .map(|mut e| {
+                let i = e.file_path.rfind('/').unwrap();
+                e.file_path.replace_range(i..(i + 37), "/<file_id>");
+                format!("\n{:?}", e)
+            })
+            .sorted()
+            .collect::<Vec<_>>()
+            .join("");
+        assert_eq!(
+            debug_format,
+            r#"
+StorageSstEntry { file_path: "test_metric_region/11_0000000001/data/<file_id>.parquet", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test_metric_region/11_0000000001/data/index/<file_id>.puffin", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test_metric_region/11_0000000001/metadata/<file_id>.parquet", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test_metric_region/11_0000000002/data/<file_id>.parquet", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test_metric_region/11_0000000002/data/index/<file_id>.puffin", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test_metric_region/11_0000000002/metadata/<file_id>.parquet", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test_metric_region/22_0000000042/data/<file_id>.parquet", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test_metric_region/22_0000000042/data/index/<file_id>.puffin", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test_metric_region/22_0000000042/metadata/<file_id>.parquet", file_size: None, last_modified_ms: None }"#
+        );
+    }
+}

src/mito2/src/access_layer.rs

@@ -15,6 +15,9 @@
 use std::sync::Arc;
 use std::time::Duration;
 
+use async_stream::try_stream;
+use common_time::Timestamp;
+use futures::{Stream, TryStreamExt};
 use object_store::services::Fs;
 use object_store::util::{join_dir, with_instrument_layers};
 use object_store::{ErrorKind, ObjectStore, ATOMIC_WRITE_DIR, OLD_ATOMIC_WRITE_DIR};
@@ -22,6 +25,7 @@ use smallvec::SmallVec;
 use snafu::ResultExt;
 use store_api::metadata::RegionMetadataRef;
 use store_api::region_request::PathType;
+use store_api::sst_entry::StorageSstEntry;
 use store_api::storage::{RegionId, SequenceNumber};
 
 use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey};
@@ -300,6 +304,46 @@ impl AccessLayer {
 
         Ok((sst_info, metrics))
     }
+
+    /// Lists the SST entries from the storage layer in the table directory.
+    pub fn storage_sst_entries(&self) -> impl Stream<Item = Result<StorageSstEntry>> {
+        let object_store = self.object_store.clone();
+        let table_dir = self.table_dir.clone();
+
+        try_stream! {
+            let mut lister = object_store
+                .lister_with(table_dir.as_str())
+                .recursive(true)
+                .await
+                .context(OpenDalSnafu)?;
+
+            while let Some(entry) = lister.try_next().await.context(OpenDalSnafu)? {
+                let metadata = entry.metadata();
+                if metadata.is_dir() {
+                    continue;
+                }
+
+                let path = entry.path();
+                if !path.ends_with(".parquet") && !path.ends_with(".puffin") {
+                    continue;
+                }
+
+                let file_size = metadata.content_length();
+                let file_size = if file_size == 0 { None } else { Some(file_size) };
+                let last_modified_ms = metadata
+                    .last_modified()
+                    .map(|ts| Timestamp::new_millisecond(ts.timestamp_millis()));
+
+                let entry = StorageSstEntry {
+                    file_path: path.to_string(),
+                    file_size,
+                    last_modified_ms,
+                };
+
+                yield entry;
+            }
+        }
+    }
 }
 
 /// `OperationType` represents the origin of the `SstWriteRequest`.

src/mito2/src/engine.rs

@@ -77,6 +77,7 @@ use common_recordbatch::SendableRecordBatchStream;
 use common_telemetry::{info, tracing};
 use common_wal::options::{WalOptions, WAL_OPTIONS_KEY};
 use futures::future::{join_all, try_join_all};
+use futures::stream::{self, Stream, StreamExt};
 use object_store::manager::ObjectStoreManagerRef;
 use snafu::{ensure, OptionExt, ResultExt};
 use store_api::codec::PrimaryKeyEncoding;
@@ -91,6 +92,7 @@ use store_api::region_engine::{
     RegionStatistic, SetRegionRoleStateResponse, SettableRegionRoleState, SyncManifestResponse,
 };
 use store_api::region_request::{AffectedRows, RegionOpenRequest, RegionRequest};
+use store_api::sst_entry::{ManifestSstEntry, StorageSstEntry};
 use store_api::storage::{RegionId, ScanRequest, SequenceNumber};
 use store_api::ManifestVersion;
 use tokio::sync::{oneshot, Semaphore};
@@ -382,6 +384,32 @@ impl MitoEngine {
             .collect::<Vec<_>>();
         Ok((memtable_stats, sst_stats))
     }
+
+    /// Lists all SSTs from the manifest of all regions in the engine.
+    pub fn all_ssts_from_manifest(&self) -> impl Iterator<Item = ManifestSstEntry> + use<'_> {
+        self.inner
+            .workers
+            .all_regions()
+            .flat_map(|region| region.manifest_sst_entries())
+    }
+
+    /// Lists all SSTs from the storage layer of all regions in the engine.
+    pub fn all_ssts_from_storage(&self) -> impl Stream<Item = Result<StorageSstEntry>> {
+        let regions = self.inner.workers.all_regions();
+
+        let mut layers_distinct_table_dirs = HashMap::new();
+        for region in regions {
+            let table_dir = region.access_layer.table_dir();
+            if !layers_distinct_table_dirs.contains_key(table_dir) {
+                layers_distinct_table_dirs
+                    .insert(table_dir.to_string(), region.access_layer.clone());
+            }
+        }
+
+        stream::iter(layers_distinct_table_dirs)
+            .map(|(_, access_layer)| access_layer.storage_sst_entries())
+            .flatten()
+    }
 }
 
 /// Check whether the region edit is valid. Only adding files to region is considered valid now.

src/mito2/src/engine/basic_test.rs

@@ -27,11 +27,13 @@ use common_recordbatch::RecordBatches;
 use common_wal::options::WAL_OPTIONS_KEY;
 use datatypes::prelude::ConcreteDataType;
 use datatypes::schema::ColumnSchema;
+use futures::TryStreamExt;
+use itertools::Itertools;
 use rstest::rstest;
 use rstest_reuse::{self, apply};
 use store_api::metadata::ColumnMetadata;
 use store_api::region_request::{
-    PathType, RegionCreateRequest, RegionOpenRequest, RegionPutRequest,
+    PathType, RegionCreateRequest, RegionFlushRequest, RegionOpenRequest, RegionPutRequest,
 };
 use store_api::storage::RegionId;
 
@@ -720,3 +722,97 @@ async fn test_cache_null_primary_key() {
 +-------+-------+---------+---------------------+";
     assert_eq!(expected, batches.pretty_print().unwrap());
 }
+
+#[tokio::test]
+async fn test_list_ssts() {
+    let mut env = TestEnv::new().await;
+    let engine = env.create_engine(MitoConfig::default()).await;
+
+    // Create 3 regions and write some rows to each of them
+    let region_ids = vec![
+        RegionId::new(11, 1),
+        RegionId::new(11, 2),
+        RegionId::new(22, 42),
+    ];
+
+    for region_id in &region_ids {
+        let mut request = CreateRequestBuilder::new().build();
+        request.column_metadatas[0]
+            .column_schema
+            .set_inverted_index(true); // set inverted index for tag_0
+        let column_schemas = rows_schema(&request);
+        engine
+            .handle_request(*region_id, RegionRequest::Create(request))
+            .await
+            .unwrap();
+
+        // write some rows
+        let rows = Rows {
+            schema: column_schemas.clone(),
+            rows: build_rows_for_key("x", 0, 10, 0),
+        };
+        put_rows(&engine, *region_id, rows).await;
+    }
+
+    // flush to generate sst
+    for region_id in &region_ids {
+        engine
+            .handle_request(
+                *region_id,
+                RegionRequest::Flush(RegionFlushRequest {
+                    row_group_size: None,
+                }),
+            )
+            .await
+            .unwrap();
+    }
+
+    // list from manifest
+    let debug_format = engine
+        .all_ssts_from_manifest()
+        .map(|mut e| {
+            e.file_path = e.file_path.replace(&e.file_id, "<file_id>");
+            e.index_file_path = e
+                .index_file_path
+                .map(|p| p.replace(&e.file_id, "<file_id>"));
+            e.file_id = "<file_id>".to_string();
+            format!("\n{:?}", e)
+        })
+        .sorted()
+        .collect::<Vec<_>>()
+        .join("");
+    assert_eq!(
+        debug_format,
+        r#"
+ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", level: 0, file_path: "test/11_0000000001/<file_id>.parquet", file_size: 2515, index_file_path: Some("test/11_0000000001/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10) }
+ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", level: 0, file_path: "test/11_0000000002/<file_id>.parquet", file_size: 2515, index_file_path: Some("test/11_0000000002/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10) }
+ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", level: 0, file_path: "test/22_0000000042/<file_id>.parquet", file_size: 2515, index_file_path: Some("test/22_0000000042/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10) }"#
+    );
+
+    // list from storage
+    let storage_entries = engine
+        .all_ssts_from_storage()
+        .try_collect::<Vec<_>>()
+        .await
+        .unwrap();
+    let debug_format = storage_entries
+        .into_iter()
+        .map(|mut e| {
+            let i = e.file_path.rfind('/').unwrap();
+            e.file_path.replace_range(i..(i + 37), "/<file_id>");
+            format!("\n{:?}", e)
+        })
+        .sorted()
+        .collect::<Vec<_>>()
+        .join("");
+    assert_eq!(
+        debug_format,
+        r#"
+StorageSstEntry { file_path: "test/11_0000000001/<file_id>.parquet", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test/11_0000000001/index/<file_id>.puffin", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test/11_0000000002/<file_id>.parquet", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test/11_0000000002/index/<file_id>.puffin", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test/22_0000000042/<file_id>.parquet", file_size: None, last_modified_ms: None }
+StorageSstEntry { file_path: "test/22_0000000042/index/<file_id>.puffin", file_size: None, last_modified_ms: None }"#
+    );
+}

src/mito2/src/region.rs

@@ -32,6 +32,7 @@ use store_api::metadata::RegionMetadataRef;
 use store_api::region_engine::{
     RegionManifestInfo, RegionRole, RegionStatistic, SettableRegionRoleState,
 };
+use store_api::sst_entry::ManifestSstEntry;
 use store_api::storage::{RegionId, SequenceNumber};
 use store_api::ManifestVersion;
 
@@ -46,6 +47,7 @@ use crate::memtable::MemtableBuilderRef;
 use crate::region::version::{VersionControlRef, VersionRef};
 use crate::request::{OnFailure, OptionOutputTx};
 use crate::sst::file_purger::FilePurgerRef;
+use crate::sst::location::{index_file_path, sst_file_path};
 use crate::time_provider::TimeProviderRef;
 
 /// This is the approximate factor to estimate the size of wal.
@@ -497,6 +499,48 @@ impl MitoRegion {
             })?;
         Ok(())
     }
+
+    /// Returns the SST entries of the region.
+    pub fn manifest_sst_entries(&self) -> Vec<ManifestSstEntry> {
+        let table_dir = self.table_dir();
+        let path_type = self.access_layer.path_type();
+        self.version()
+            .ssts
+            .levels()
+            .iter()
+            .flat_map(|level| {
+                level.files().map(|file| {
+                    let meta = file.meta_ref();
+                    let region_id = meta.region_id;
+                    let (index_file_path, index_file_size) = if meta.index_file_size > 0 {
+                        let index_file_path = index_file_path(table_dir, meta.file_id(), path_type);
+                        (Some(index_file_path), Some(meta.index_file_size))
+                    } else {
+                        (None, None)
+                    };
+                    ManifestSstEntry {
+                        table_dir: table_dir.to_string(),
+                        region_id,
+                        table_id: region_id.table_id(),
+                        region_number: region_id.region_number(),
+                        region_group: region_id.region_group(),
+                        region_sequence: region_id.region_sequence(),
+                        file_id: meta.file_id.to_string(),
+                        level: meta.level,
+                        file_path: sst_file_path(table_dir, meta.file_id(), path_type),
+                        file_size: meta.file_size,
+                        index_file_path,
+                        index_file_size,
+                        num_rows: meta.num_rows,
+                        num_row_groups: meta.num_row_groups,
+                        min_ts: meta.time_range.0,
+                        max_ts: meta.time_range.1,
+                        sequence: meta.sequence.map(|s| s.get()),
+                    }
+                })
+            })
+            .collect()
+    }
 }
 
 /// Context to update the region manifest.

src/mito2/src/worker.rs

@@ -272,6 +272,12 @@ impl WorkerGroup {
 
         &self.workers[index]
     }
+
+    pub(crate) fn all_regions(&self) -> impl Iterator<Item = MitoRegionRef> + use<'_> {
+        self.workers
+            .iter()
+            .flat_map(|worker| worker.regions.list_regions())
+    }
 }
 
 // Tests methods.

src/store-api/src/lib.rs

@@ -27,6 +27,7 @@ pub mod mito_engine_options;
 pub mod path_utils;
 pub mod region_engine;
 pub mod region_request;
+pub mod sst_entry;
 pub mod storage;
 
 pub type ManifestVersion = u64;

src/store-api/src/sst_entry.rs

@@ -0,0 +1,68 @@
+// 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 common_time::Timestamp;
+use serde::{Deserialize, Serialize};
+
+use crate::storage::{RegionGroup, RegionId, RegionNumber, RegionSeq, TableId};
+
+/// An entry describing a SST file known by the engine's manifest.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct ManifestSstEntry {
+    /// The table directory this file belongs to.
+    pub table_dir: String,
+    /// The region id this file belongs to.
+    pub region_id: RegionId,
+    /// The table id this file belongs to.
+    pub table_id: TableId,
+    /// The region number this file belongs to.
+    pub region_number: RegionNumber,
+    /// The region group this file belongs to.
+    pub region_group: RegionGroup,
+    /// The region sequence this file belongs to.
+    pub region_sequence: RegionSeq,
+    /// Engine-specific file identifier (string form).
+    pub file_id: String,
+    /// SST level.
+    pub level: u8,
+    /// Full path of the SST file in object store.
+    pub file_path: String,
+    /// File size in bytes.
+    pub file_size: u64,
+    /// Full path of the index file in object store.
+    pub index_file_path: Option<String>,
+    /// File size of the index file in object store.
+    pub index_file_size: Option<u64>,
+    /// Number of rows in the SST.
+    pub num_rows: u64,
+    /// Number of row groups in the SST.
+    pub num_row_groups: u64,
+    /// Min timestamp.
+    pub min_ts: Timestamp,
+    /// Max timestamp.
+    pub max_ts: Timestamp,
+    /// The sequence number associated with this file.
+    pub sequence: Option<u64>,
+}
+
+/// An entry describing a SST file listed from storage layer directly.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct StorageSstEntry {
+    /// Full path of the SST file in object store.
+    pub file_path: String,
+    /// File size in bytes.
+    pub file_size: Option<u64>,
+    /// Last modified time in milliseconds since epoch, if available from storage.
+    pub last_modified_ms: Option<Timestamp>,
+}