diff --git a/be/src/cloud/cloud_warm_up_manager.cpp b/be/src/cloud/cloud_warm_up_manager.cpp index 02df3e9113739d..a397498a56d6cf 100644 --- a/be/src/cloud/cloud_warm_up_manager.cpp +++ b/be/src/cloud/cloud_warm_up_manager.cpp @@ -19,6 +19,8 @@ #include #include +#include +#include #include #include @@ -803,10 +805,39 @@ void CloudWarmUpManager::record_balanced_tablet(int64_t tablet_id, const std::st meta.brpc_port = brpc_port; shard.tablets.emplace(tablet_id, std::move(meta)); g_balance_tablet_be_mapping_size << 1; + schedule_remove_balanced_tablet(tablet_id); VLOG_DEBUG << "Recorded balanced warm up cache tablet: tablet_id=" << tablet_id << ", host=" << host << ":" << brpc_port; } +void CloudWarmUpManager::schedule_remove_balanced_tablet(int64_t tablet_id) { + // Use std::make_unique to avoid raw pointer allocation + auto tablet_id_ptr = std::make_unique(tablet_id); + unsigned long expired_ms = g_tablet_report_inactive_duration_ms; + if (doris::config::cache_read_from_peer_expired_seconds > 0 && + doris::config::cache_read_from_peer_expired_seconds <= + g_tablet_report_inactive_duration_ms / 1000) { + expired_ms = doris::config::cache_read_from_peer_expired_seconds * 1000; + } + bthread_timer_t timer_id; + // ATTN: The timer callback will reclaim ownership of the tablet_id_ptr, so we need to release it after the timer is added. + if (const int rc = bthread_timer_add(&timer_id, butil::milliseconds_from_now(expired_ms), + clean_up_expired_mappings, tablet_id_ptr.get()); + rc == 0) { + tablet_id_ptr.release(); + } else { + LOG(WARNING) << "Fail to add timer for clean up expired mappings for tablet_id=" + << tablet_id << " rc=" << rc; + } +} + +void CloudWarmUpManager::clean_up_expired_mappings(void* arg) { + std::unique_ptr tid(static_cast(arg)); + auto& manager = ExecEnv::GetInstance()->storage_engine().to_cloud().cloud_warm_up_manager(); + manager.remove_balanced_tablet(*tid); + VLOG_DEBUG << "Removed expired balanced warm up cache tablet: tablet_id=" << *tid; +} + std::optional> CloudWarmUpManager::get_balanced_tablet_info( int64_t tablet_id) { auto& shard = get_shard(tablet_id); diff --git a/be/src/cloud/cloud_warm_up_manager.h b/be/src/cloud/cloud_warm_up_manager.h index 8725dc069398e0..f5b00394a2768f 100644 --- a/be/src/cloud/cloud_warm_up_manager.h +++ b/be/src/cloud/cloud_warm_up_manager.h @@ -98,6 +98,8 @@ class CloudWarmUpManager { std::unordered_map> get_all_balanced_tablets() const; private: + void schedule_remove_balanced_tablet(int64_t tablet_id); + static void clean_up_expired_mappings(void* arg); void handle_jobs(); Status _do_warm_up_rowset(RowsetMeta& rs_meta, std::vector& replicas, diff --git a/be/src/io/cache/block_file_cache_downloader.cpp b/be/src/io/cache/block_file_cache_downloader.cpp index a74ddcad6279de..9c19083dbe2c06 100644 --- a/be/src/io/cache/block_file_cache_downloader.cpp +++ b/be/src/io/cache/block_file_cache_downloader.cpp @@ -173,14 +173,6 @@ std::unordered_map snapshot_rs_metas(BaseTable return id_to_rowset_meta_map; } -static void clean_up_expired_mappings(void* arg) { - // Reclaim ownership with unique_ptr for automatic memory management - std::unique_ptr tablet_id(static_cast(arg)); - auto& manager = ExecEnv::GetInstance()->storage_engine().to_cloud().cloud_warm_up_manager(); - manager.remove_balanced_tablet(*tablet_id); - VLOG_DEBUG << "Removed expired balanced warm up cache tablet: tablet_id=" << *tablet_id; -} - void FileCacheBlockDownloader::download_file_cache_block( const DownloadTask::FileCacheBlockMetaVec& metas) { std::unordered_set synced_tablets; @@ -189,9 +181,34 @@ void FileCacheBlockDownloader::download_file_cache_block( << ", rowset_id=" << meta.rowset_id() << ", segment_id=" << meta.segment_id() << ", offset=" << meta.offset() << ", size=" << meta.size() << ", type=" << meta.cache_type(); + + // Helper to decrease inflight count on early return. + // NOTE: This lambda captures 'this' pointer. It's safe because: + // 1. download_segment_file() calls download_done synchronously + // 2. ~FileCacheBlockDownloader() waits for all workers to finish via _workers->shutdown() + // If this assumption changes (e.g., async callback), consider using shared_from_this pattern. + auto decrease_inflight_count = [this, tablet_id = meta.tablet_id()]() { + std::lock_guard lock(_inflight_mtx); + auto it = _inflight_tablets.find(tablet_id); + if (it == _inflight_tablets.end()) { + LOG(WARNING) << "inflight ref cnt not exist, tablet id " << tablet_id; + } else { + it->second--; + VLOG_DEBUG << "download_file_cache_block: inflight_tablets[" << tablet_id + << "] = " << it->second; + if (it->second <= 0) { + DCHECK_EQ(it->second, 0) << it->first; + _inflight_tablets.erase(it); + VLOG_DEBUG << "download_file_cache_block: erase inflight_tablets[" << tablet_id + << "]"; + } + } + }; + CloudTabletSPtr tablet; if (auto res = _engine.tablet_mgr().get_tablet(meta.tablet_id(), false); !res.has_value()) { LOG(INFO) << "failed to find tablet " << meta.tablet_id() << " : " << res.error(); + decrease_inflight_count(); return; } else { tablet = std::move(res).value(); @@ -210,12 +227,14 @@ void FileCacheBlockDownloader::download_file_cache_block( if (find_it == id_to_rowset_meta_map.end()) { LOG(WARNING) << "download_file_cache_block: tablet_id=" << meta.tablet_id() << " rowset_id not found, rowset_id=" << meta.rowset_id(); + decrease_inflight_count(); return; } auto storage_resource = find_it->second->remote_storage_resource(); if (!storage_resource) { LOG(WARNING) << storage_resource.error(); + decrease_inflight_count(); return; } // Use RowsetMeta::fs() instead of storage_resource->fs to support packed file. @@ -224,49 +243,19 @@ void FileCacheBlockDownloader::download_file_cache_block( // to their actual locations within packed files. auto file_system = find_it->second->fs(); if (!file_system) { + decrease_inflight_count(); LOG(WARNING) << "download_file_cache_block: failed to get file system for tablet_id=" << meta.tablet_id() << ", rowset_id=" << meta.rowset_id(); return; } - auto download_done = [&, tablet_id = meta.tablet_id()](Status st) { - std::lock_guard lock(_inflight_mtx); - auto it = _inflight_tablets.find(tablet_id); + // Capture decrease_inflight_count by value to ensure lifetime safety + // even if download_done is called asynchronously in the future + auto download_done = [decrease_inflight_count, tablet_id = meta.tablet_id()](Status st) { TEST_SYNC_POINT_CALLBACK("FileCacheBlockDownloader::download_file_cache_block"); - if (it == _inflight_tablets.end()) { - LOG(WARNING) << "inflight ref cnt not exist, tablet id " << tablet_id; - } else { - it->second--; - VLOG_DEBUG << "download_file_cache_block: inflight_tablets[" << tablet_id - << "] = " << it->second; - if (it->second <= 0) { - DCHECK_EQ(it->second, 0) << it->first; - _inflight_tablets.erase(it); - VLOG_DEBUG << "download_file_cache_block: erase inflight_tablets[" << tablet_id - << "]"; - } - } - // Use std::make_unique to avoid raw pointer allocation - auto tablet_id_ptr = std::make_unique(tablet_id); - unsigned long expired_ms = g_tablet_report_inactive_duration_ms; - if (doris::config::cache_read_from_peer_expired_seconds > 0 && - doris::config::cache_read_from_peer_expired_seconds <= - g_tablet_report_inactive_duration_ms / 1000) { - expired_ms = doris::config::cache_read_from_peer_expired_seconds * 1000; - } - bthread_timer_t timer_id; - // ATTN: The timer callback will reclaim ownership of the tablet_id_ptr, so we need to release it after the timer is added. - if (const int rc = - bthread_timer_add(&timer_id, butil::milliseconds_from_now(expired_ms), - clean_up_expired_mappings, tablet_id_ptr.get()); - rc == 0) { - tablet_id_ptr.release(); - } else { - LOG(WARNING) << "Fail to add timer for clean up expired mappings for tablet_id=" - << tablet_id << " rc=" << rc; - } + decrease_inflight_count(); LOG(INFO) << "download_file_cache_block: download_done, tablet_Id=" << tablet_id - << " status=" << st.to_string() << " expired_ms=" << expired_ms; + << " status=" << st.to_string(); }; std::string path; diff --git a/be/test/io/cache/block_file_cache_downloader_test.cpp b/be/test/io/cache/block_file_cache_downloader_test.cpp new file mode 100644 index 00000000000000..79974ef4860ad1 --- /dev/null +++ b/be/test/io/cache/block_file_cache_downloader_test.cpp @@ -0,0 +1,240 @@ +// +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +#include "io/cache/block_file_cache_downloader.h" + +#include +#include + +#include +#include + +#include "cloud/cloud_storage_engine.h" +#include "common/config.h" + +namespace doris::io { + +class FileCacheBlockDownloaderTest : public testing::Test { +public: + FileCacheBlockDownloaderTest() : _engine(CloudStorageEngine(EngineOptions {})) {} + + void SetUp() override { + // Enable file cache for testing + config::enable_file_cache = true; + // Set reasonable thread pool size for testing + config::file_cache_downloader_thread_num_min = 2; + config::file_cache_downloader_thread_num_max = 4; + } + + void TearDown() override { config::enable_file_cache = false; } + + // Helper to wait for inflight tasks to complete with timeout + bool wait_for_task_done(FileCacheBlockDownloader& downloader, int64_t tablet_id, + int timeout_seconds = 10) { + auto start = std::chrono::steady_clock::now(); + while (true) { + std::vector tablets = {tablet_id}; + std::map done; + downloader.check_download_task(tablets, &done); + + if (done.contains(tablet_id) && done[tablet_id]) { + return true; + } + + auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count(); + if (elapsed >= timeout_seconds) { + return false; + } + + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + } + + // Helper to wait for multiple tablets + bool wait_for_tasks_done(FileCacheBlockDownloader& downloader, + const std::vector& tablet_ids, int timeout_seconds = 10) { + auto start = std::chrono::steady_clock::now(); + while (true) { + std::map done; + downloader.check_download_task(tablet_ids, &done); + + bool all_done = true; + for (int64_t tablet_id : tablet_ids) { + if (!done.contains(tablet_id) || !done[tablet_id]) { + all_done = false; + break; + } + } + + if (all_done) { + return true; + } + + auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count(); + if (elapsed >= timeout_seconds) { + return false; + } + + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + } + +protected: + CloudStorageEngine _engine; +}; + +// Test that inflight count is correctly decremented when tablet is not found +// This tests the bug fix where early return did not decrement inflight count +TEST_F(FileCacheBlockDownloaderTest, TestInflightCountDecrementOnTabletNotFound) { + FileCacheBlockDownloader downloader(_engine); + + // Create a download task with a non-existent tablet + DownloadTask::FileCacheBlockMetaVec metas; + auto* meta = metas.Add(); + meta->set_tablet_id(99999999); // Non-existent tablet ID + meta->set_rowset_id("non_existent_rowset"); + meta->set_segment_id(0); + meta->set_offset(0); + meta->set_size(1024); + + DownloadTask task(std::move(metas)); + + // Submit the task - this increments inflight count + downloader.submit_download_task(std::move(task)); + + // Wait for the task to be processed with timeout + bool done = wait_for_task_done(downloader, 99999999); + + // The task should be marked as done because inflight count was decremented + EXPECT_TRUE(done) << "Inflight count should be decremented even when tablet not found"; +} + +// Test that inflight count is correctly decremented for multiple metas +// when some fail and some succeed (or all fail) +TEST_F(FileCacheBlockDownloaderTest, TestInflightCountDecrementOnMultipleMetasFail) { + FileCacheBlockDownloader downloader(_engine); + + // Create a download task with multiple non-existent tablets + DownloadTask::FileCacheBlockMetaVec metas; + + std::vector tablet_ids; + // Add multiple metas for different tablets + for (int i = 0; i < 5; i++) { + auto* meta = metas.Add(); + int64_t tablet_id = 88888880 + i; + meta->set_tablet_id(tablet_id); // Non-existent tablet IDs + meta->set_rowset_id("non_existent_rowset_" + std::to_string(i)); + meta->set_segment_id(0); + meta->set_offset(0); + meta->set_size(1024); + tablet_ids.push_back(tablet_id); + } + + DownloadTask task(std::move(metas)); + + // Submit the task + downloader.submit_download_task(std::move(task)); + + // Wait for all tasks to be processed + bool all_done = wait_for_tasks_done(downloader, tablet_ids); + + // All tasks should be marked as done + EXPECT_TRUE(all_done) << "All inflight counts should be decremented for non-existent tablets"; +} + +// Test that inflight count is correctly decremented for same tablet with multiple blocks +TEST_F(FileCacheBlockDownloaderTest, TestInflightCountDecrementSameTabletMultipleBlocks) { + FileCacheBlockDownloader downloader(_engine); + + // Create a download task with multiple blocks for the same tablet + DownloadTask::FileCacheBlockMetaVec metas; + + int64_t tablet_id = 77777777; + for (int i = 0; i < 3; i++) { + auto* meta = metas.Add(); + meta->set_tablet_id(tablet_id); // Same tablet + meta->set_rowset_id("non_existent_rowset"); + meta->set_segment_id(i); + meta->set_offset(i * 1024); + meta->set_size(1024); + } + + DownloadTask task(std::move(metas)); + + // Submit the task - this increments inflight count by 3 for the same tablet + downloader.submit_download_task(std::move(task)); + + // Wait for the task to be processed + bool done = wait_for_task_done(downloader, tablet_id); + + // The task should be marked as done + // All 3 blocks should have decremented the count, resulting in removal from map + EXPECT_TRUE(done) << "Inflight count should be zero after all blocks processed for tablet " + << tablet_id; +} + +// Test check_download_task returns true for tablets that were never submitted +TEST_F(FileCacheBlockDownloaderTest, TestCheckDownloadTaskNonExistentTablet) { + FileCacheBlockDownloader downloader(_engine); + + // Check a tablet that was never submitted + std::vector tablets = {11111111}; + std::map done; + downloader.check_download_task(tablets, &done); + + // Should return true (done) because it's not in inflight map + ASSERT_TRUE(done.contains(11111111)); + EXPECT_TRUE(done[11111111]) << "Non-existent tablet should be reported as done"; +} + +// Test that check_download_task correctly reports in-flight status +TEST_F(FileCacheBlockDownloaderTest, TestCheckDownloadTaskInflightStatus) { + FileCacheBlockDownloader downloader(_engine); + + int64_t tablet_id = 66666666; + + // Initially, should be done (not in inflight) + { + std::vector tablets = {tablet_id}; + std::map done; + downloader.check_download_task(tablets, &done); + EXPECT_TRUE(done[tablet_id]) << "Initially tablet should not be in inflight"; + } + + // Submit a task + DownloadTask::FileCacheBlockMetaVec metas; + auto* meta = metas.Add(); + meta->set_tablet_id(tablet_id); + meta->set_rowset_id("non_existent_rowset"); + meta->set_segment_id(0); + meta->set_offset(0); + meta->set_size(1024); + + DownloadTask task(std::move(metas)); + downloader.submit_download_task(std::move(task)); + + // After task completes (should be quick since tablet doesn't exist) + bool done = wait_for_task_done(downloader, tablet_id); + EXPECT_TRUE(done) << "Task should complete with decremented inflight count"; +} + +} // namespace doris::io