From be1109098be4d1bca1df5fa7c94d78f6faa76aef Mon Sep 17 00:00:00 2001
From: Gavin Chou <gavineaglechou@gmail.com>
Date: Wed, 14 Aug 2024 18:19:15 +0800
Subject: [PATCH] [fix](file cache) Fix data race of rowset meta when download
 segment data

The original impl. _rs_metas is not protected correctly when doing traversal.
We have to access the rowset meta via tablet level interface to ensure integrity.
---
 be/src/cloud/cloud_warm_up_manager.cpp          | 12 +++++++++++-
 be/src/io/cache/block_file_cache_downloader.cpp | 12 +++++++++++-
 be/src/olap/tablet_meta.h                       | 12 ------------
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/be/src/cloud/cloud_warm_up_manager.cpp b/be/src/cloud/cloud_warm_up_manager.cpp
index 47046de36985ad..07beeaeb078a46 100644
--- a/be/src/cloud/cloud_warm_up_manager.cpp
+++ b/be/src/cloud/cloud_warm_up_manager.cpp
@@ -49,6 +49,16 @@ CloudWarmUpManager::~CloudWarmUpManager() {
     }
 }
 
+std::unordered_map<std::string, RowsetMetaSharedPtr> snapshot_rs_metas(BaseTablet* tablet) {
+    std::unordered_map<std::string, RowsetMetaSharedPtr> id_to_rowset_meta_map;
+    auto visitor = [&id_to_rowset_meta_map](const RowsetSharedPtr& r) {
+        id_to_rowset_meta_map.emplace(r->rowset_meta()->rowset_id().to_string(), r->rowset_meta());
+    };
+    constexpr bool include_stale = false;
+    tablet->traverse_rowsets(visitor, include_stale);
+    return id_to_rowset_meta_map;
+}
+
 void CloudWarmUpManager::handle_jobs() {
 #ifndef BE_TEST
     constexpr int WAIT_TIME_SECONDS = 600;
@@ -78,7 +88,7 @@ void CloudWarmUpManager::handle_jobs() {
             std::shared_ptr<bthread::CountdownEvent> wait =
                     std::make_shared<bthread::CountdownEvent>(0);
             auto tablet_meta = tablet->tablet_meta();
-            auto rs_metas = tablet_meta->snapshot_rs_metas();
+            auto rs_metas = snapshot_rs_metas(tablet.get());
             for (auto& [_, rs] : rs_metas) {
                 for (int64_t seg_id = 0; seg_id < rs->num_segments(); seg_id++) {
                     auto storage_resource = rs->remote_storage_resource();
diff --git a/be/src/io/cache/block_file_cache_downloader.cpp b/be/src/io/cache/block_file_cache_downloader.cpp
index 585c0ff015993b..026f7e2a01741d 100644
--- a/be/src/io/cache/block_file_cache_downloader.cpp
+++ b/be/src/io/cache/block_file_cache_downloader.cpp
@@ -130,6 +130,16 @@ void FileCacheBlockDownloader::check_download_task(const std::vector<int64_t>& t
     }
 }
 
+std::unordered_map<std::string, RowsetMetaSharedPtr> snapshot_rs_metas(BaseTablet* tablet) {
+    std::unordered_map<std::string, RowsetMetaSharedPtr> id_to_rowset_meta_map;
+    auto visitor = [&id_to_rowset_meta_map](const RowsetSharedPtr& r) {
+        id_to_rowset_meta_map.emplace(r->rowset_meta()->rowset_id().to_string(), r->rowset_meta());
+    };
+    constexpr bool include_stale = false;
+    tablet->traverse_rowsets(visitor, include_stale);
+    return id_to_rowset_meta_map;
+}
+
 void FileCacheBlockDownloader::download_file_cache_block(
         const DownloadTask::FileCacheBlockMetaVec& metas) {
     std::ranges::for_each(metas, [&](const FileCacheBlockMeta& meta) {
@@ -141,7 +151,7 @@ void FileCacheBlockDownloader::download_file_cache_block(
             tablet = std::move(res).value();
         }
 
-        auto id_to_rowset_meta_map = tablet->tablet_meta()->snapshot_rs_metas();
+        auto id_to_rowset_meta_map = snapshot_rs_metas(tablet.get());
         auto find_it = id_to_rowset_meta_map.find(meta.rowset_id());
         if (find_it == id_to_rowset_meta_map.end()) {
             return;
diff --git a/be/src/olap/tablet_meta.h b/be/src/olap/tablet_meta.h
index bb6b5b8cd51725..41455c051c7f44 100644
--- a/be/src/olap/tablet_meta.h
+++ b/be/src/olap/tablet_meta.h
@@ -192,9 +192,6 @@ class TabletMeta {
     void revise_delete_bitmap_unlocked(const DeleteBitmap& delete_bitmap);
 
     const std::vector<RowsetMetaSharedPtr>& all_stale_rs_metas() const;
-    // return the snapshot of rowset_meta
-    // the return value is map<rowset_id_str, rowset_meta_sptr>
-    std::unordered_map<std::string, RowsetMetaSharedPtr> snapshot_rs_metas() const;
     RowsetMetaSharedPtr acquire_rs_meta_by_version(const Version& version) const;
     void delete_stale_rs_meta_by_version(const Version& version);
     RowsetMetaSharedPtr acquire_stale_rs_meta_by_version(const Version& version) const;
@@ -698,15 +695,6 @@ inline bool TabletMeta::all_beta() const {
     return true;
 }
 
-inline std::unordered_map<std::string, RowsetMetaSharedPtr> TabletMeta::snapshot_rs_metas() const {
-    std::unordered_map<std::string, RowsetMetaSharedPtr> id_to_rowset_meta_map;
-    std::shared_lock rlock(_meta_lock);
-    std::for_each(_rs_metas.cbegin(), _rs_metas.cend(), [&](const auto& rowset_meta) {
-        id_to_rowset_meta_map.emplace(rowset_meta->rowset_id().to_string(), rowset_meta);
-    });
-    return id_to_rowset_meta_map;
-}
-
 std::string tablet_state_name(TabletState state);
 
 // Only for unit test now.