fix: preserve repartitioned file refs during GC (#8445)

* fix: preserve repartitioned file refs during GC

Signed-off-by: discord9 <discord9@163.com>

* refactor: clarify file refs target region

Signed-off-by: discord9 <discord9@163.com>

---------

Signed-off-by: discord9 <discord9@163.com>

src/mito2/src/engine.rs

@@ -39,6 +39,8 @@ mod drop_test;
 #[cfg(test)]
 mod edit_region_test;
 #[cfg(test)]
+mod file_ref_test;
+#[cfg(test)]
 mod filter_deleted_test;
 #[cfg(test)]
 mod flush_test;
@@ -356,10 +358,13 @@ impl MitoEngine {
                 );
             let mut dst_region_to_src_regions = Vec::with_capacity(dst2src.len());
             for (dst_region, srcs) in dst2src {
-                let Some(dst_region) = self.find_region(dst_region) else {
-                    continue;
+                let Some(region) = self.find_region(dst_region) else {
+                    return RegionNotFoundSnafu {
+                        region_id: dst_region,
+                    }
+                    .fail();
                 };
-                dst_region_to_src_regions.push((dst_region, srcs));
+                dst_region_to_src_regions.push((region, srcs));
             }
             dst_region_to_src_regions
         };

src/mito2/src/engine/file_ref_test.rs

@@ -0,0 +1,55 @@
+// 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, HashSet};
+
+use common_error::ext::ErrorExt;
+use common_error::status_code::StatusCode;
+use store_api::storage::RegionId;
+
+use crate::config::MitoConfig;
+use crate::test_util::TestEnv;
+
+#[tokio::test]
+async fn get_snapshot_of_file_refs_errors_on_missing_related_region() {
+    let mut env = TestEnv::new().await;
+    let engine = env.create_engine(MitoConfig::default()).await;
+
+    let source_region = RegionId::new(1, 1);
+    let missing_target_region = RegionId::new(1, 2);
+    let related_regions = HashMap::from([(source_region, HashSet::from([missing_target_region]))]);
+
+    let err = engine
+        .get_snapshot_of_file_refs([], related_regions)
+        .await
+        .expect_err("missing related target region should fail file ref snapshot");
+
+    assert_eq!(StatusCode::RegionNotFound, err.status_code());
+}
+
+#[tokio::test]
+async fn get_snapshot_of_file_refs_skips_missing_file_handle_region() {
+    let mut env = TestEnv::new().await;
+    let engine = env.create_engine(MitoConfig::default()).await;
+
+    let missing_region = RegionId::new(1, 1);
+    let manifest = engine
+        .get_snapshot_of_file_refs([missing_region], HashMap::new())
+        .await
+        .expect("missing file handle region should be skipped");
+
+    assert!(manifest.file_refs.is_empty());
+    assert!(manifest.manifest_version.is_empty());
+    assert!(manifest.cross_region_refs.is_empty());
+}

tests-integration/tests/repartition.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeSet;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -33,7 +34,7 @@ use servers::query_handler::sql::SqlQueryHandler;
 use session::context::{QueryContext, QueryContextRef};
 use store_api::codec::PrimaryKeyEncoding;
 use store_api::storage::RegionId;
-use tests_integration::cluster::GreptimeDbClusterBuilder;
+use tests_integration::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder};
 use tests_integration::test_util::{StorageType, get_test_store_config};
 use tokio::sync::oneshot;
 
@@ -115,6 +116,217 @@ macro_rules! repartition_tests {
     };
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_repartition_multi_hop_fulltext_index_gc_file() {
+    if StorageType::File.test_on() {
+        common_telemetry::init_default_ut_logging();
+        test_repartition_multi_hop_fulltext_index_gc().await;
+    }
+}
+
+async fn test_repartition_multi_hop_fulltext_index_gc() {
+    let store_type = StorageType::File;
+    let cluster_name = "test_repartition_multi_hop_fulltext_index_gc";
+    let table_name = "repartition_multi_hop_index_table";
+    let (store_config, _guard) = get_test_store_config(&store_type);
+    let datanodes = 3u64;
+    let home_dir = create_temp_dir("test_repartition_multi_hop_fulltext_index_gc_data_home");
+
+    let cluster = GreptimeDbClusterBuilder::new(cluster_name)
+        .await
+        .with_shared_home_dir(Arc::new(home_dir))
+        .with_datanodes(datanodes as u32)
+        .with_store_config(store_config)
+        .with_datanode_wal_config(DatanodeWalConfig::Noop)
+        .with_metasrv_gc_config(GcSchedulerOptions {
+            enable: true,
+            gc_cooldown_period: Duration::from_nanos(1),
+            ..Default::default()
+        })
+        .with_datanode_gc_config(GcConfig {
+            enable: true,
+            lingering_time: Some(Duration::from_secs(0)),
+            unknown_file_lingering_time: Duration::from_secs(0),
+            ..Default::default()
+        })
+        .build(true)
+        .await;
+    let metasrv = &cluster.metasrv;
+    let ticker = metasrv.gc_ticker().unwrap();
+
+    let query_ctx = QueryContext::arc();
+    let instance = cluster.fe_instance();
+
+    run_sql(
+        instance,
+        r#"
+        CREATE TABLE `repartition_multi_hop_index_table`(
+          `id` INT,
+          `msg` STRING FULLTEXT INDEX,
+          `ts` TIMESTAMP TIME INDEX,
+          PRIMARY KEY(`id`)
+        ) PARTITION ON COLUMNS (`id`) (
+          `id` < 100,
+          `id` >= 100
+        ) ENGINE = mito
+          WITH ('sst_format' = 'flat');
+    "#,
+        query_ctx.clone(),
+    )
+    .await
+    .unwrap();
+
+    run_sql(
+        instance,
+        r#"
+        INSERT INTO `repartition_multi_hop_index_table` VALUES
+          (10, 'quick fox in low partition', '2022-01-01 00:00:00'),
+          (60, 'quick fox in middle partition', '2022-01-01 00:00:01'),
+          (90, 'quick fox in high partition', '2022-01-01 00:00:02'),
+          (120, 'quick fox in outer partition', '2022-01-01 00:00:03');
+    "#,
+        query_ctx.clone(),
+    )
+    .await
+    .unwrap();
+    run_sql(
+        instance,
+        "ADMIN FLUSH_TABLE('repartition_multi_hop_index_table')",
+        query_ctx.clone(),
+    )
+    .await
+    .unwrap();
+
+    run_sql(
+        instance,
+        r#"
+        ALTER TABLE `repartition_multi_hop_index_table` SPLIT PARTITION (
+          `id` < 100
+        ) INTO (
+          `id` < 50,
+          `id` >= 50 AND `id` < 100
+        );
+    "#,
+        query_ctx.clone(),
+    )
+    .await
+    .unwrap();
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    run_sql(
+        instance,
+        r#"
+        ALTER TABLE `repartition_multi_hop_index_table` SPLIT PARTITION (
+          `id` >= 50 AND `id` < 100
+        ) INTO (
+          `id` >= 50 AND `id` < 75,
+          `id` >= 75 AND `id` < 100
+        );
+    "#,
+        query_ctx.clone(),
+    )
+    .await
+    .unwrap();
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    run_sql(
+        instance,
+        r#"
+        INSERT INTO `repartition_multi_hop_index_table` VALUES
+          (80, 'fresh fox in soon dropped source', '2022-01-01 00:00:04');
+    "#,
+        query_ctx.clone(),
+    )
+    .await
+    .unwrap();
+    run_sql(
+        instance,
+        "ADMIN FLUSH_TABLE('repartition_multi_hop_index_table')",
+        query_ctx.clone(),
+    )
+    .await
+    .unwrap();
+
+    let fox_query = "SELECT id, msg FROM `repartition_multi_hop_index_table` WHERE MATCHES(msg, 'fox') ORDER BY id";
+    let expected = "\
++-----+----------------------------------+
+| id  | msg                              |
++-----+----------------------------------+
+| 10  | quick fox in low partition       |
+| 60  | quick fox in middle partition    |
+| 80  | fresh fox in soon dropped source |
+| 90  | quick fox in high partition      |
+| 120 | quick fox in outer partition     |
++-----+----------------------------------+";
+    let result = run_sql(instance, fox_query, query_ctx.clone())
+        .await
+        .unwrap();
+    check_output_stream(result.data, expected).await;
+
+    let manifest_entries = cluster_manifest_index_entries(&cluster).await;
+    assert!(
+        manifest_entries
+            .iter()
+            .any(|(_, origin_region_id, region_id)| origin_region_id != region_id),
+        "expected at least one cross-origin indexed SST before GC, got {manifest_entries:?}"
+    );
+
+    run_sql(
+        instance,
+        r#"
+        ALTER TABLE `repartition_multi_hop_index_table` MERGE PARTITION (
+          `id` >= 50 AND `id` < 75,
+          `id` >= 75 AND `id` < 100
+        );
+    "#,
+        query_ctx.clone(),
+    )
+    .await
+    .unwrap();
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    trigger_full_gc(&ticker).await;
+    trigger_table_gc(metasrv, table_name).await;
+
+    let result = run_sql(instance, fox_query, query_ctx.clone())
+        .await
+        .unwrap();
+    check_output_stream(result.data, expected).await;
+
+    let manifest_index_paths = cluster_manifest_index_paths(&cluster).await;
+    assert!(
+        !manifest_index_paths.is_empty(),
+        "expected manifest index paths after GC"
+    );
+    let storage_paths = cluster.list_sst_files_from_all_datanodes().await;
+    assert!(
+        manifest_index_paths.is_subset(&storage_paths),
+        "manifest index paths should exist in storage after GC, missing: {:?}",
+        manifest_index_paths
+            .difference(&storage_paths)
+            .collect::<Vec<_>>()
+    );
+
+    let cross_origin_index_paths = cluster_manifest_index_entries(&cluster)
+        .await
+        .into_iter()
+        .filter_map(|(index_file_path, origin_region_id, region_id)| {
+            (origin_region_id != region_id).then_some(index_file_path)
+        })
+        .collect::<BTreeSet<_>>();
+    assert!(
+        !cross_origin_index_paths.is_empty(),
+        "expected cross-origin indexed SSTs after GC"
+    );
+    assert!(
+        cross_origin_index_paths.is_subset(&storage_paths),
+        "cross-origin manifest index paths should exist in storage after GC, missing: {:?}",
+        cross_origin_index_paths
+            .difference(&storage_paths)
+            .collect::<Vec<_>>()
+    );
+}
+
 pub async fn test_partition_unpartitioned_mito(store_type: StorageType) {
     info!(
         "test_partition_unpartitioned_mito: store_type: {:?}",
@@ -583,6 +795,35 @@ async fn trigger_full_gc(ticker: &GcTickerRef) {
     let _ = rx.await.unwrap().unwrap();
 }
 
+async fn cluster_manifest_index_entries(
+    cluster: &GreptimeDbCluster,
+) -> Vec<(String, RegionId, RegionId)> {
+    let mut entries = Vec::new();
+    for datanode in cluster.datanode_instances.values() {
+        let region_server = datanode.region_server();
+        let mito = region_server.mito_engine().unwrap();
+        entries.extend(
+            mito.all_ssts_from_manifest()
+                .await
+                .into_iter()
+                .filter_map(|entry| {
+                    entry.index_file_path.map(|index_file_path| {
+                        (index_file_path, entry.origin_region_id, entry.region_id)
+                    })
+                }),
+        );
+    }
+    entries
+}
+
+async fn cluster_manifest_index_paths(cluster: &GreptimeDbCluster) -> BTreeSet<String> {
+    cluster_manifest_index_entries(cluster)
+        .await
+        .into_iter()
+        .map(|(index_file_path, _, _)| index_file_path)
+        .collect()
+}
+
 fn query_partitions_sql(table_name: &str) -> String {
     // We query information_schema.partitions to assert repartition results across engines,
     // rather than relying on SHOW CREATE TABLE formatting differences.