diff --git a/be/src/cloud/cloud_engine_calc_delete_bitmap_task.cpp b/be/src/cloud/cloud_engine_calc_delete_bitmap_task.cpp index e4c815c913bcaf..7eca5005085c13 100644 --- a/be/src/cloud/cloud_engine_calc_delete_bitmap_task.cpp +++ b/be/src/cloud/cloud_engine_calc_delete_bitmap_task.cpp @@ -26,8 +26,10 @@ #include "cloud/cloud_meta_mgr.h" #include "cloud/cloud_tablet.h" +#include "cloud/cloud_txn_delete_bitmap_cache.h" #include "common/status.h" #include "runtime/memory/mem_tracker_limiter.h" +#include "storage/delete/calc_delete_bitmap_executor.h" #include "storage/olap_common.h" #include "storage/rowset/beta_rowset.h" #include "storage/rowset/rowset.h" @@ -67,9 +69,7 @@ Status CloudEngineCalcDeleteBitmapTask::execute() { int64_t transaction_id = _cal_delete_bitmap_req.transaction_id; OlapStopWatch watch; VLOG_NOTICE << "begin to calculate delete bitmap. transaction_id=" << transaction_id; - std::unique_ptr token = - _engine.calc_tablet_delete_bitmap_task_thread_pool().new_token( - ThreadPool::ExecutionMode::CONCURRENT); + std::vector> tokens; DBUG_EXECUTE_IF("CloudEngineCalcDeleteBitmapTask.execute.enable_wait", { auto sleep_time = DebugPoints::instance()->get_debug_param_or_default( "CloudEngineCalcDeleteBitmapTask.execute.enable_wait", "sleep_time", 3); @@ -93,6 +93,19 @@ Status CloudEngineCalcDeleteBitmapTask::execute() { if (has_tablet_states) { tablet_calc_delete_bitmap_ptr->set_tablet_state(partition.tablet_states[i]); } + // A partition's first subtransaction need not have written this tablet. + auto wg = + _engine.txn_delete_bitmap_cache().get_workload_group(transaction_id, tablet_id); + for (auto sub_txn_id : partition.sub_txn_ids) { + if (wg) { + break; + } + wg = _engine.txn_delete_bitmap_cache().get_workload_group(sub_txn_id, tablet_id); + } + auto& token = + tokens.emplace_back(_engine.calc_delete_bitmap_executor()->create_load_token( + transaction_id, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT, + std::move(wg))); const auto submit_time_us = MonotonicMicros(); auto submit_st = token->submit_func( [tablet_id, tablet_calc_delete_bitmap_ptr, this, submit_time_us]() { @@ -104,16 +117,28 @@ Status CloudEngineCalcDeleteBitmapTask::execute() { LOG(WARNING) << "handle calc delete bitmap fail, st=" << st.to_string(); add_error_tablet_id(tablet_id, st); } + return Status::OK(); }); VLOG_DEBUG << "submit TabletCalcDeleteBitmapTask for tablet=" << tablet_id; if (!submit_st.ok()) { - _res = submit_st; + add_error_tablet_id(tablet_id, submit_st); break; } } } - // wait for all finished - token->wait(); + // Drain every submitted token before reading _res, which running tablet + // callbacks may still update. Keep the recorded tablet/submission error + // ahead of a generic cancellation reported by another token's wait(). + Status wait_status; + for (auto& token : tokens) { + auto st = token->wait(); + if (wait_status.ok() && !st.ok()) { + wait_status = st; + } + } + if (_res.ok()) { + _res = wait_status; + } LOG(INFO) << "finish to calculate delete bitmap on transaction." << "transaction_id=" << transaction_id << ", cost(us): " << watch.get_elapse_time_us() @@ -149,7 +174,8 @@ void CloudTabletCalcDeleteBitmapTask::set_tablet_state(int64_t tablet_state) { Status CloudTabletCalcDeleteBitmapTask::handle(int64_t queue_time_us) const { VLOG_DEBUG << "start calculate delete bitmap on tablet " << _tablet_id << ", txn_id=" << _transaction_id; - SCOPED_ATTACH_TASK(_mem_tracker); + // The bitmap token attaches the request context at the worker entry. + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(_mem_tracker); int64_t t1 = MonotonicMicros(); auto base_tablet = DORIS_TRY(_engine.get_tablet(_tablet_id)); auto get_tablet_time_us = MonotonicMicros() - t1; diff --git a/be/src/cloud/cloud_meta_mgr.cpp b/be/src/cloud/cloud_meta_mgr.cpp index be2124d2c50e00..159942ecd39737 100644 --- a/be/src/cloud/cloud_meta_mgr.cpp +++ b/be/src/cloud/cloud_meta_mgr.cpp @@ -125,7 +125,9 @@ Status bthread_fork_join(const std::vector>& tasks, int }); bthread_t bthread_id; - if (bthread_start_background(&bthread_id, nullptr, run_bthread_work, fn) != 0) { + if (SYNC_POINT_HOOK_RETURN_VALUE( + bthread_start_background(&bthread_id, nullptr, run_bthread_work, fn), + "bthread_fork_join::start_background") != 0) { run_bthread_work(fn); } } diff --git a/be/src/cloud/cloud_rowset_builder.cpp b/be/src/cloud/cloud_rowset_builder.cpp index d9065c340ece9e..b2ac3652342dba 100644 --- a/be/src/cloud/cloud_rowset_builder.cpp +++ b/be/src/cloud/cloud_rowset_builder.cpp @@ -117,7 +117,8 @@ Status CloudRowsetBuilder::init() { _rowset_writer = DORIS_TRY(_tablet->create_rowset_writer(context, false)); _rowset_id = context.rowset_id; - _calc_delete_bitmap_token = _engine.calc_delete_bitmap_executor()->create_token(); + _calc_delete_bitmap_token = _engine.calc_delete_bitmap_executor()->create_load_token( + _req.txn_id, LoadTaskPriority::HIGH, LoadTaskType::LEAF); if (!_skip_writing_rowset_metadata) { RETURN_IF_ERROR(_engine.meta_mgr().prepare_rowset(*_rowset_writer->rowset_meta(), "", @@ -254,7 +255,7 @@ Status CloudRowsetBuilder::set_txn_related_info() { // For empty rowsets when skip_writing_empty_rowset_metadata=true, // store only a lightweight marker instead of full rowset info. // This allows CalcDeleteBitmapTask to detect and skip gracefully, - // while using minimal memory (~16 bytes per entry). + // while retaining the workload group for publish routing. if (_skip_writing_rowset_metadata) { _engine.txn_delete_bitmap_cache().mark_empty_rowset(_req.txn_id, _tablet->tablet_id(), _req.txn_expiration); diff --git a/be/src/cloud/cloud_rowset_writer.cpp b/be/src/cloud/cloud_rowset_writer.cpp index 9753ca3dd99cd9..79e8d758c400d4 100644 --- a/be/src/cloud/cloud_rowset_writer.cpp +++ b/be/src/cloud/cloud_rowset_writer.cpp @@ -89,7 +89,11 @@ Status CloudRowsetWriter::init(const RowsetWriterContext& rowset_writer_context) _context.segment_collector = std::make_shared>(this); _context.file_writer_creator = std::make_shared>(this); if (_context.mow_context != nullptr) { - _calc_delete_bitmap_token = _engine.calc_delete_bitmap_executor_for_load()->create_token(); + _calc_delete_bitmap_token = _engine.calc_delete_bitmap_executor()->create_load_token( + _context.txn_id, + _context.is_transient_rowset_writer ? LoadTaskPriority::HIGHEST + : LoadTaskPriority::MID, + LoadTaskType::LEAF); } return Status::OK(); } diff --git a/be/src/cloud/cloud_storage_engine.cpp b/be/src/cloud/cloud_storage_engine.cpp index b9d88260ee4e11..459796b114cfa6 100644 --- a/be/src/cloud/cloud_storage_engine.cpp +++ b/be/src/cloud/cloud_storage_engine.cpp @@ -220,14 +220,8 @@ Status CloudStorageEngine::open() { _calc_delete_bitmap_executor = std::make_unique(); _calc_delete_bitmap_executor->init("TabletCalcDeleteBitmapThreadPool", - config::calc_delete_bitmap_max_thread); - - _calc_delete_bitmap_executor_for_load = std::make_unique(); - _calc_delete_bitmap_executor_for_load->init( - "LoadCalcDeleteBitmapThreadPool", - config::calc_delete_bitmap_for_load_max_thread > 0 - ? config::calc_delete_bitmap_for_load_max_thread - : std::max(1, CpuInfo::num_cores() / 2)); + config::calc_delete_bitmap_max_thread, + _memtable_flush_executor->flush_pool()); // The default cache is set to 100MB, use memory limit to dynamic adjustment bool is_percent = false; @@ -275,18 +269,15 @@ Status CloudStorageEngine::open() { #ifdef BE_TEST void CloudStorageEngine::init_calc_delete_bitmap_executor_for_UT() { + if (_memtable_flush_executor == nullptr) { + _memtable_flush_executor = std::make_unique(); + _memtable_flush_executor->init(1); + } if (_calc_delete_bitmap_executor == nullptr) { _calc_delete_bitmap_executor = std::make_unique(); _calc_delete_bitmap_executor->init("TabletCalcDeleteBitmapThreadPool", - config::calc_delete_bitmap_max_thread); - } - if (_calc_delete_bitmap_executor_for_load == nullptr) { - _calc_delete_bitmap_executor_for_load = std::make_unique(); - _calc_delete_bitmap_executor_for_load->init( - "LoadCalcDeleteBitmapThreadPool", - config::calc_delete_bitmap_for_load_max_thread > 0 - ? config::calc_delete_bitmap_for_load_max_thread - : std::max(1, CpuInfo::num_cores() / 2)); + config::calc_delete_bitmap_max_thread, + _memtable_flush_executor->flush_pool()); } } #endif @@ -317,9 +308,6 @@ void CloudStorageEngine::stop() { _adaptive_thread_controller.stop(); LOG(INFO) << "Cloud storage engine is stopped."; - if (_calc_tablet_delete_bitmap_task_thread_pool) { - _calc_tablet_delete_bitmap_task_thread_pool->shutdown(); - } if (_sync_delete_bitmap_thread_pool) { _sync_delete_bitmap_thread_pool->shutdown(); } @@ -391,11 +379,6 @@ Status CloudStorageEngine::start_bg_threads(std::shared_ptr wg_sp &_id_file_map_gc_thread)); LOG(INFO) << "id file map gc thread started"; - // add calculate tablet delete bitmap task thread pool - RETURN_IF_ERROR(ThreadPoolBuilder("TabletCalDeleteBitmapThreadPool") - .set_min_threads(config::calc_tablet_delete_bitmap_task_max_thread) - .set_max_threads(config::calc_tablet_delete_bitmap_task_max_thread) - .build(&_calc_tablet_delete_bitmap_task_thread_pool)); RETURN_IF_ERROR(ThreadPoolBuilder("SyncDeleteBitmapThreadPool") .set_min_threads(config::sync_delete_bitmap_task_max_thread) .set_max_threads(config::sync_delete_bitmap_task_max_thread) diff --git a/be/src/cloud/cloud_storage_engine.h b/be/src/cloud/cloud_storage_engine.h index 330850b5ed8192..f367aaa6a2c9b5 100644 --- a/be/src/cloud/cloud_storage_engine.h +++ b/be/src/cloud/cloud_storage_engine.h @@ -97,9 +97,6 @@ class CloudStorageEngine final : public BaseStorageEngine { CloudCommittedRSMgr& committed_rs_mgr() const { return *_committed_rs_mgr; } - ThreadPool& calc_tablet_delete_bitmap_task_thread_pool() const { - return *_calc_tablet_delete_bitmap_task_thread_pool; - } ThreadPool& sync_delete_bitmap_thread_pool() const { return *_sync_delete_bitmap_thread_pool; } std::optional get_storage_resource(const std::string& vault_id) { @@ -234,7 +231,6 @@ class CloudStorageEngine final : public BaseStorageEngine { std::unique_ptr _tablet_mgr; std::unique_ptr _txn_delete_bitmap_cache; std::unique_ptr _committed_rs_mgr; - std::unique_ptr _calc_tablet_delete_bitmap_task_thread_pool; std::unique_ptr _sync_delete_bitmap_thread_pool; // Components for cache warmup diff --git a/be/src/cloud/cloud_tablets_channel.cpp b/be/src/cloud/cloud_tablets_channel.cpp index 9a5f25c0439aeb..47da15c0f71aad 100644 --- a/be/src/cloud/cloud_tablets_channel.cpp +++ b/be/src/cloud/cloud_tablets_channel.cpp @@ -25,6 +25,7 @@ #include "cloud/config.h" #include "load/channel/tablets_channel.h" #include "load/delta_writer/delta_writer.h" +#include "runtime/thread_context.h" #include "storage/tablet_info.h" namespace doris { @@ -278,7 +279,17 @@ Status CloudTabletsChannel::close(LoadChannel* parent, const PTabletWriterAddBlo std::vector> tasks; tasks.reserve(writers_to_commit.size()); for (auto* writer : writers_to_commit) { - tasks.emplace_back([writer] { return writer->commit_rowset(); }); + tasks.emplace_back([writer, caller_bthread_id = bthread_self()] { + // bthread_fork_join runs inline if starting a bthread fails. The caller + // already has the load context attached; do not attach it twice. + if (bthread_self() == caller_bthread_id) { + return writer->commit_rowset(); + } + // Empty writers initialize their rowset and bitmap token during commit. + // The new bthread must inherit the load's memory tracker and workload group. + SCOPED_ATTACH_TASK(writer->resource_context()); + return writer->commit_rowset(); + }); } _close_status = cloud::bthread_fork_join(tasks, 10); if (!_close_status.ok()) { diff --git a/be/src/cloud/cloud_txn_delete_bitmap_cache.cpp b/be/src/cloud/cloud_txn_delete_bitmap_cache.cpp index 83a7528a6c662b..7788daa9b6d83c 100644 --- a/be/src/cloud/cloud_txn_delete_bitmap_cache.cpp +++ b/be/src/cloud/cloud_txn_delete_bitmap_cache.cpp @@ -26,6 +26,8 @@ #include "cloud/config.h" #include "common/status.h" #include "cpp/sync_point.h" +#include "runtime/thread_context.h" +#include "runtime/workload_management/resource_context.h" #include "storage/olap_common.h" #include "storage/rowset/rowset_fwd.h" #include "storage/tablet/tablet_meta.h" @@ -187,6 +189,19 @@ Status CloudTxnDeleteBitmapCache::get_delete_bitmap( return Status::OK(); } +std::shared_ptr CloudTxnDeleteBitmapCache::get_workload_group( + TTransactionId transaction_id, int64_t tablet_id) { + std::shared_lock rlock(_rwlock); + TxnKey key(transaction_id, tablet_id); + auto it = _txn_map.find(key); + if (it != _txn_map.end()) { + return it->second.workload_group; + } + auto marker = _empty_rowset_markers.find(key); + // A retried request on another BE has no local owner. + return marker == _empty_rowset_markers.end() ? nullptr : marker->second.workload_group; +} + void CloudTxnDeleteBitmapCache::set_tablet_txn_info( TTransactionId transaction_id, int64_t tablet_id, DeleteBitmapPtr delete_bitmap, const RowsetIdUnorderedSet& rowset_ids, RowsetSharedPtr rowset, int64_t txn_expiration, @@ -204,6 +219,11 @@ void CloudTxnDeleteBitmapCache::set_tablet_txn_info( std::make_shared(PublishStatus::INIT); _txn_map[txn_key] = TxnVal(rowset, txn_expiration, std::move(partial_update_info), std::move(publish_status), attach_row_binlog); + // Cloud DELETE agent tasks have no attached resource context and use the + // default load pool. Only capture a workload group from an attached task. + auto* ctx = thread_context(); + _txn_map[txn_key].workload_group = + ctx->is_attach_task() ? ctx->resource_ctx()->workload_group() : nullptr; _expiration_txn.emplace(txn_expiration, txn_key); } std::string key_str = fmt::format("{}/{}", transaction_id, tablet_id); @@ -304,7 +324,8 @@ void CloudTxnDeleteBitmapCache::remove_expired_tablet_txn_info() { } // Clean from _empty_rowset_markers if exists auto marker_iter = _empty_rowset_markers.find(iter->second); - if (marker_iter != _empty_rowset_markers.end()) { + if (marker_iter != _empty_rowset_markers.end() && + iter->first == marker_iter->second.txn_expiration) { LOG_INFO("clean expired empty rowset marker") .tag("txn_id", iter->second.txn_id) .tag("tablet_id", iter->second.tablet_id); @@ -329,6 +350,7 @@ void CloudTxnDeleteBitmapCache::remove_unused_tablet_txn_info(TTransactionId tra erase(cache_key); _txn_map.erase(txn_key); } + _empty_rowset_markers.erase(txn_key); } void CloudTxnDeleteBitmapCache::mark_empty_rowset(TTransactionId txn_id, int64_t tablet_id, @@ -347,7 +369,10 @@ void CloudTxnDeleteBitmapCache::mark_empty_rowset(TTransactionId txn_id, int64_t } std::unique_lock wlock(_rwlock); TxnKey txn_key(txn_id, tablet_id); - _empty_rowset_markers.emplace(txn_key); + auto* ctx = thread_context(); + _empty_rowset_markers[txn_key] = { + ctx->is_attach_task() ? ctx->resource_ctx()->workload_group() : nullptr, + txn_expiration}; _expiration_txn.emplace(txn_expiration, txn_key); } @@ -364,4 +389,4 @@ void CloudTxnDeleteBitmapCache::_clean_thread_callback() { std::chrono::seconds(config::remove_expired_tablet_txn_info_interval_seconds))); } -} // namespace doris \ No newline at end of file +} // namespace doris diff --git a/be/src/cloud/cloud_txn_delete_bitmap_cache.h b/be/src/cloud/cloud_txn_delete_bitmap_cache.h index 6be96990ed7f2b..a8e8ec630be2bf 100644 --- a/be/src/cloud/cloud_txn_delete_bitmap_cache.h +++ b/be/src/cloud/cloud_txn_delete_bitmap_cache.h @@ -29,6 +29,8 @@ namespace doris { +class WorkloadGroup; + // Record transaction related delete bitmaps using a lru cache. class CloudTxnDeleteBitmapCache : public LRUCachePolicy { public: @@ -38,6 +40,9 @@ class CloudTxnDeleteBitmapCache : public LRUCachePolicy { Status init(); + std::shared_ptr get_workload_group(TTransactionId transaction_id, + int64_t tablet_id); + Status get_tablet_txn_info(TTransactionId transaction_id, int64_t tablet_id, RowsetSharedPtr* rowset, DeleteBitmapPtr* delete_bitmap, RowsetIdUnorderedSet* rowset_ids, int64_t* txn_expiration, @@ -68,7 +73,7 @@ class CloudTxnDeleteBitmapCache : public LRUCachePolicy { // Check if this is a known empty/skipped rowset // Returns true if was marked as empty rowset // Note: Does not remove the marker, as CalcDeleteBitmapTask may retry. - // Cleanup is handled by expiration-based removal in remove_expired_tablet_txn_info() + // Cleanup is handled by remove_expired_tablet_txn_info() or remove_unused_tablet_txn_info(). bool is_empty_rowset(TTransactionId txn_id, int64_t tablet_id); // !!!ATTENTION!!!: the delete bitmap stored in CloudTxnDeleteBitmapCache contains sentinel marks, @@ -105,6 +110,8 @@ class CloudTxnDeleteBitmapCache : public LRUCachePolicy { }; struct TxnVal { + // Preserve write-stage resource isolation through commit/retries. + std::shared_ptr workload_group; RowsetSharedPtr rowset; int64_t txn_expiration; std::shared_ptr partial_update_info; @@ -124,11 +131,16 @@ class CloudTxnDeleteBitmapCache : public LRUCachePolicy { attach_row_binlog(attach_row_binlog_) {} }; + struct EmptyRowsetMarker { + std::shared_ptr workload_group; + int64_t txn_expiration; + }; + std::map _txn_map; std::multimap _expiration_txn; - // Lightweight markers for empty/skipped rowsets (only stores TxnKey, ~16 bytes per entry) - // Used to track empty rowsets that were not committed to meta-service - std::set _empty_rowset_markers; + // Empty/skipped rowsets have no metadata, but their publish tasks still need + // the write-stage resource domain until the marker is removed or expires. + std::map _empty_rowset_markers; std::shared_mutex _rwlock; std::shared_ptr _clean_thread; CountDownLatch _stop_latch; diff --git a/be/src/common/config.cpp b/be/src/common/config.cpp index a0b479a4eeaefa..b15062337d0009 100644 --- a/be/src/common/config.cpp +++ b/be/src/common/config.cpp @@ -857,14 +857,14 @@ DEFINE_mInt32(storage_flood_stage_usage_percent, "90"); // 90% DEFINE_mInt64(storage_flood_stage_left_capacity_bytes, "1073741824"); // 1GB // number of thread for flushing memtable per store DEFINE_mInt32(flush_thread_num_per_store, "6"); -// number of thread for flushing memtable per store, for high priority load task +// Deprecated compatibility setting; foreground load tasks share the normal flush pool. DEFINE_mInt32(high_priority_flush_thread_num_per_store, "6"); -// number of threads = min(flush_thread_num_per_store * num_store, -// max_flush_thread_num_per_cpu * num_cpu) -DEFINE_mInt32(max_flush_thread_num_per_cpu, "4"); +// Maximum shared foreground load threads per CPU (default 8). +// Without adaptive flushing, also capped by flush_thread_num_per_store * num_store. +DEFINE_mInt32(max_flush_thread_num_per_cpu, "8"); -// minimum flush threads per cpu when adaptive flush is enabled (default 0.5) -DEFINE_mDouble(min_flush_thread_num_per_cpu, "0.5"); +// Minimum shared foreground load threads per CPU when adaptive flushing is enabled (default 1). +DEFINE_mDouble(min_flush_thread_num_per_cpu, "1.0"); // Whether to enable adaptive flush thread adjustment DEFINE_mBool(enable_adaptive_flush_threads, "true"); @@ -2438,9 +2438,7 @@ void update_config(const std::string& field, const std::string& value) { if ("sys_log_level" == field) { // update log level update_logging(field, value); - } else if ("flush_thread_num_per_store" == field || - "high_priority_flush_thread_num_per_store" == field || - "max_flush_thread_num_per_cpu" == field) { + } else if ("flush_thread_num_per_store" == field || "max_flush_thread_num_per_cpu" == field) { // update memtable flush thread pool size auto* exec_env = ExecEnv::GetInstance(); if (exec_env != nullptr) { diff --git a/be/src/common/config.h b/be/src/common/config.h index 6bda64a14553a2..99955ee2bdab4f 100644 --- a/be/src/common/config.h +++ b/be/src/common/config.h @@ -938,10 +938,10 @@ DECLARE_mInt64(storage_flood_stage_left_capacity_bytes); // 1GB DECLARE_mInt32(flush_thread_num_per_store); // number of thread for flushing memtable per store, for high priority load task DECLARE_mInt32(high_priority_flush_thread_num_per_store); -// number of threads = min(flush_thread_num_per_store * num_store, -// max_flush_thread_num_per_cpu * num_cpu) +// Maximum shared foreground load threads per CPU (default 8). +// Without adaptive flushing, also capped by flush_thread_num_per_store * num_store. DECLARE_mInt32(max_flush_thread_num_per_cpu); -// minimum flush threads per cpu when adaptive flush is enabled (default 0.5) +// Minimum shared foreground load threads per CPU when adaptive flushing is enabled (default 1). DECLARE_mDouble(min_flush_thread_num_per_cpu); // Whether to enable adaptive flush thread adjustment diff --git a/be/src/load/channel/load_stream_writer.cpp b/be/src/load/channel/load_stream_writer.cpp index df67245ef375d7..4e1b34081cc10d 100644 --- a/be/src/load/channel/load_stream_writer.cpp +++ b/be/src/load/channel/load_stream_writer.cpp @@ -316,6 +316,10 @@ Status LoadStreamWriter::close() { if (!_pre_closed) { RETURN_IF_ERROR(_pre_close()); } + // The heavy-work callback has no attached load context, and _pre_close() + // detaches its context on return. Keep the owner attached through commit so + // the transaction retains its workload group for later publish work. + SCOPED_ATTACH_TASK(_resource_ctx); RETURN_IF_ERROR(_rowset_builder->wait_calc_delete_bitmap()); // FIXME(plat1ko): No `commit_txn` operation in cloud mode, need better abstractions RETURN_IF_ERROR(static_cast(_rowset_builder.get())->commit_txn()); diff --git a/be/src/load/memtable/memtable_flush_executor.cpp b/be/src/load/memtable/memtable_flush_executor.cpp index 2ad61661d4f854..e5ba0abc54e217 100644 --- a/be/src/load/memtable/memtable_flush_executor.cpp +++ b/be/src/load/memtable/memtable_flush_executor.cpp @@ -145,7 +145,8 @@ Status FlushToken::_submit_sub_tasks(ThreadPool* pool, return _flush_status; } } - Status submit_st = pool->submit(std::move(sub_tasks[i])); + Status submit_st = pool->submit_load( + std::move(sub_tasks[i]), _rowset_writer->context().txn_id, LoadTaskPriority::LOW); if (UNLIKELY(!submit_st.ok())) { { std::lock_guard wrlk(_flush_status_lock); @@ -489,13 +490,6 @@ void MemTableFlushExecutor::init(int num_disk) { .set_min_threads(min_threads) .set_max_threads(max_threads) .build(&_flush_pool)); - - auto [hi_min, hi_max] = calc_flush_thread_count( - num_cpus, _num_disk, config::high_priority_flush_thread_num_per_store); - static_cast(ThreadPoolBuilder("MemTableHighPriorityFlushThreadPool") - .set_min_threads(hi_min) - .set_max_threads(hi_max) - .build(&_high_prio_flush_pool)); } void MemTableFlushExecutor::update_memtable_flush_threads() { @@ -506,18 +500,13 @@ void MemTableFlushExecutor::update_memtable_flush_threads() { // Update max_threads first to avoid constraint violation when increasing min_threads static_cast(_flush_pool->set_max_threads(max_threads)); static_cast(_flush_pool->set_min_threads(min_threads)); - - auto [hi_min, hi_max] = calc_flush_thread_count( - num_cpus, _num_disk, config::high_priority_flush_thread_num_per_store); - // Update max_threads first to avoid constraint violation when increasing min_threads - static_cast(_high_prio_flush_pool->set_max_threads(hi_max)); - static_cast(_high_prio_flush_pool->set_min_threads(hi_min)); } -// NOTE: we use SERIAL mode here to ensure all mem-tables from one tablet are flushed in order. +// Each resource domain shares workers across foreground load tasks. Stage priority is +// applied within a load; is_high_priority no longer selects a separate pool. Status MemTableFlushExecutor::create_flush_token( std::shared_ptr& flush_token, std::shared_ptr rowset_writer, - bool is_high_priority, std::shared_ptr wg_sptr, + bool /*is_high_priority*/, std::shared_ptr wg_sptr, std::shared_ptr table_schema_param) { switch (rowset_writer->type()) { case ALPHA_ROWSET: @@ -525,7 +514,7 @@ Status MemTableFlushExecutor::create_flush_token( return Status::InternalError("not support alpha rowset load now."); case BETA_ROWSET: { // beta rowset can be flush in CONCURRENT, because each memtable using a new segment writer. - ThreadPool* pool = is_high_priority ? _high_prio_flush_pool.get() : _flush_pool.get(); + ThreadPool* pool = _flush_pool.get(); flush_token = FlushToken::create_shared(pool, wg_sptr); flush_token->set_rowset_writer(rowset_writer); flush_token->set_table_schema_param(std::move(table_schema_param)); diff --git a/be/src/load/memtable/memtable_flush_executor.h b/be/src/load/memtable/memtable_flush_executor.h index c3053d483e5e36..44e1072c9c0ddd 100644 --- a/be/src/load/memtable/memtable_flush_executor.h +++ b/be/src/load/memtable/memtable_flush_executor.h @@ -186,10 +186,7 @@ class FlushToken : public std::enable_shared_from_this { class MemTableFlushExecutor { public: MemTableFlushExecutor() = default; - ~MemTableFlushExecutor() { - _flush_pool->shutdown(); - _high_prio_flush_pool->shutdown(); - } + ~MemTableFlushExecutor() { _flush_pool->shutdown(); } // init should be called after storage engine is opened, // because it needs path hash of each data dir. @@ -218,8 +215,6 @@ class MemTableFlushExecutor { ThreadPool* flush_pool() { return _flush_pool.get(); } - ThreadPool* high_prio_flush_pool() { return _high_prio_flush_pool.get(); } - void update_memtable_flush_threads(); // Returns {min_threads, max_threads} for a flush thread pool. @@ -229,7 +224,6 @@ class MemTableFlushExecutor { private: std::unique_ptr _flush_pool; - std::unique_ptr _high_prio_flush_pool; std::atomic _flushing_task_count = 0; int _num_disk = 0; }; diff --git a/be/src/storage/delete/calc_delete_bitmap_executor.cpp b/be/src/storage/delete/calc_delete_bitmap_executor.cpp index af1b4bfb23a830..cfe15212a04853 100644 --- a/be/src/storage/delete/calc_delete_bitmap_executor.cpp +++ b/be/src/storage/delete/calc_delete_bitmap_executor.cpp @@ -21,8 +21,12 @@ #include +#include "common/check.h" #include "common/logging.h" #include "load/memtable/memtable.h" +#include "runtime/thread_context.h" +#include "runtime/workload_group/workload_group.h" +#include "runtime/workload_management/resource_context.h" #include "storage/tablet/base_tablet.h" #include "util/time.h" @@ -38,13 +42,11 @@ Status CalcDeleteBitmapToken::submit(BaseTabletSPtr tablet, RowsetSharedPtr cur_ { std::shared_lock rlock(_lock); RETURN_IF_ERROR(_status); - _resource_ctx = thread_context()->resource_ctx(); } const auto submit_time_us = MonotonicMicros(); - return _thread_token->submit_func([=, this]() { + return _submit_func([=, this]() { const auto queue_time_us = MonotonicMicros() - submit_time_us; - SCOPED_ATTACH_TASK(_resource_ctx); auto st = tablet->calc_segment_delete_bitmap(cur_rowset, cur_segment, target_rowsets, delete_bitmap, end_version, rowset_writer, tablet_delete_bitmap, queue_time_us); @@ -68,12 +70,10 @@ Status CalcDeleteBitmapToken::submit(BaseTabletSPtr tablet, TabletSchemaSPtr sch { std::shared_lock rlock(_lock); RETURN_IF_ERROR(_status); - _resource_ctx = thread_context()->resource_ctx(); } const auto submit_time_us = MonotonicMicros(); - return _thread_token->submit_func([=, this]() { + return _submit_func([=, this]() { const auto queue_time_us = MonotonicMicros() - submit_time_us; - SCOPED_ATTACH_TASK(_resource_ctx); auto st = tablet->calc_delete_bitmap_between_segments(schema, rowset_id, segments, delete_bitmap, queue_time_us); if (!st.ok()) { @@ -89,12 +89,68 @@ Status CalcDeleteBitmapToken::submit(BaseTabletSPtr tablet, TabletSchemaSPtr sch } Status CalcDeleteBitmapToken::wait() { - _thread_token->wait(); - // all tasks complete here, don't need lock; - return _status; + if (_help_while_wait) { + auto st = _thread_token->wait_and_help(); + if (!st.ok()) { + // Do not return while callbacks can still access this token or its + // caller's state. Factory-created nested tokens always own leaves. + { + std::lock_guard wlock(_lock); + if (_status.ok()) { + _status = st; + } + } + _thread_token->shutdown(); + } + } else { + _thread_token->wait(); + } + std::shared_lock rlock(_lock); + RETURN_IF_ERROR(_status); + // A workload-group shutdown may remove queued tasks without executing them. + if (_finished_tasks.load() != _submitted_tasks.load()) { + return Status::Cancelled("delete bitmap tasks were cancelled before completion"); + } + return Status::OK(); } -void CalcDeleteBitmapExecutor::init(const std::string& name, int max_threads) { +Status CalcDeleteBitmapToken::_submit_func(std::function func) { + ++_submitted_tasks; + auto task = [this, func = std::move(func)]() { + func(); + ++_finished_tasks; + }; + auto st = _thread_token->submit_func( + [task = std::move(task), resource_ctx = thread_context()->resource_ctx(), + tracker = thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker_sptr()]() { + if (ThreadPool::is_helping_load_task()) { + // The parent is already attached. AttachTask cannot be nested. + DCHECK(thread_context()->is_attach_task()); + DCHECK(thread_context()->resource_ctx() == resource_ctx); + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tracker); + task(); + } else { + SCOPED_ATTACH_TASK(resource_ctx); + // A cloud parent switches to a tablet tracker after attaching + // the request context. Preserve that tracker on spare workers too. + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tracker); + task(); + } + }); + if (!st.ok()) { + // Preserve the submission error before wait() checks for callbacks + // that did not finish; a rejection is not a generic cancellation. + std::lock_guard wlock(_lock); + if (_status.ok()) { + _status = st; + } + } + return st; +} + +void CalcDeleteBitmapExecutor::init(const std::string& name, int max_threads, + ThreadPool* load_pool) { + _load_pool = load_pool; static_cast(ThreadPoolBuilder(name) .set_min_threads(1) .set_max_threads(max_threads) @@ -106,4 +162,31 @@ std::unique_ptr CalcDeleteBitmapExecutor::create_token() _thread_pool->new_token(ThreadPool::ExecutionMode::CONCURRENT)); } +std::unique_ptr CalcDeleteBitmapExecutor::create_load_token( + int64_t load_id, LoadTaskPriority priority, LoadTaskType type) { + return create_load_token(load_id, priority, type, + thread_context()->resource_ctx()->workload_group()); +} + +std::unique_ptr CalcDeleteBitmapExecutor::create_load_token( + int64_t load_id, LoadTaskPriority priority, LoadTaskType type, + std::shared_ptr wg) { + // Nested segment calculations belong to the parent's actual pool. Its + // attached request context may not carry the workload group used to route it. + if (auto* pool = ThreadPool::current_load_pool()) { + DORIS_CHECK(type == LoadTaskType::LEAF) << "Nested bitmap tasks must be leaves"; + return std::make_unique( + pool->new_load_token(load_id, priority, type), std::move(wg), true); + } + // A commit retry can outlive a dropped workload group. Its pool is stopped; + // use the default domain in that case. A concurrent stop is reported by submit/wait. + ThreadPool* pool = wg && !wg->can_be_dropped() ? wg->get_memtable_flush_pool() : nullptr; + if (pool == nullptr) { + pool = _load_pool; + } + DCHECK(pool != nullptr); + return std::make_unique(pool->new_load_token(load_id, priority, type), + std::move(wg)); +} + } // namespace doris diff --git a/be/src/storage/delete/calc_delete_bitmap_executor.h b/be/src/storage/delete/calc_delete_bitmap_executor.h index 4f0f0fd2c7fff8..5364a76885ecf1 100644 --- a/be/src/storage/delete/calc_delete_bitmap_executor.h +++ b/be/src/storage/delete/calc_delete_bitmap_executor.h @@ -37,6 +37,7 @@ namespace doris { class DataDir; +class WorkloadGroup; class Tablet; enum RowsetTypePB : int; @@ -48,8 +49,15 @@ enum RowsetTypePB : int; // 4. call `get_delete_bitmap()` to get the result of all tasks class CalcDeleteBitmapToken { public: - explicit CalcDeleteBitmapToken(std::unique_ptr thread_token) - : _thread_token(std::move(thread_token)), _status(Status::OK()) {} + explicit CalcDeleteBitmapToken(std::unique_ptr thread_token, + std::shared_ptr workload_group = nullptr, + bool help_while_wait = false) + : _workload_group(std::move(workload_group)), + _thread_token(std::move(thread_token)), + _help_while_wait(help_while_wait), + _status(Status::OK()) {} + + ~CalcDeleteBitmapToken() { cancel(); } // calculate delete bitmap of `cur_segment` to historical `target_rowsets` Status submit(BaseTabletSPtr tablet, RowsetSharedPtr cur_rowset, @@ -69,10 +77,8 @@ class CalcDeleteBitmapToken { { std::shared_lock rlock(_lock); RETURN_IF_ERROR(_status); - _resource_ctx = thread_context()->resource_ctx(); } - return _thread_token->submit_func([this, func = std::forward(func)]() { - SCOPED_ATTACH_TASK(_resource_ctx); + return _submit_func([this, func = std::forward(func)]() { auto st = func(); if (!st.ok()) { std::lock_guard wlock(_lock); @@ -89,13 +95,19 @@ class CalcDeleteBitmapToken { void cancel() { _thread_token->shutdown(); } private: + Status _submit_func(std::function func); + + // Keep the selected workload-group pool alive until the token is destroyed. + std::shared_ptr _workload_group; std::unique_ptr _thread_token; + const bool _help_while_wait; std::shared_mutex _lock; // Records the current status of the calc delete bitmap job. // Note: Once its value is set to Failed, it cannot return to SUCCESS. Status _status; - std::shared_ptr _resource_ctx; + std::atomic _submitted_tasks {0}; + std::atomic _finished_tasks {0}; }; // CalcDeleteBitmapExecutor is responsible for calc delete bitmap concurrently. @@ -106,12 +118,21 @@ class CalcDeleteBitmapExecutor { ~CalcDeleteBitmapExecutor() { _thread_pool->shutdown(); } // init should be called after storage engine is opened, - void init(const std::string& name, int max_threads); + void init(const std::string& name, int max_threads, ThreadPool* load_pool); std::unique_ptr create_token(); + std::unique_ptr create_load_token(int64_t load_id, + LoadTaskPriority priority, + LoadTaskType type); + std::unique_ptr create_load_token(int64_t load_id, + LoadTaskPriority priority, + LoadTaskType type, + std::shared_ptr wg); + private: std::unique_ptr _thread_pool; + ThreadPool* _load_pool = nullptr; }; } // namespace doris diff --git a/be/src/storage/rowset/beta_rowset_writer.cpp b/be/src/storage/rowset/beta_rowset_writer.cpp index 265448a831472b..a2b8b7c947cf97 100644 --- a/be/src/storage/rowset/beta_rowset_writer.cpp +++ b/be/src/storage/rowset/beta_rowset_writer.cpp @@ -507,7 +507,11 @@ Status BetaRowsetWriter::init(const RowsetWriterContext& rowset_writer_context) _segcompaction_worker->init_mem_tracker(rowset_writer_context); } if (_context.mow_context != nullptr) { - _calc_delete_bitmap_token = _engine.calc_delete_bitmap_executor_for_load()->create_token(); + _calc_delete_bitmap_token = _engine.calc_delete_bitmap_executor()->create_load_token( + _context.txn_id, + _context.is_transient_rowset_writer ? LoadTaskPriority::HIGHEST + : LoadTaskPriority::MID, + LoadTaskType::LEAF); } return Status::OK(); } diff --git a/be/src/storage/rowset_builder.cpp b/be/src/storage/rowset_builder.cpp index 0741ed4582e678..1259bb3a9810ba 100644 --- a/be/src/storage/rowset_builder.cpp +++ b/be/src/storage/rowset_builder.cpp @@ -257,7 +257,8 @@ Status RowsetBuilder::init() { tmp_pending_rowset_ids.begin() + 1); _pending_rs_guard = _engine.pending_local_rowsets().add(tmp_pending_rowset_ids); - _calc_delete_bitmap_token = _engine.calc_delete_bitmap_executor()->create_token(); + _calc_delete_bitmap_token = _engine.calc_delete_bitmap_executor()->create_load_token( + _req.txn_id, LoadTaskPriority::HIGH, LoadTaskType::LEAF); _is_init = true; return Status::OK(); diff --git a/be/src/storage/storage_engine.cpp b/be/src/storage/storage_engine.cpp index dcb81f3f111159..62a385b640f3a5 100644 --- a/be/src/storage/storage_engine.cpp +++ b/be/src/storage/storage_engine.cpp @@ -151,8 +151,7 @@ void BaseStorageEngine::_start_adaptive_thread_controller() { if (_memtable_flush_executor) { auto* flush_pool = _memtable_flush_executor->flush_pool(); - auto* high_prio_pool = _memtable_flush_executor->high_prio_flush_pool(); - _adaptive_thread_controller.add("flush", {flush_pool, high_prio_pool}, + _adaptive_thread_controller.add("flush", {flush_pool}, AdaptiveThreadPoolController::make_flush_adjust_func( &_adaptive_thread_controller, flush_pool), config::max_flush_thread_num_per_cpu, @@ -327,14 +326,8 @@ Status StorageEngine::_open() { _calc_delete_bitmap_executor = std::make_unique(); _calc_delete_bitmap_executor->init("TabletCalcDeleteBitmapThreadPool", - config::calc_delete_bitmap_max_thread); - - _calc_delete_bitmap_executor_for_load = std::make_unique(); - _calc_delete_bitmap_executor_for_load->init( - "LoadCalcDeleteBitmapThreadPool", - config::calc_delete_bitmap_for_load_max_thread > 0 - ? config::calc_delete_bitmap_for_load_max_thread - : std::max(1, CpuInfo::num_cores() / 2)); + config::calc_delete_bitmap_max_thread, + _memtable_flush_executor->flush_pool()); _parse_default_rowset_type(); @@ -795,7 +788,6 @@ void StorageEngine::stop() { _adaptive_thread_controller.stop(); _memtable_flush_executor.reset(nullptr); _calc_delete_bitmap_executor.reset(nullptr); - _calc_delete_bitmap_executor_for_load.reset(); _stopped = true; LOG(INFO) << "Storage engine is stopped."; diff --git a/be/src/storage/storage_engine.h b/be/src/storage/storage_engine.h index d029d43662403a..8712cb836795d6 100644 --- a/be/src/storage/storage_engine.h +++ b/be/src/storage/storage_engine.h @@ -147,10 +147,6 @@ class BaseStorageEngine { return _calc_delete_bitmap_executor.get(); } - CalcDeleteBitmapExecutor* calc_delete_bitmap_executor_for_load() { - return _calc_delete_bitmap_executor_for_load.get(); - } - int64_t memory_limitation_bytes_per_thread_for_schema_change() const; int get_disk_num() { return _disk_num; } @@ -178,7 +174,6 @@ class BaseStorageEngine { std::unique_ptr _memtable_flush_executor; AdaptiveThreadPoolController _adaptive_thread_controller; std::unique_ptr _calc_delete_bitmap_executor; - std::unique_ptr _calc_delete_bitmap_executor_for_load; CountDownLatch _stop_background_threads_latch; std::shared_ptr _id_file_map_gc_thread; diff --git a/be/src/storage/tablet/base_tablet.cpp b/be/src/storage/tablet/base_tablet.cpp index c7ad7a6c639e93..b1d770d517ba21 100644 --- a/be/src/storage/tablet/base_tablet.cpp +++ b/be/src/storage/tablet/base_tablet.cpp @@ -1696,7 +1696,8 @@ Status BaseTablet::update_delete_bitmap(const BaseTabletSPtr& self, TabletTxnInf }); if (!rowsets_skip_alignment.empty()) { - auto token = self->calc_delete_bitmap_executor()->create_token(); + auto token = self->calc_delete_bitmap_executor()->create_load_token( + txn_id, LoadTaskPriority::HIGHEST, LoadTaskType::LEAF, txn_info->workload_group); // set rowset_writer to nullptr to skip the alignment process RETURN_IF_ERROR(calc_delete_bitmap(self, rowset, segments, rowsets_skip_alignment, delete_bitmap, cur_version - 1, token.get(), nullptr, @@ -1745,15 +1746,15 @@ Status BaseTablet::update_delete_bitmap(const BaseTabletSPtr& self, TabletTxnInf transient_rs_writer = std::move(group_writer); } - // When there is only one segment, it will be calculated in the current thread. - // Otherwise, it will be submitted to the thread pool for calculation. - if (segments.size() <= 1) { + // Preserve the local single-segment fast path while holding the tablet lock. + // Load workers submit leaves and help them without another resource-context attach. + if (segments.size() <= 1 && ThreadPool::current_load_pool() == nullptr) { RETURN_IF_ERROR(calc_delete_bitmap(self, rowset, segments, specified_rowsets, delete_bitmap, cur_version - 1, nullptr, transient_rs_writer.get(), tablet_delete_bitmap)); - } else { - auto token = self->calc_delete_bitmap_executor()->create_token(); + auto token = self->calc_delete_bitmap_executor()->create_load_token( + txn_id, LoadTaskPriority::HIGHEST, LoadTaskType::LEAF, txn_info->workload_group); RETURN_IF_ERROR(calc_delete_bitmap(self, rowset, segments, specified_rowsets, delete_bitmap, cur_version - 1, token.get(), transient_rs_writer.get(), tablet_delete_bitmap)); diff --git a/be/src/storage/txn/txn_manager.cpp b/be/src/storage/txn/txn_manager.cpp index 4c079eda1bc12a..3b56e3bc525c63 100644 --- a/be/src/storage/txn/txn_manager.cpp +++ b/be/src/storage/txn/txn_manager.cpp @@ -37,6 +37,8 @@ #include "common/logging.h" #include "common/status.h" #include "load/delta_writer/delta_writer.h" +#include "runtime/thread_context.h" +#include "runtime/workload_management/resource_context.h" #include "storage/binlog.h" #include "storage/data_dir.h" #include "storage/olap_common.h" @@ -478,6 +480,15 @@ Status TxnManager::commit_txn(OlapMeta* meta, TPartitionId partition_id, { std::lock_guard wrlock(_get_txn_map_lock(transaction_id)); auto load_info = std::make_shared(load_id, rowset_ptr); + // Publish workers attach their own memory context. Preserve the writer's + // resource domain here; recovered/contextless transactions use the default pool. + if (!is_recovery) { + SCOPED_INIT_THREAD_CONTEXT(); + auto* ctx = thread_context(); + if (ctx->is_attach_task()) { + load_info->workload_group = ctx->resource_ctx()->workload_group(); + } + } load_info->attach_row_binlog = attach_row_binlog; // resolve the independent binlog tablet in advance for the later publish phase. if (load_info->attach_row_binlog.rowset != nullptr && diff --git a/be/src/storage/txn/txn_manager.h b/be/src/storage/txn/txn_manager.h index b0195bb9e3586a..6d52ade470f60e 100644 --- a/be/src/storage/txn/txn_manager.h +++ b/be/src/storage/txn/txn_manager.h @@ -49,6 +49,7 @@ namespace doris { class OlapMeta; +class WorkloadGroup; struct TabletPublishStatistics; struct PartialUpdateInfo; @@ -79,6 +80,8 @@ struct TxnPublishInfo { struct TabletTxnInfo { PUniqueId load_id; + // Runtime-only write-stage owner, retained through local publish/retries. + std::shared_ptr workload_group; RowsetSharedPtr rowset; // The row-binlog tablet is attached while PREPARED; its rowset is filled at commit. RowBinlogTxnInfo attach_row_binlog; diff --git a/be/src/util/load_task_queue.h b/be/src/util/load_task_queue.h new file mode 100644 index 00000000000000..80395ad57f2fec --- /dev/null +++ b/be/src/util/load_task_queue.h @@ -0,0 +1,146 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace doris { + +// One-task round robin between ready loads, strict priority within each load. +// The caller serializes push/pop/erase/remove with the same lock. +// Only empty-to-nonempty transitions put a load in _ready_loads. There is no +// per-load concurrency cap. Push/pop are +// amortized O(1), as is erasing a task using its handle. +template +class LoadTaskQueue { +public: + static constexpr size_t NUM_PRIORITIES = 4; + + LoadTaskQueue() = default; + // Ready positions refer to this queue's list and must not be copied. + LoadTaskQueue(const LoadTaskQueue&) = delete; + LoadTaskQueue& operator=(const LoadTaskQueue&) = delete; + + // Valid until this entry is popped or erased. Other entries retain their handles. + struct Handle { + int64_t load_id; + size_t priority; + typename std::list::iterator position; + }; + + Handle push(int64_t load_id, size_t priority, T task) { + assert(priority < NUM_PRIORITIES); + auto [it, inserted] = _loads.try_emplace(load_id); + if (inserted) { + it->second.ready_position = _ready_loads.insert(_ready_loads.end(), load_id); + } + auto& queue = it->second.queues[priority]; + auto position = queue.insert(queue.end(), std::move(task)); + ++_size; + return {load_id, priority, position}; + } + + T pop() { + assert(!empty()); + auto load_id = _ready_loads.front(); + auto it = _loads.find(load_id); + auto& queues = it->second.queues; + size_t p = 0; + while (queues[p].empty()) { + ++p; + } + T task = std::move(queues[p].front()); + queues[p].pop_front(); + --_size; + if (queues_empty(queues)) { + _ready_loads.pop_front(); + _loads.erase(it); + } else { + // Rotate before the caller executes the task, without reallocating + // the ready entry or waiting for a running task of this load. + _ready_loads.splice(_ready_loads.end(), _ready_loads, _ready_loads.begin()); + } + return task; + } + + // Remove an already selected task without consuming another load's turn. + void erase(const Handle& handle) { + auto it = _loads.find(handle.load_id); + assert(it != _loads.end()); + it->second.queues[handle.priority].erase(handle.position); + --_size; + if (queues_empty(it->second.queues)) { + _ready_loads.erase(it->second.ready_position); + _loads.erase(it); + } + } + + // Return removed tasks so owners can destroy callbacks outside their lock. + template + std::vector remove_if(int64_t load_id, Predicate predicate) { + std::vector removed; + auto it = _loads.find(load_id); + if (it == _loads.end()) { + return removed; // The token may have only running tasks. + } + for (auto& queue : it->second.queues) { + for (auto entry = queue.begin(); entry != queue.end();) { + if (predicate(*entry)) { + removed.push_back(std::move(*entry)); + entry = queue.erase(entry); + --_size; + } else { + ++entry; + } + } + } + if (queues_empty(it->second.queues)) { + _ready_loads.erase(it->second.ready_position); + _loads.erase(it); + } + return removed; + } + + bool empty() const { return _size == 0; } + size_t size() const { return _size; } + +private: + using Queues = std::array, NUM_PRIORITIES>; + static bool queues_empty(const Queues& queues) { + return std::all_of(queues.begin(), queues.end(), [](const auto& q) { return q.empty(); }); + } + + struct Load { + Queues queues; + std::list::iterator ready_position; + }; + + std::unordered_map _loads; + std::list _ready_loads; + size_t _size = 0; +}; + +} // namespace doris diff --git a/be/src/util/threadpool.cpp b/be/src/util/threadpool.cpp index 194b80a8423402..a84aa74d72d26d 100644 --- a/be/src/util/threadpool.cpp +++ b/be/src/util/threadpool.cpp @@ -23,16 +23,20 @@ #include #include #include +#include #include #include #include #include "absl/strings/substitute.h" +#include "common/check.h" #include "common/exception.h" #include "common/logging.h" #include "common/metrics/doris_metrics.h" #include "common/metrics/metrics.h" #include "util/debug_points.h" +#include "util/defer_op.h" +#include "util/load_task_queue.h" #include "util/stopwatch.hpp" #include "util/thread.h" @@ -63,6 +67,36 @@ class FunctionRunnable : public Runnable { std::function _func; }; +struct ThreadPool::ScheduledLoadTask { + ThreadPoolToken* token; + Task task; + std::list::iterator token_position; + LoadTaskQueue::Handle queue_position; +}; + +class ThreadPool::LoadQueue : public LoadTaskQueue {}; + +class ThreadPoolToken::LoadEntries { +public: + std::list tasks; +}; + +static thread_local ThreadPoolToken* executing_load_token = nullptr; + +static thread_local bool helping_load_task = false; + +ThreadPool* ThreadPool::current_load_pool() { + return executing_load_token == nullptr ? nullptr : executing_load_token->_pool; +} + +bool ThreadPool::is_helping_load_task() { + return helping_load_task; +} + +bool ThreadPool::queues_empty() const { + return _queue.empty() && _load_queue->empty(); +} + ThreadPoolBuilder::ThreadPoolBuilder(string name, string workload_group) : _name(std::move(name)), _workload_group(std::move(workload_group)), @@ -114,7 +148,7 @@ ThreadPoolToken::~ThreadPoolToken() { } Status ThreadPoolToken::submit(std::shared_ptr r) { - return _pool->do_submit(std::move(r), this); + return _pool->do_submit(std::move(r), this, _load_id, _load_priority); } Status ThreadPoolToken::submit_func(std::function f) { @@ -122,8 +156,28 @@ Status ThreadPoolToken::submit_func(std::function f) { } void ThreadPoolToken::shutdown() { + // Declare before the lock: callback destruction must run after unlocking. + std::list removed_load_tasks; std::unique_lock l(_pool->_lock); - _pool->check_not_pool_thread_unlocked(); + // Flush cleanup can release the last rowset-writer reference on a worker. + // Its bitmap jobs are leaves: after removing queued jobs below, only + // already-running independent leaves remain to join. Never join our own token. + const bool join_bitmap_leaves = _is_load_token && _is_leaf && executing_load_token != nullptr && + executing_load_token != this; + if (!_is_load_token || (_active_threads != 0 && !join_bitmap_leaves)) { + _pool->check_not_pool_thread_unlocked(); + } + if (_active_threads != 0 && join_bitmap_leaves && ThreadPool::current_load_pool() == _pool) { + DORIS_CHECK(!executing_load_token->_is_leaf) << "Load children must not join other tasks"; + } + if (_queued_load_tasks != 0) { + for (const auto& entry : _load_entries->tasks) { + _pool->_load_queue->erase(entry.queue_position); + } + removed_load_tasks.splice(removed_load_tasks.end(), _load_entries->tasks); + _pool->_total_queued_tasks -= _queued_load_tasks; + _queued_load_tasks = 0; + } // Clear the queue under the lock, but defer the releasing of the tasks // outside the lock, in case there are concurrent threads wanting to access @@ -147,11 +201,13 @@ void ThreadPoolToken::shutdown() { // Plus doing it this way (rather than switching to QUIESCING and waiting // for a worker thread to process the queue entry) helps retain state // transition symmetry with ThreadPool::shutdown. - for (auto it = _pool->_queue.begin(); it != _pool->_queue.end();) { - if (*it == this) { - it = _pool->_queue.erase(it); - } else { - it++; + if (!_is_load_token) { + for (auto it = _pool->_queue.begin(); it != _pool->_queue.end();) { + if (*it == this) { + it = _pool->_queue.erase(it); + } else { + it++; + } } } @@ -169,6 +225,10 @@ void ThreadPoolToken::shutdown() { default: break; } + if (_pool->_total_queued_tasks == 0 && _pool->_active_threads == 0) { + _pool->_idle_cond.notify_all(); + } + l.unlock(); } void ThreadPoolToken::wait() { @@ -177,6 +237,42 @@ void ThreadPoolToken::wait() { _not_running_cond.wait(l, [this]() { return !is_active(); }); } +Status ThreadPoolToken::wait_and_help() { + if (executing_load_token == nullptr || executing_load_token->_pool != _pool) { + return Status::InvalidArgument( + "wait_and_help requires a load worker from the token's pool"); + } + if (executing_load_token == this || executing_load_token->_is_leaf) { + return Status::InvalidArgument("wait_and_help requires a distinct non-leaf parent token"); + } + if (!_is_load_token || !_is_leaf) { + return Status::InvalidArgument("wait_and_help requires a load leaf token"); + } + std::unique_lock l(_pool->_lock); + while (is_active()) { + if (_queued_load_tasks == 0) { + ++_waiting_helpers; + Defer stop_waiting = [&] { --_waiting_helpers; }; + _not_running_cond.wait(l, [this] { return !is_active() || _queued_load_tasks != 0; }); + continue; + } + auto* entry = &_load_entries->tasks.front(); + _pool->_load_queue->erase(entry->queue_position); + auto task = _pool->take_load_task_unlocked(entry); + ++_active_threads; + // The parent already counts as a busy physical worker in the pool. + l.unlock(); + Defer finish = [&] { + // Also release captures and retire the task if a callback throws. + task.runnable.reset(); + l.lock(); + _pool->finish_task_unlocked(this); + }; + _pool->run_task(this, task, true); + } + return Status::OK(); +} + void ThreadPoolToken::transition(State new_state) { #ifndef NDEBUG CHECK_NE(_state, new_state); @@ -185,16 +281,16 @@ void ThreadPoolToken::transition(State new_state) { case State::IDLE: CHECK(new_state == State::RUNNING || new_state == State::QUIESCED); if (new_state == State::RUNNING) { - CHECK(!_entries.empty()); + CHECK(!tasks_empty()); } else { - CHECK(_entries.empty()); + CHECK(tasks_empty()); CHECK_EQ(_active_threads, 0); } break; case State::RUNNING: CHECK(new_state == State::IDLE || new_state == State::QUIESCING || new_state == State::QUIESCED); - CHECK(_entries.empty()); + CHECK(tasks_empty()); if (new_state == State::QUIESCING) { CHECK_GT(_active_threads, 0); } @@ -249,7 +345,8 @@ bool ThreadPoolToken::need_dispatch() { } ThreadPool::ThreadPool(const ThreadPoolBuilder& builder) - : _name(builder._name), + : _load_queue(std::make_unique()), + _name(builder._name), _workload_group(builder._workload_group), _min_threads(builder._min_threads), _max_threads(builder._max_threads), @@ -262,11 +359,12 @@ ThreadPool::ThreadPool(const ThreadPoolBuilder& builder) _total_queued_tasks(0), _cgroup_cpu_ctl(builder._cgroup_cpu_ctl), _tokenless(new_token(ExecutionMode::CONCURRENT)), + _load_tokenless(new_load_token(0, LoadTaskPriority::LOW, LoadTaskType::PARENT)), _id(UniqueId::gen_uid()) {} ThreadPool::~ThreadPool() { - // There should only be one live token: the one used in tokenless submission. - CHECK_EQ(1, _tokens.size()) << absl::Substitute( + // Only the two internal tokens used by tokenless submission may remain. + CHECK_EQ(2, _tokens.size()) << absl::Substitute( "Threadpool $0 destroyed with $1 allocated tokens", _name, _tokens.size()); shutdown(); VLOG_DEBUG << fmt::format("Thread pool {} destroyed", _name); @@ -360,9 +458,15 @@ void ThreadPool::shutdown() { // wanting to access the ThreadPool. The task's destructors may acquire // locks, etc, so this also prevents lock inversions. _queue.clear(); + std::list load_tasks_to_release; + _load_queue = std::make_unique(); std::deque> to_release; for (auto* t : _tokens) { + if (t->_queued_load_tasks != 0) { + load_tasks_to_release.splice(load_tasks_to_release.end(), t->_load_entries->tasks); + t->_queued_load_tasks = 0; + } if (!t->_entries.empty()) { to_release.emplace_back(std::move(t->_entries)); } @@ -400,6 +504,7 @@ void ThreadPool::shutdown() { DCHECK(t->state() == ThreadPoolToken::State::IDLE || t->state() == ThreadPoolToken::State::QUIESCED); } + l.unlock(); } std::unique_ptr ThreadPool::new_token(ExecutionMode mode, int max_concurrency) { @@ -411,6 +516,27 @@ std::unique_ptr ThreadPool::new_token(ExecutionMode mode, int m return t; } +std::unique_ptr ThreadPool::new_load_token(int64_t load_id, + LoadTaskPriority priority, + LoadTaskType type) { + DORIS_CHECK(current_load_pool() != this || !executing_load_token->_is_leaf) + << "Load children must be leaf tasks"; + std::lock_guard l(_lock); + std::unique_ptr token(new ThreadPoolToken(this, ExecutionMode::CONCURRENT)); + token->_load_entries = std::make_unique(); + token->_is_leaf = type == LoadTaskType::LEAF; + token->_is_load_token = true; + token->_load_id = load_id; + token->_load_priority = priority; + CHECK(_tokens.insert(token.get()).second); + return token; +} + +Status ThreadPool::submit_load(std::shared_ptr r, int64_t load_id, + LoadTaskPriority priority) { + return do_submit(std::move(r), _load_tokenless.get(), load_id, priority); +} + void ThreadPool::release_token(ThreadPoolToken* t) { std::lock_guard l(_lock); CHECK(!t->is_active()) << absl::Substitute("Token with state $0 may not be released", @@ -426,7 +552,8 @@ Status ThreadPool::submit_func(std::function f) { return submit(std::make_shared(std::move(f))); } -Status ThreadPool::do_submit(std::shared_ptr r, ThreadPoolToken* token) { +Status ThreadPool::do_submit(std::shared_ptr r, ThreadPoolToken* token, int64_t load_id, + LoadTaskPriority priority) { DCHECK(token); std::unique_lock l(_lock); @@ -468,8 +595,8 @@ Status ThreadPool::do_submit(std::shared_ptr r, ThreadPoolToken* token int threads_from_this_submit = token->is_active() && token->mode() == ExecutionMode::SERIAL ? 0 : 1; int inactive_threads = _num_threads + _num_threads_pending_start - _active_threads; - int additional_threads = - static_cast(_queue.size()) + threads_from_this_submit - inactive_threads; + int additional_threads = static_cast(_queue.size() + _load_queue->size()) + + threads_from_this_submit - inactive_threads; bool need_a_thread = false; if (additional_threads > 0 && _num_threads + _num_threads_pending_start < _max_threads) { need_a_thread = true; @@ -483,23 +610,42 @@ Status ThreadPool::do_submit(std::shared_ptr r, ThreadPoolToken* token // Add the task to the token's queue. ThreadPoolToken::State state = token->state(); DCHECK(state == ThreadPoolToken::State::IDLE || state == ThreadPoolToken::State::RUNNING); - token->_entries.emplace_back(std::move(task)); - // When we need to execute the task in the token, we submit the token object to the queue. - // There are currently two places where tokens will be submitted to the queue: - // 1. When submitting a new task, if the token is still in the IDLE state, - // or the concurrency of the token has not reached the online level, it will be added to the queue. - // 2. When the dispatch thread finishes executing a task: - // 1. If it is a SERIAL token, and there are unsubmitted tasks, submit them to the queue. - // 2. If it is a CONCURRENT token, and there are still unsubmitted tasks, and the upper limit of concurrency is not reached, - // then submitted to the queue. - if (token->need_dispatch()) { - _queue.emplace_back(token); - ++token->_num_submitted_tasks; + if (token->_is_load_token) { + auto& entries = token->_load_entries->tasks; + auto position = entries.emplace(entries.end()); + position->token = token; + position->task = std::move(task); + position->token_position = position; + position->queue_position = + _load_queue->push(load_id, static_cast(priority), &*position); + ++token->_queued_load_tasks; + // Only a sleeping helper needs an enqueue notification. Completion + // waiters are still all notified on transitions to IDLE or QUIESCED. + if (token->_waiting_helpers != 0) { + token->_not_running_cond.notify_all(); + } if (state == ThreadPoolToken::State::IDLE) { token->transition(ThreadPoolToken::State::RUNNING); } } else { - ++token->_num_unsubmitted_tasks; + token->_entries.emplace_back(std::move(task)); + // When we need to execute the task in the token, we submit the token object to the queue. + // There are currently two places where tokens will be submitted to the queue: + // 1. When submitting a new task, if the token is still in the IDLE state, + // or the concurrency of the token has not reached the online level, it will be added to the queue. + // 2. When the dispatch thread finishes executing a task: + // 1. If it is a SERIAL token, and there are unsubmitted tasks, submit them to the queue. + // 2. If it is a CONCURRENT token, and there are still unsubmitted tasks, and the upper limit of concurrency is not reached, + // then submitted to the queue. + if (token->need_dispatch()) { + _queue.emplace_back(token); + ++token->_num_submitted_tasks; + if (state == ThreadPoolToken::State::IDLE) { + token->transition(ThreadPoolToken::State::RUNNING); + } + } else { + ++token->_num_unsubmitted_tasks; + } } _total_queued_tasks++; @@ -540,6 +686,52 @@ void ThreadPool::wait() { _idle_cond.wait(l, [this]() { return _total_queued_tasks == 0 && _active_threads == 0; }); } +ThreadPool::Task ThreadPool::take_load_task_unlocked(ScheduledLoadTask* entry) { + auto* token = entry->token; + Task task = std::move(entry->task); + token->_load_entries->tasks.erase(entry->token_position); + --token->_queued_load_tasks; + --_total_queued_tasks; + return task; +} + +void ThreadPool::run_task(ThreadPoolToken* token, Task& task, bool helping) { + thread_pool_task_wait_worker_time_ns_total->increment(task.submit_time_wather.elapsed_time()); + thread_pool_task_wait_worker_count_total->increment(1); + MonotonicStopWatch execution_time; + execution_time.start(); + auto* previous_token = executing_load_token; + const bool previous_helping = helping_load_task; + executing_load_token = token->_is_load_token ? token : nullptr; + helping_load_task = helping; + Defer restore = [&] { + executing_load_token = previous_token; + helping_load_task = previous_helping; + }; + task.runnable->run(); + // Release captures outside the pool lock, while still identifying this load worker. + task.runnable.reset(); + thread_pool_task_execution_time_ns_total->increment(execution_time.elapsed_time()); + thread_pool_task_execution_count_total->increment(1); +} + +void ThreadPool::finish_task_unlocked(ThreadPoolToken* token) { + auto state = token->state(); + DCHECK(state == ThreadPoolToken::State::RUNNING || state == ThreadPoolToken::State::QUIESCING); + --token->_active_threads; + if (!token->_is_load_token) { + --token->_num_submitted_tasks; + } + if (token->_active_threads == 0) { + if (state == ThreadPoolToken::State::QUIESCING) { + DCHECK(token->tasks_empty()); + token->transition(ThreadPoolToken::State::QUIESCED); + } else if (token->tasks_empty()) { + token->transition(ThreadPoolToken::State::IDLE); + } + } +} + void ThreadPool::dispatch_thread() { std::unique_lock l(_lock); if (!_threads.insert(Thread::current_thread()).second) { @@ -567,7 +759,7 @@ void ThreadPool::dispatch_thread() { break; } - if (_queue.empty()) { + if (queues_empty()) { // There's no work to do, let's go idle. // // Note: if FIFO behavior is desired, it's as simple as changing this to push_back(). @@ -587,7 +779,7 @@ void ThreadPool::dispatch_thread() { // brief period during which another thread may actually grab the internal mutex // protecting the state, signal, and release again before we get the mutex. So, // we'll recheck the empty queue case regardless. - if (_queue.empty() && _num_threads + _num_threads_pending_start > _min_threads) { + if (queues_empty() && _num_threads + _num_threads_pending_start > _min_threads) { VLOG_NOTICE << "Releasing worker thread from pool " << _name << " after " << std::chrono::duration_cast( _idle_timeout) @@ -599,56 +791,31 @@ void ThreadPool::dispatch_thread() { continue; } - MonotonicStopWatch task_execution_time_watch; - task_execution_time_watch.start(); // Get the next token and task to execute. - ThreadPoolToken* token = _queue.front(); - _queue.pop_front(); + ThreadPoolToken* token; + Task task; + if (_queue.empty()) { + auto* entry = _load_queue->pop(); + token = entry->token; + task = take_load_task_unlocked(entry); + } else { + token = _queue.front(); + _queue.pop_front(); + task = std::move(token->_entries.front()); + token->_entries.pop_front(); + } DCHECK_EQ(ThreadPoolToken::State::RUNNING, token->state()); - DCHECK(!token->_entries.empty()); - Task task = std::move(token->_entries.front()); - thread_pool_task_wait_worker_time_ns_total->increment( - task.submit_time_wather.elapsed_time()); - thread_pool_task_wait_worker_count_total->increment(1); - token->_entries.pop_front(); token->_active_threads++; - --_total_queued_tasks; + if (!token->_is_load_token) { + --_total_queued_tasks; + } ++_active_threads; l.unlock(); - - // Execute the task - task.runnable->run(); - // Destruct the task while we do not hold the lock. - // - // The task's destructor may be expensive if it has a lot of bound - // objects, and we don't want to block submission of the threadpool. - // In the worst case, the destructor might even try to do something - // with this threadpool, and produce a deadlock. - task.runnable.reset(); + run_task(token, task); l.lock(); - thread_pool_task_execution_time_ns_total->increment( - task_execution_time_watch.elapsed_time()); - thread_pool_task_execution_count_total->increment(1); - // Possible states: - // 1. The token was shut down while we ran its task. Transition to QUIESCED. - // 2. The token has no more queued tasks. Transition back to IDLE. - // 3. The token has more tasks. Requeue it and transition back to RUNNABLE. + finish_task_unlocked(token); ThreadPoolToken::State state = token->state(); - DCHECK(state == ThreadPoolToken::State::RUNNING || - state == ThreadPoolToken::State::QUIESCING); - --token->_active_threads; - --token->_num_submitted_tasks; - - // handle shutdown && idle - if (token->_active_threads == 0) { - if (state == ThreadPoolToken::State::QUIESCING) { - DCHECK(token->_entries.empty()); - token->transition(ThreadPoolToken::State::QUIESCED); - } else if (token->_entries.empty()) { - token->transition(ThreadPoolToken::State::IDLE); - } - } // We decrease _num_submitted_tasks holding lock, so the following DCHECK works. DCHECK(token->_num_submitted_tasks < token->_max_concurrency); @@ -681,7 +848,7 @@ void ThreadPool::dispatch_thread() { // Sanity check: if we're the last thread exiting, the queue ought to be // empty. Otherwise it will never get processed. - CHECK(_queue.empty()); + CHECK(queues_empty()); DCHECK_EQ(0, _total_queued_tasks); } } diff --git a/be/src/util/threadpool.h b/be/src/util/threadpool.h index ab97322a7bbd9b..166bb0be0fde22 100644 --- a/be/src/util/threadpool.h +++ b/be/src/util/threadpool.h @@ -50,6 +50,20 @@ class Thread; class ThreadPool; class ThreadPoolToken; +// Priority within a load. Callers map task stages to levels; lower values run first. +enum class LoadTaskPriority : uint8_t { + HIGHEST = 0, + HIGH = 1, + MID = 2, + LOW = 3, +}; + +// Dependency role is independent of scheduling priority. +enum class LoadTaskType : uint8_t { + PARENT, + LEAF, +}; + class Runnable { public: virtual void run() = 0; @@ -200,6 +214,17 @@ class ThreadPool { // Submits a function bound using std::bind(&FuncName, args...). Status submit_func(std::function f); + // Take one task per transaction turn on this pool (resource domain). Existing tokenless + // and SERIAL/CONCURRENT token submissions retain their original policy. + Status submit_load(std::shared_ptr r, int64_t load_id, LoadTaskPriority priority); + // Leaf tokens must never wait for other work in this pool. + std::unique_ptr new_load_token(int64_t load_id, LoadTaskPriority priority, + LoadTaskType type); + // Null outside a load callback; helping children remain in the parent pool. + static ThreadPool* current_load_pool(); + // True only while a parent runs a queued child via wait_and_help(). + static bool is_helping_load_task(); + // Waits until all the tasks are completed. void wait(); @@ -312,7 +337,15 @@ class ThreadPool { void check_not_pool_thread_unlocked(); // Submits a task to be run via token. - Status do_submit(std::shared_ptr r, ThreadPoolToken* token); + Status do_submit(std::shared_ptr r, ThreadPoolToken* token, int64_t load_id = 0, + LoadTaskPriority priority = LoadTaskPriority::LOW); + bool queues_empty() const; + struct ScheduledLoadTask; + Task take_load_task_unlocked(ScheduledLoadTask* entry); + void run_task(ThreadPoolToken* token, Task& task, bool helping = false); + void finish_task_unlocked(ThreadPoolToken* token); + class LoadQueue; + std::unique_ptr _load_queue; // Releases token 't' and invalidates it. void release_token(ThreadPoolToken* t); @@ -406,6 +439,7 @@ class ThreadPool { // ExecutionMode::CONCURRENT token used by the pool for tokenless submission. std::unique_ptr _tokenless; + std::unique_ptr _load_tokenless; const UniqueId _id; std::shared_ptr _metric_entity; @@ -448,6 +482,12 @@ class ThreadPoolToken { // Waits until all the tasks submitted via this token are completed. void wait(); + // Only a non-leaf load task of this pool may help a distinct leaf token. + // Runs this token's queued tasks on the caller, then joins running leaves. + // Invalid callers receive an error without changing this token. Its owner + // must still drain or cancel it before releasing resources used by tasks. + Status wait_and_help(); + // Waits for all submissions using this token are complete, or until 'delta' // time elapses. // @@ -463,7 +503,7 @@ class ThreadPoolToken { size_t num_tasks() { std::lock_guard l(_pool->_lock); - return _entries.size(); + return _entries.size() + _queued_load_tasks; } ThreadPoolToken(const ThreadPoolToken&) = delete; @@ -537,6 +577,19 @@ class ThreadPoolToken { // Queued client tasks. std::deque _entries; + // Immutable scheduling identity; writer/tablet tokens of one transaction + // share an outer FIFO entry while retaining independent wait/shutdown. + // Load tasks are owned here; the scheduler holds removable references only. + class LoadEntries; + std::unique_ptr _load_entries; + bool _is_load_token = false; + bool _is_leaf = false; + int64_t _load_id = 0; + LoadTaskPriority _load_priority = LoadTaskPriority::LOW; + size_t _queued_load_tasks = 0; + size_t _waiting_helpers = 0; // Protected by the pool lock; only counts CV waits. + bool tasks_empty() const { return _entries.empty() && _queued_load_tasks == 0; } + // Condition variable for "token is idle". Waiters wake up when the token // transitions to IDLE or QUIESCED. std::condition_variable _not_running_cond; diff --git a/be/test/cloud/cloud_engine_calc_delete_bitmap_task_test.cpp b/be/test/cloud/cloud_engine_calc_delete_bitmap_task_test.cpp new file mode 100644 index 00000000000000..dd3caea4ebf78c --- /dev/null +++ b/be/test/cloud/cloud_engine_calc_delete_bitmap_task_test.cpp @@ -0,0 +1,132 @@ +// 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 "cloud/cloud_engine_calc_delete_bitmap_task.h" + +#include + +#include +#include +#include + +#include "runtime/thread_context.h" +#include "runtime/workload_group/workload_group.h" +#include "runtime/workload_management/resource_context.h" +#include "storage/delete/calc_delete_bitmap_executor.h" +#include "util/threadpool.h" + +namespace doris { + +TEST(CloudEngineCalcDeleteBitmapTaskTest, SubmissionFailureKeepsOriginalError) { + std::unique_ptr pool; + ASSERT_TRUE( + ThreadPoolBuilder("cloud_bitmap_rejected_test").set_max_threads(1).build(&pool).ok()); + CloudStorageEngine engine {EngineOptions()}; + engine._txn_delete_bitmap_cache = std::make_unique(1024 * 1024); + ASSERT_TRUE(engine._txn_delete_bitmap_cache->init().ok()); + engine._calc_delete_bitmap_executor = std::make_unique(); + engine._calc_delete_bitmap_executor->init("cloud_bitmap_background_test", 1, pool.get()); + SCOPED_ATTACH_TASK(MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "cloud_bitmap_rejected_test")); + pool->shutdown(); + + TCalcDeleteBitmapRequest request; + request.transaction_id = 1; + TCalcDeleteBitmapPartitionInfo partition; + partition.partition_id = 1; + partition.version = 2; + partition.tablet_ids = {10}; + request.partitions = {partition}; + std::vector errors, successes; + CloudEngineCalcDeleteBitmapTask task(engine, request, &errors, &successes); + auto st = task.execute(); + EXPECT_TRUE(st.is()); + EXPECT_NE(st.to_string().find("cloud_bitmap_rejected_test"), std::string::npos); + EXPECT_NE(st.to_string().find("shut down"), std::string::npos); + EXPECT_EQ(errors, (std::vector {10})); + EXPECT_TRUE(successes.empty()); + + errors.clear(); + CloudEngineCalcDeleteBitmapTask failed_task(engine, request, &errors, &successes); + auto original = Status::InternalError("earlier tablet bitmap failure"); + failed_task.add_error_tablet_id(9, original); + // Even when wait() fails, retain the tablet error already selected by the + // cloud task instead of replacing it with the later submission failure. + EXPECT_EQ(failed_task.execute().to_string(), original.to_string()); + EXPECT_EQ(errors, (std::vector {9, 10})); + EXPECT_TRUE(successes.empty()); +} + +TEST(CloudEngineCalcDeleteBitmapTaskTest, EmptyRowsetPublishUsesCachedWorkloadGroup) { + std::unique_ptr default_pool; + ASSERT_TRUE(ThreadPoolBuilder("empty_publish_default") + .set_max_threads(1) + .build(&default_pool) + .ok()); + auto wg = std::make_shared( + WorkloadGroupInfo {.id = 68388, .name = "empty_publish_wg"}); + ASSERT_TRUE(ThreadPoolBuilder("empty_publish_wg") + .set_max_threads(1) + .build(&wg->_memtable_flush_pool) + .ok()); + CloudStorageEngine engine {EngineOptions()}; + engine._txn_delete_bitmap_cache = std::make_unique(1024 * 1024); + ASSERT_TRUE(engine._txn_delete_bitmap_cache->init().ok()); + engine._calc_delete_bitmap_executor = std::make_unique(); + engine._calc_delete_bitmap_executor->init("empty_publish_background", 1, default_pool.get()); + auto ctx = ResourceContext::create_shared(); + ctx->memory_context()->set_mem_tracker(MemTrackerLimiter::create_shared( + MemTrackerLimiter::Type::OTHER, "empty_publish_writer")); + ctx->set_workload_group(wg); + { + SCOPED_ATTACH_TASK(ctx); + engine.txn_delete_bitmap_cache().mark_empty_rowset(1, 10, INT64_MAX); + engine.txn_delete_bitmap_cache().mark_empty_rowset(3, 10, INT64_MAX); + } + // The publish request itself has no workload group. The first empty + // subtransaction has no owner, so routing must continue to the second one. + SCOPED_ATTACH_TASK(MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "empty_publish_request")); + engine.txn_delete_bitmap_cache().mark_empty_rowset(2, 10, INT64_MAX); + default_pool->shutdown(); + wg->_memtable_flush_pool->shutdown(); + // Distinct shutdown errors identify the pool selected by the real execute + // path without requiring tablet metadata or meta-service RPCs. + for (bool subtransactions : {false, true}) { + TCalcDeleteBitmapRequest request; + request.transaction_id = subtransactions ? 4 : 1; + TCalcDeleteBitmapPartitionInfo partition; + partition.partition_id = 1; + partition.version = 2; + partition.tablet_ids = {10}; + if (subtransactions) { + partition.__set_sub_txn_ids({2, 3}); + } + request.partitions = {partition}; + for (int retry = 0; retry < 2; ++retry) { + std::vector errors, successes; + CloudEngineCalcDeleteBitmapTask task(engine, request, &errors, &successes); + auto st = task.execute(); + EXPECT_TRUE(st.is()); + EXPECT_NE(st.to_string().find("empty_publish_wg"), std::string::npos); + EXPECT_EQ(errors, (std::vector {10})); + EXPECT_TRUE(successes.empty()); + } + } +} + +} // namespace doris diff --git a/be/test/cloud/cloud_tablets_channel_context_test.cpp b/be/test/cloud/cloud_tablets_channel_context_test.cpp new file mode 100644 index 00000000000000..15b0c3b47a7855 --- /dev/null +++ b/be/test/cloud/cloud_tablets_channel_context_test.cpp @@ -0,0 +1,185 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include + +#include "cloud/cloud_delta_writer.h" +#include "cloud/cloud_rowset_builder.h" +#include "cloud/cloud_storage_engine.h" +#include "cloud/cloud_tablets_channel.h" +#include "common/signal_handler.h" +#include "cpp/sync_point.h" +#include "runtime/thread_context.h" +#include "runtime/workload_group/workload_group.h" +#include "runtime/workload_management/resource_context.h" +#include "storage/delete/calc_delete_bitmap_executor.h" +#include "util/threadpool.h" + +namespace doris { +namespace { + +// Exercise close -> commit_rowset -> _commit_empty_rowset without metadata RPCs. +// Stop after creating and exercising the same bitmap token as the real builder. +class EmptyRowsetContextProbe : public CloudRowsetBuilder { +public: + using CloudRowsetBuilder::CloudRowsetBuilder; + Status init() override { return on_init(); } + std::function on_init; +}; + +} // namespace + +class CloudTabletsChannelContextTest : public testing::TestWithParam> {}; + +TEST_P(CloudTabletsChannelContextTest, EmptyRowsetCommitInheritsLoadContext) { + const auto [force_inline, bthread_caller] = GetParam(); + auto* sync_point = SyncPoint::get_instance(); + SyncPoint::CallbackGuard guard; + if (force_inline) { + sync_point->set_call_back( + "bthread_fork_join::start_background", + [](auto&& args) { + auto* result = try_any_cast*>(args.back()); + result->first = EAGAIN; + result->second = true; + }, + &guard); + sync_point->enable_processing(); + } + + int initialized = 0; + std::function run = [&] { + SCOPED_INIT_THREAD_CONTEXT(); + EXPECT_FALSE(thread_context()->is_attach_task()); + auto ctx = ResourceContext::create_shared(); + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::LOAD, + "empty_rowset_context_test"); + ctx->memory_context()->set_mem_tracker(tracker); + TUniqueId task_id; + task_id.hi = 68385; + task_id.lo = 67674; + ctx->task_controller()->set_task_id(task_id); + auto wg = std::make_shared( + WorkloadGroupInfo {.id = 68385, .name = "empty_rowset_context_test"}); + ASSERT_TRUE(ThreadPoolBuilder("empty_rowset_workload_group_test") + .set_max_threads(1) + .build(&wg->_memtable_flush_pool) + .ok()); + ctx->set_workload_group(wg); + { + SCOPED_ATTACH_TASK(ctx); + CloudStorageEngine engine {EngineOptions {}}; + engine.init_calc_delete_bitmap_executor_for_UT(); + WriteRequest req; + req.txn_id = 123; + req.tablet_id = 10; + req.partition_id = 20; + req.index_id = 30; + PUniqueId load_id; + load_id.set_hi(task_id.hi); + load_id.set_lo(task_id.lo); + CloudTabletsChannel channel(engine, TabletsChannelKey(load_id, req.index_id), + UniqueId(load_id), false, nullptr); + channel._num_remaining_senders = 1; + auto writer = + std::make_unique(engine, req, nullptr, UniqueId(load_id)); + ASSERT_FALSE(writer->is_init()); + auto builder = std::make_unique(engine, req, nullptr); + const auto caller_id = bthread_self(); + const auto expected_error = Status::InternalError("stop after empty rowset token init"); + builder->on_init = [&]() -> Status { + ++initialized; + EXPECT_EQ(bthread_self() == caller_id, force_inline); + // Fail cleanly on the unfixed code even when a fresh bthread has no TLS. + SCOPED_INIT_THREAD_CONTEXT(); + EXPECT_TRUE(thread_context()->is_attach_task()); + if (!thread_context()->is_attach_task()) { + return Status::InternalError("empty rowset commit lost its load context"); + } + EXPECT_EQ(thread_context()->resource_ctx(), ctx); + EXPECT_EQ(thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker(), + tracker.get()); + EXPECT_EQ(signal::query_id_hi, task_id.hi); + EXPECT_EQ(signal::query_id_lo, task_id.lo); + auto token = engine.calc_delete_bitmap_executor()->create_load_token( + req.txn_id, LoadTaskPriority::HIGH, LoadTaskType::LEAF); + EXPECT_EQ(token->_thread_token->_pool, wg->get_memtable_flush_pool()); + EXPECT_EQ(token->_thread_token->_load_id, req.txn_id); + RETURN_IF_ERROR(token->submit_func([&] { + EXPECT_EQ(thread_context()->resource_ctx(), ctx); + EXPECT_EQ(thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker(), + tracker.get()); + return Status::OK(); + })); + RETURN_IF_ERROR(token->wait()); + return expected_error; + }; + writer->_rowset_builder = std::move(builder); + channel._tablet_writers.emplace(req.tablet_id, std::move(writer)); + ++BaseTabletsChannel::_s_tablet_writer_count; + + PTabletWriterAddBlockRequest close_req; + close_req.set_sender_id(0); + close_req.add_partition_ids(req.partition_id); + PTabletWriterAddBlockResult result; + bool finished = false; + EXPECT_EQ(channel.close(nullptr, close_req, &result, &finished), expected_error); + EXPECT_TRUE(finished); + // The synchronous fallback must not detach or replace the parent's context. + EXPECT_EQ(thread_context()->resource_ctx(), ctx); + EXPECT_EQ(thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker(), + tracker.get()); + EXPECT_EQ(signal::query_id_hi, task_id.hi); + EXPECT_EQ(signal::query_id_lo, task_id.lo); + } + EXPECT_FALSE(thread_context()->is_attach_task()); + }; + if (bthread_caller) { + bthread_t id; + const int start_result = bthread_start_background( + &id, nullptr, + [](void* arg) -> void* { + (*static_cast*>(arg))(); + return nullptr; + }, + &run); + EXPECT_EQ(start_result, 0); + if (start_result == 0) { + EXPECT_EQ(bthread_join(id, nullptr), 0); + } + } else { + std::thread caller(run); + caller.join(); + } + if (force_inline) { + sync_point->disable_processing(); + } + EXPECT_EQ(initialized, 1); +} + +INSTANTIATE_TEST_SUITE_P(AsyncAndFallback, CloudTabletsChannelContextTest, + testing::Combine(testing::Bool(), testing::Bool())); + +} // namespace doris diff --git a/be/test/cloud/cloud_txn_delete_bitmap_cache_test.cpp b/be/test/cloud/cloud_txn_delete_bitmap_cache_test.cpp new file mode 100644 index 00000000000000..66c31e07884581 --- /dev/null +++ b/be/test/cloud/cloud_txn_delete_bitmap_cache_test.cpp @@ -0,0 +1,144 @@ +// 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 "cloud/cloud_txn_delete_bitmap_cache.h" + +#include + +#include +#include + +#include "runtime/thread_context.h" +#include "runtime/workload_group/workload_group.h" +#include "runtime/workload_management/resource_context.h" + +namespace doris { + +TEST(CloudTxnDeleteBitmapCacheTest, ContextlessDeleteUsesDefaultWorkloadGroup) { + CloudTxnDeleteBitmapCache cache(1024 * 1024); + ASSERT_TRUE(cache.init().ok()); + auto wg = std::make_shared( + WorkloadGroupInfo {.id = 68385, .name = "bitmap_cache_test"}); + auto ctx = ResourceContext::create_shared(); + ctx->memory_context()->set_mem_tracker( + MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, "bitmap_cache_test")); + ctx->set_workload_group(wg); + // Agent DELETE workers have no attached task context. A fresh thread also + // verifies that a prior load's workload group is not retained after detach. + std::thread worker([&] { + SCOPED_INIT_THREAD_CONTEXT(); + EXPECT_FALSE(thread_context()->is_attach_task()); + { + SCOPED_ATTACH_TASK(ctx); + cache.set_tablet_txn_info(1, 10, std::make_shared(10), {}, nullptr, 0, + nullptr); + EXPECT_EQ(cache.get_workload_group(1, 10), wg); + } + EXPECT_FALSE(thread_context()->is_attach_task()); + cache.set_tablet_txn_info(2, 10, std::make_shared(10), {}, nullptr, 0, + nullptr); + EXPECT_EQ(cache.get_workload_group(2, 10), nullptr); + EXPECT_EQ(cache.get_workload_group(1, 10), wg); + // Replacing an entry from a contextless caller clears its old owner too. + cache.set_tablet_txn_info(1, 10, std::make_shared(10), {}, nullptr, 0, + nullptr); + EXPECT_EQ(cache.get_workload_group(1, 10), nullptr); + EXPECT_FALSE(thread_context()->is_attach_task()); + }); + worker.join(); +} + +TEST(CloudTxnDeleteBitmapCacheTest, EmptyRowsetOwnerIsRetainedUntilCleanup) { + CloudTxnDeleteBitmapCache cache(1024 * 1024); + ASSERT_TRUE(cache.init().ok()); + auto wg = std::make_shared( + WorkloadGroupInfo {.id = 68386, .name = "empty_bitmap_cache_test"}); + std::weak_ptr weak_wg = wg; + auto ctx = ResourceContext::create_shared(); + ctx->memory_context()->set_mem_tracker(MemTrackerLimiter::create_shared( + MemTrackerLimiter::Type::OTHER, "empty_bitmap_cache_test")); + ctx->set_workload_group(wg); + { + SCOPED_ATTACH_TASK(ctx); + cache.mark_empty_rowset(1, 10, INT64_MAX); + cache.mark_empty_rowset(2, 10, INT64_MAX); + } + EXPECT_TRUE(cache.is_empty_rowset(1, 10)); + EXPECT_EQ(cache.get_workload_group(1, 10), wg); + EXPECT_EQ(cache.get_workload_group(2, 10), wg); + EXPECT_EQ(cache.get_workload_group(3, 10), nullptr); + auto rowset_and_bitmap = cache.get_rowset_and_delete_bitmap(1, 10); + ASSERT_TRUE(rowset_and_bitmap.has_value()); + EXPECT_EQ(rowset_and_bitmap->first, nullptr); + EXPECT_EQ(rowset_and_bitmap->second, nullptr); + ctx.reset(); + wg.reset(); + EXPECT_FALSE(weak_wg.expired()); + + cache.remove_unused_tablet_txn_info(1, 10); + EXPECT_FALSE(cache.is_empty_rowset(1, 10)); + EXPECT_EQ(cache.get_workload_group(1, 10), nullptr); + EXPECT_FALSE(weak_wg.expired()); + { + std::unique_lock lock(cache._rwlock); + CloudTxnDeleteBitmapCache::TxnKey key(2, 10); + // A stale expiration entry must not remove a renewed marker or its owner. + cache._expiration_txn.emplace(0, key); + } + cache.remove_expired_tablet_txn_info(); + EXPECT_TRUE(cache.is_empty_rowset(2, 10)); + EXPECT_FALSE(weak_wg.expired()); + { + std::unique_lock lock(cache._rwlock); + CloudTxnDeleteBitmapCache::TxnKey key(2, 10); + cache._empty_rowset_markers.at(key).txn_expiration = 0; + cache._expiration_txn.emplace(0, key); + } + cache.remove_expired_tablet_txn_info(); + EXPECT_FALSE(cache.is_empty_rowset(2, 10)); + EXPECT_EQ(cache.get_workload_group(2, 10), nullptr); + EXPECT_TRUE(weak_wg.expired()); +} + +TEST(CloudTxnDeleteBitmapCacheTest, ContextlessEmptyRowsetClearsPreviousOwner) { + CloudTxnDeleteBitmapCache cache(1024 * 1024); + ASSERT_TRUE(cache.init().ok()); + auto wg = std::make_shared( + WorkloadGroupInfo {.id = 68387, .name = "contextless_empty_bitmap_cache_test"}); + auto ctx = ResourceContext::create_shared(); + ctx->memory_context()->set_mem_tracker(MemTrackerLimiter::create_shared( + MemTrackerLimiter::Type::OTHER, "contextless_empty_bitmap_cache_test")); + ctx->set_workload_group(wg); + std::thread worker([&] { + SCOPED_INIT_THREAD_CONTEXT(); + { + SCOPED_ATTACH_TASK(ctx); + cache.mark_empty_rowset(1, 10, INT64_MAX); + EXPECT_EQ(cache.get_workload_group(1, 10), wg); + } + EXPECT_FALSE(thread_context()->is_attach_task()); + cache.mark_empty_rowset(1, 10, INT64_MAX); + cache.mark_empty_rowset(2, 10, INT64_MAX); + EXPECT_EQ(cache.get_workload_group(1, 10), nullptr); + EXPECT_EQ(cache.get_workload_group(2, 10), nullptr); + EXPECT_TRUE(cache.is_empty_rowset(1, 10)); + EXPECT_TRUE(cache.is_empty_rowset(2, 10)); + }); + worker.join(); +} + +} // namespace doris diff --git a/be/test/load/delta_writer/delta_writer_test.cpp b/be/test/load/delta_writer/delta_writer_test.cpp index 1c6d7bd4522011..a387dadea990ca 100644 --- a/be/test/load/delta_writer/delta_writer_test.cpp +++ b/be/test/load/delta_writer/delta_writer_test.cpp @@ -33,6 +33,7 @@ #include #include "common/config.h" +#include "common/metrics/metrics.h" #include "common/object_pool.h" #include "core/block/block.h" #include "core/block/column_with_type_and_name.h" @@ -44,15 +45,22 @@ #include "exprs/function/cast/cast_to_datev2_impl.hpp" #include "gtest/gtest_pred_impl.h" #include "io/fs/local_file_system.h" +#include "load/channel/load_stream.h" +#include "load/channel/load_stream_mgr.h" #include "load/memtable/memtable_memory_limiter.h" #include "runtime/descriptor_helper.h" #include "runtime/descriptors.h" #include "runtime/exec_env.h" +#include "runtime/thread_context.h" +#include "runtime/workload_group/workload_group.h" +#include "runtime/workload_management/resource_context.h" #include "storage/data_dir.h" +#include "storage/delete/calc_delete_bitmap_executor.h" #include "storage/iterators.h" #include "storage/olap_define.h" #include "storage/options.h" #include "storage/rowset/beta_rowset.h" +#include "storage/rowset/rowset_writer.h" #include "storage/rowset_builder.h" #include "storage/schema.h" #include "storage/segment/segment.h" @@ -62,6 +70,8 @@ #include "storage/tablet_info.h" #include "storage/task/engine_publish_version_task.h" #include "storage/txn/txn_manager.h" +#include "util/defer_op.h" +#include "util/threadpool.h" namespace doris { class OlapMeta; @@ -850,6 +860,191 @@ TEST_F(TestDeltaWriter, vec_sequence_col) { ASSERT_TRUE(res.ok()); } +TEST_F(TestDeltaWriter, LocalPublishRetainsWorkloadGroup) { + const auto old_segcompaction = config::enable_segcompaction; + config::enable_segcompaction = false; + Defer restore_config {[&] { config::enable_segcompaction = old_segcompaction; }}; + enum class CommitPath { CHANNEL, STREAM_PRE_CLOSE, STREAM_CLOSE }; + int64_t case_id = 0; + for (const auto& [path, async_publish] : + {std::pair {CommitPath::CHANNEL, false}, std::pair {CommitPath::CHANNEL, true}, + std::pair {CommitPath::STREAM_PRE_CLOSE, false}, + std::pair {CommitPath::STREAM_PRE_CLOSE, true}, + std::pair {CommitPath::STREAM_CLOSE, false}, + std::pair {CommitPath::STREAM_CLOSE, true}}) { + for (bool drop_group : {false, true}) { + SCOPED_TRACE(testing::Message() << "path=" << static_cast(path) << ", async=" + << async_publish << ", dropped=" << drop_group); + ++case_id; + RuntimeProfile profile("local_publish_workload_group"); + TCreateTabletReq request; + create_tablet_request_with_sequence_col(168385 + case_id, 270068377, &request, true); + ASSERT_TRUE(engine_ref->create_tablet(request, &profile).ok()); + auto tablet = engine_ref->tablet_manager()->get_tablet(request.tablet_id); + ASSERT_NE(tablet, nullptr); + + TDescriptorTable tdesc_tbl = create_descriptor_tablet_with_sequence_col(); + ObjectPool obj_pool; + DescriptorTbl* desc_tbl = nullptr; + ASSERT_TRUE(DescriptorTbl::create(&obj_pool, tdesc_tbl, &desc_tbl).ok()); + auto* tuple_desc = desc_tbl->get_tuple_descriptor(0); + auto wg = std::make_shared( + WorkloadGroupInfo {.id = 68385, .name = "local_publish_test"}); + ASSERT_TRUE(ThreadPoolBuilder("local_publish_owner") + .set_max_threads(1) + .build(&wg->_memtable_flush_pool) + .ok()); + auto* owner_pool = wg->get_memtable_flush_pool(); + auto ctx = ResourceContext::create_shared(); + ctx->memory_context()->set_mem_tracker(MemTrackerLimiter::create_shared( + MemTrackerLimiter::Type::LOAD, "local_publish_writer")); + ctx->set_workload_group(wg); + const int64_t seed_txn = 268385 + case_id * 2; + const int64_t txn_id = seed_txn + 1; + auto write_and_commit = [&](int64_t id, int segments) -> Status { + SCOPED_ATTACH_TASK(ctx); + WriteRequest req; + req.tablet_id = request.tablet_id; + req.schema_hash = request.tablet_schema.schema_hash; + req.partition_id = request.partition_id; + req.txn_id = id; + req.load_id.set_hi(0); + req.load_id.set_lo(id); + req.tuple_desc = tuple_desc; + req.slots = &tuple_desc->slots(); + req.table_schema_param = std::make_shared(); + DeltaWriter writer(*engine_ref, req, &profile, UniqueId(req.load_id)); + for (int i = 0; i < segments; ++i) { + Block block; + for (const auto* slot : tuple_desc->slots()) { + block.insert(ColumnWithTypeAndName(slot->get_empty_mutable_column(), + slot->type(), slot->col_name())); + } + generate_data(&block, static_cast(10 + i), 123, 100); + RETURN_IF_ERROR(writer.write(&block, TabletAddRowsPayload {.row_idxs = {0}})); + RETURN_IF_ERROR(writer.flush_memtable_async()); + RETURN_IF_ERROR(writer.wait_flush()); + } + RETURN_IF_ERROR(writer.close()); + RETURN_IF_ERROR(writer.wait_flush()); + RETURN_IF_ERROR(writer.build_rowset()); + RETURN_IF_ERROR(writer.submit_calc_delete_bitmap_task()); + RETURN_IF_ERROR(writer.wait_calc_delete_bitmap()); + return writer.commit_txn(); + }; + ASSERT_TRUE(write_and_commit(seed_txn, 1).ok()); + auto stream_write_and_commit = [&]() -> Status { + FifoThreadPool heavy_work_pool(1, 16, "publish_stream_close"); + LoadStreamMgr stream_mgr(1); + stream_mgr.set_heavy_work_pool(&heavy_work_pool); + PUniqueId load_id; + load_id.set_hi(0); + load_id.set_lo(txn_id); + TabletStream stream(load_id, request.tablet_id, txn_id, &stream_mgr, &profile); + { + SCOPED_ATTACH_TASK(ctx); + RETURN_IF_ERROR(stream.init(std::make_shared(), 0, + request.partition_id)); + // Generate real segment files directly; the regression under test + // starts at the stream's close-to-commit worker handoff. + auto rowset_writer = + stream._load_stream_writer->_rowset_builder->rowset_writer(); + for (int i = 0; i < 2; ++i) { + Block block; + for (const auto* slot : tuple_desc->slots()) { + block.insert(ColumnWithTypeAndName(slot->get_empty_mutable_column(), + slot->type(), slot->col_name())); + } + generate_data(&block, static_cast(10 + i), 123, 100); + RETURN_IF_ERROR(rowset_writer->add_block(&block)); + RETURN_IF_ERROR(rowset_writer->flush()); + } + } + // The input files bypassed stream append counters. Both close paths + // must preserve the saved owner after _pre_close() detaches it. + stream.disable_num_segments_check(); + if (path == CommitPath::STREAM_PRE_CLOSE) { + stream.pre_close(); + } + RETURN_IF_ERROR(stream.close()); + // Reuse the same worker to verify close did not leave its context attached. + return stream._run_in_heavy_work_pool([]() { + SCOPED_INIT_THREAD_CONTEXT(); + EXPECT_FALSE(thread_context()->is_attach_task()); + return Status::OK(); + }); + }; + auto commit_status = path == CommitPath::CHANNEL ? write_and_commit(txn_id, 2) + : stream_write_and_commit(); + ASSERT_TRUE(commit_status.ok()) << commit_status; + ctx.reset(); + + // Publish another rowset after the target committed, so publishing the + // two-segment target must submit new bitmap work against that rowset. + std::map rowsets; + engine_ref->txn_manager()->get_txn_related_tablets(seed_txn, request.partition_id, + &rowsets); + ASSERT_EQ(rowsets.size(), 1); + TabletPublishTxnTask seed_task(*engine_ref, nullptr, tablet, rowsets.begin()->second, + {}, request.partition_id, seed_txn, Version(2, 2), + rowsets.begin()->first, -1); + seed_task.handle(); + ASSERT_TRUE(seed_task.result().ok()) << seed_task.result(); + + rowsets.clear(); + std::map> infos; + engine_ref->txn_manager()->get_txn_related_tablets(txn_id, request.partition_id, + &rowsets, &infos); + ASSERT_EQ(rowsets.size(), 1); + ASSERT_EQ(infos.size(), 1); + ASSERT_EQ(rowsets.begin()->second->num_segments(), 2); + ASSERT_EQ(infos.begin()->second->workload_group, wg); + if (drop_group) { + wg->shutdown(); + ASSERT_TRUE(wg->can_be_dropped()); + owner_pool->shutdown(); + } + owner_pool->wait(); + + // Give the default domain a dedicated pool so execution counters prove + // the selected domain without interference from other engine tasks. + std::unique_ptr default_pool; + ASSERT_TRUE(ThreadPoolBuilder("local_publish_default") + .set_max_threads(1) + .build(&default_pool) + .ok()); + auto* executor = engine_ref->calc_delete_bitmap_executor(); + auto* old_load_pool = executor->_load_pool; + executor->_load_pool = default_pool.get(); + Defer restore_pool {[&] { executor->_load_pool = old_load_pool; }}; + const auto owner_before = owner_pool->thread_pool_task_execution_count_total->value(); + if (async_publish) { + AsyncTabletPublishTask task(*engine_ref, tablet, request.partition_id, txn_id, 3, + -1); + task.handle(); + } else { + TabletPublishTxnTask task(*engine_ref, nullptr, tablet, rowsets.begin()->second, {}, + request.partition_id, txn_id, Version(3, 3), + rowsets.begin()->first, -1); + task.handle(); + ASSERT_TRUE(task.result().ok()) << task.result(); + } + owner_pool->wait(); + default_pool->wait(); + EXPECT_EQ(tablet->get_rowset_with_max_version()->end_version(), 3); + const auto owner_tasks = + owner_pool->thread_pool_task_execution_count_total->value() - owner_before; + const auto default_tasks = + default_pool->thread_pool_task_execution_count_total->value(); + EXPECT_EQ(owner_tasks, drop_group ? 0 : 2); + EXPECT_EQ(default_tasks, drop_group ? 2 : 0); + ASSERT_TRUE(engine_ref->tablet_manager() + ->drop_tablet(request.tablet_id, request.replica_id, false) + .ok()); + } + } +} + TEST_F(TestDeltaWriter, vec_sequence_col_concurrent_write) { RuntimeProfile profile("CreateTablet"); TCreateTabletReq request; diff --git a/be/test/load/memtable/memtable_flush_executor_test.cpp b/be/test/load/memtable/memtable_flush_executor_test.cpp index ad6e27dd0c580b..d53888fafb3db1 100644 --- a/be/test/load/memtable/memtable_flush_executor_test.cpp +++ b/be/test/load/memtable/memtable_flush_executor_test.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include "common/config.h" #include "exec/sink/autoinc_buffer.h" @@ -48,11 +49,22 @@ #include "storage/tablet/tablet_meta_manager.h" #include "storage/utils.h" #include "testutil/creators.h" +#include "util/countdown_latch.h" +#include "util/defer_op.h" namespace doris { namespace { +class FlushOrderTask final : public Runnable { +public: + explicit FlushOrderTask(std::vector* order) : _order(order) {} + void run() override { _order->push_back(3); } + +private: + std::vector* _order; +}; + class MockRowsetWriter final : public RowsetWriter { public: // `flush_enter_cnt`, when set, is incremented at the top of flush_memtable() BEFORE the @@ -327,6 +339,43 @@ void tear_down() { .ok()); } +TEST(MemTableFlushExecutorTest, FlushSharesItsTransactionTurnForEveryKeyType) { + using namespace std::chrono_literals; + for (auto keys_type : {DUP_KEYS, UNIQUE_KEYS, AGG_KEYS}) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("flush_load_fifo_test").set_max_threads(1).build(&pool).ok()); + std::atomic flush_count = 0; + auto writer = std::make_shared(&flush_count); + RowsetWriterContext context; + context.txn_id = 1; + context.tablet_schema = std::make_shared(); + context.tablet_schema->_keys_type = keys_type; + ASSERT_TRUE(writer->init(context).ok()); + auto flush = FlushToken::create_shared(pool.get(), nullptr); + flush->set_rowset_writer(writer); + auto own_bitmap = + pool->new_load_token(context.txn_id, LoadTaskPriority::MID, LoadTaskType::LEAF); + auto other_bitmap = pool->new_load_token(2, LoadTaskPriority::HIGHEST, LoadTaskType::LEAF); + CountDownLatch entered(1), release(1); + std::vector order; + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(pool->submit_func([&] { + entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(entered.wait_for(5s)); + EXPECT_TRUE(flush->_submit_sub_tasks(pool.get(), {std::make_shared(&order), + std::make_shared(&order)}) + .ok()); + EXPECT_TRUE(other_bitmap->submit_func([&] { order.push_back(2); }).ok()); + EXPECT_TRUE(own_bitmap->submit_func([&] { order.push_back(1); }).ok()); + release.count_down(); + pool->wait(); + // Our bitmap precedes our flushes, but another load gets the next turn. + EXPECT_EQ(order, (std::vector {1, 2, 3, 3})); + } +} + TEST(MemTableFlushExecutorTest, TestDynamicThreadPoolUpdate) { // Setup set_up(); @@ -368,10 +417,11 @@ TEST(MemTableFlushExecutorTest, TestDynamicThreadPoolUpdate) { EXPECT_EQ(actual_max, expected_max); } - // Test 4: Update high_priority_flush_thread_num_per_store + // Test 4: The retired high-priority setting does not resize the shared pool. + int shared_max_threads = flush_executor->flush_pool()->max_threads(); config::high_priority_flush_thread_num_per_store = 8; flush_executor->update_memtable_flush_threads(); - // Note: We can't directly access _high_prio_flush_pool, but update should not crash + EXPECT_EQ(flush_executor->flush_pool()->max_threads(), shared_max_threads); // Test 5: Set very small values config::flush_thread_num_per_store = 0; // Should be adjusted to 1 by std::max diff --git a/be/test/storage/adaptive_thread_pool_controller_test.cpp b/be/test/storage/adaptive_thread_pool_controller_test.cpp index 1590ff5411d2ba..7eb7c5dc836686 100644 --- a/be/test/storage/adaptive_thread_pool_controller_test.cpp +++ b/be/test/storage/adaptive_thread_pool_controller_test.cpp @@ -24,10 +24,15 @@ #include #include +#include "cloud/cloud_storage_engine.h" #include "common/config.h" #include "common/metrics/metrics.h" #include "common/metrics/system_metrics.h" #include "cpp/sync_point.h" +#include "load/memtable/memtable_flush_executor.h" +#include "runtime/exec_env.h" +#include "runtime/workload_group/workload_group.h" +#include "runtime/workload_group/workload_group_manager.h" #include "testutil/test_util.h" #include "util/defer_op.h" #include "util/threadpool.h" @@ -142,6 +147,71 @@ class AdaptiveThreadPoolControllerTest : public testing::Test { std::unique_ptr _pool2; }; +TEST_F(AdaptiveThreadPoolControllerTest, RetiredConfigPreservesAdaptiveReduction) { + config::enable_adaptive_flush_threads = true; + const auto old_flush_threads = config::flush_thread_num_per_store; + const auto old_high_priority_threads = config::high_priority_flush_thread_num_per_store; + const auto old_max_per_cpu = config::max_flush_thread_num_per_cpu; + const auto old_min_per_cpu = config::min_flush_thread_num_per_cpu; + Defer restore_config {[&] { + config::flush_thread_num_per_store = old_flush_threads; + config::max_flush_thread_num_per_cpu = old_max_per_cpu; + config::min_flush_thread_num_per_cpu = old_min_per_cpu; + EXPECT_TRUE(config::set_config("high_priority_flush_thread_num_per_store", + std::to_string(old_high_priority_threads)) + .ok()); + }}; + config::flush_thread_num_per_store = 8; + config::max_flush_thread_num_per_cpu = 4; + config::min_flush_thread_num_per_cpu = 0; + + auto engine = std::make_unique(EngineOptions()); + engine->_memtable_flush_executor = std::make_unique(); + auto* default_pool = _pool.get(); + engine->_memtable_flush_executor->_flush_pool = std::move(_pool); + auto wg = std::make_shared( + WorkloadGroupInfo {.id = 68389, .name = "retired_flush_config_test"}); + auto* wg_pool = _pool2.get(); + wg->_memtable_flush_pool = std::move(_pool2); + WorkloadGroupMgr wg_mgr; + wg_mgr._workload_groups.emplace(wg->id(), wg); + auto* env = ExecEnv::GetInstance(); + auto old_engine = std::move(env->_storage_engine); + auto* old_wg_mgr = env->_workload_group_manager; + env->set_storage_engine(std::move(engine)); + env->_workload_group_manager = &wg_mgr; + Defer restore_env {[&] { + env->_workload_group_manager = old_wg_mgr; + env->set_storage_engine(std::move(old_engine)); + }}; + + bool reduce = true; + AdaptiveThreadPoolController controller; + auto adjust = [&](int current, int min_threads, int, std::string&) { + return reduce ? min_threads : current; + }; + // Drive ticks explicitly; the long timer interval keeps the sequence deterministic. + controller.add("default", {default_pool}, adjust, 4, 0, 3600000); + controller.add("workload", {wg_pool}, adjust, 4, 0, 3600000); + controller.adjust_once(); + reduce = false; + ASSERT_EQ(default_pool->max_threads(), 1); + ASSERT_EQ(wg_pool->max_threads(), 1); + ASSERT_EQ(controller.get_current_threads("default"), 1); + ASSERT_EQ(controller.get_current_threads("workload"), 1); + + // set_config invokes the real update_config callback for both resource domains. + ASSERT_TRUE(config::set_config("high_priority_flush_thread_num_per_store", "7").ok()); + EXPECT_EQ(config::high_priority_flush_thread_num_per_store, 7); + EXPECT_EQ(default_pool->max_threads(), 1); + EXPECT_EQ(wg_pool->max_threads(), 1); + controller.adjust_once(); // No new pressure signal: keep the reduced limits. + EXPECT_EQ(default_pool->max_threads(), 1); + EXPECT_EQ(wg_pool->max_threads(), 1); + EXPECT_EQ(controller.get_current_threads("default"), 1); + EXPECT_EQ(controller.get_current_threads("workload"), 1); +} + // Test basic add and get_current_threads TEST_F(AdaptiveThreadPoolControllerTest, TestAddPoolGroup) { AdaptiveThreadPoolController controller; diff --git a/be/test/storage/txn/txn_manager_test.cpp b/be/test/storage/txn/txn_manager_test.cpp index 299e9cc73335b5..f12243f2a1b4b9 100644 --- a/be/test/storage/txn/txn_manager_test.cpp +++ b/be/test/storage/txn/txn_manager_test.cpp @@ -31,9 +31,13 @@ #include #include #include +#include #include "common/config.h" #include "gtest/gtest_pred_impl.h" +#include "runtime/thread_context.h" +#include "runtime/workload_group/workload_group.h" +#include "runtime/workload_management/resource_context.h" #include "storage/olap_meta.h" #include "storage/options.h" #include "storage/rowset/rowset.h" @@ -323,6 +327,85 @@ TEST_F(TxnManagerTest, CommitTxnTwiceWithSameRowsetId) { EXPECT_TRUE(k_engine->pending_local_rowsets().contains(_rowset->rowset_id())); } +TEST_F(TxnManagerTest, CommitRetainsWorkloadGroupAcrossRetryAndCleanup) { + auto wg = std::make_shared( + WorkloadGroupInfo {.id = 68385, .name = "local_publish_owner"}); + std::weak_ptr weak_wg = wg; + auto ctx = ResourceContext::create_shared(); + ctx->memory_context()->set_mem_tracker( + MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::LOAD, "local_publish_owner")); + ctx->set_workload_group(wg); + { + SCOPED_ATTACH_TASK(ctx); + auto guard = k_engine->pending_local_rowsets().add(_rowset->rowset_id()); + ASSERT_TRUE(k_engine->txn_manager() + ->commit_txn(_meta.get(), partition_id, transaction_id, tablet_id, + _tablet_uid, load_id, _rowset, std::move(guard), false) + .ok()); + } + ctx.reset(); + wg.reset(); + ASSERT_FALSE(weak_wg.expired()); + + // A repeated commit from a different request context must retain the original owner. + { + SCOPED_ATTACH_TASK(MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "local_publish_retry")); + auto guard = k_engine->pending_local_rowsets().add(_rowset_same_id->rowset_id()); + ASSERT_TRUE(k_engine->txn_manager() + ->commit_txn(_meta.get(), partition_id, transaction_id, tablet_id, + _tablet_uid, load_id, _rowset_same_id, std::move(guard), + false) + .ok()); + } + std::map rowsets; + std::map> infos; + k_engine->txn_manager()->get_txn_related_tablets(transaction_id, partition_id, &rowsets, + &infos); + ASSERT_EQ(infos.size(), 1); + EXPECT_EQ(infos.begin()->second->workload_group, weak_wg.lock()); + infos.clear(); + ASSERT_TRUE( + k_engine->txn_manager() + ->delete_txn(_meta.get(), partition_id, transaction_id, tablet_id, _tablet_uid) + .ok()); + EXPECT_TRUE(weak_wg.expired()); +} + +TEST_F(TxnManagerTest, ContextlessAndRecoveredCommitsHaveNoWorkloadGroup) { + auto wg = std::make_shared( + WorkloadGroupInfo {.id = 68386, .name = "local_publish_recovery"}); + auto ctx = ResourceContext::create_shared(); + ctx->memory_context()->set_mem_tracker(MemTrackerLimiter::create_shared( + MemTrackerLimiter::Type::OTHER, "local_publish_recovery")); + ctx->set_workload_group(wg); + std::thread worker([&] { + SCOPED_INIT_THREAD_CONTEXT(); + auto guard = k_engine->pending_local_rowsets().add(_rowset->rowset_id()); + ASSERT_FALSE(thread_context()->is_attach_task()); + ASSERT_TRUE(k_engine->txn_manager() + ->commit_txn(_meta.get(), partition_id, transaction_id, tablet_id, + _tablet_uid, load_id, _rowset, std::move(guard), false) + .ok()); + // Recovery must not assign an unrelated caller's group to a persisted transaction. + SCOPED_ATTACH_TASK(ctx); + auto recovery_guard = k_engine->pending_local_rowsets().add(_rowset_diff_id->rowset_id()); + ASSERT_TRUE(k_engine->txn_manager() + ->commit_txn(_meta.get(), partition_id, transaction_id + 1, tablet_id, + _tablet_uid, load_id, _rowset_diff_id, + std::move(recovery_guard), true) + .ok()); + }); + worker.join(); + for (auto txn_id : {transaction_id, transaction_id + 1}) { + std::map rowsets; + std::map> infos; + k_engine->txn_manager()->get_txn_related_tablets(txn_id, partition_id, &rowsets, &infos); + ASSERT_EQ(infos.size(), 1); + EXPECT_EQ(infos.begin()->second->workload_group, nullptr); + } +} + // 1. prepare twice should be success TEST_F(TxnManagerTest, PrepareNewTxnTwice) { auto st = k_engine->txn_manager()->prepare_txn(partition_id, transaction_id, tablet_id, diff --git a/be/test/util/load_task_queue_test.cpp b/be/test/util/load_task_queue_test.cpp new file mode 100644 index 00000000000000..8029e1a2a9a455 --- /dev/null +++ b/be/test/util/load_task_queue_test.cpp @@ -0,0 +1,163 @@ +// 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 "util/load_task_queue.h" + +#include + +#include +#include + +namespace doris { + +TEST(LoadTaskQueueTest, LoadFifoThenPriority) { + LoadTaskQueue queue; + queue.push(1, 3, 13); + queue.push(1, 2, 12); + queue.push(1, 1, 11); + queue.push(1, 0, 10); + queue.push(2, 3, 23); // DUP load gets a turn despite load 1's bitmap work. + std::vector actual; + while (!queue.empty()) { + actual.push_back(queue.pop()); + } + EXPECT_EQ(actual, (std::vector {10, 23, 11, 12, 13})); +} + +TEST(LoadTaskQueueTest, PriorityDoesNotCrossLoadsAndSamePriorityIsFifo) { + LoadTaskQueue queue; + queue.push(1, 3, 1); + queue.push(2, 0, 2); + queue.push(1, 3, 3); + queue.push(2, 0, 4); + for (int expected : {1, 2, 3, 4}) { + EXPECT_EQ(queue.pop(), expected); + } + EXPECT_TRUE(queue.empty()); +} + +TEST(LoadTaskQueueTest, RequeueBeforeExecutionAndReactivateEmptyLoad) { + LoadTaskQueue queue; + queue.push(1, 3, 1); + queue.push(1, 3, 2); + queue.push(2, 3, 3); + EXPECT_EQ(queue.pop(), 1); // Task 1 need not finish before the next turn. + EXPECT_EQ(queue.pop(), 3); + EXPECT_EQ(queue.pop(), 2); + queue.push(2, 3, 4); + queue.push(1, 0, 5); // Empty -> nonempty puts load 1 at the tail exactly once. + queue.push(1, 0, 6); + for (int expected : {4, 5, 6}) { + EXPECT_EQ(queue.pop(), expected); + } + EXPECT_TRUE(queue.empty()); +} + +TEST(LoadTaskQueueTest, CancelOneTokenPreservesOtherTasksAndLoadOrder) { + LoadTaskQueue queue; + queue.push(1, 0, 1); + queue.push(1, 2, 2); + queue.push(2, 3, 3); + queue.push(1, 3, 4); + EXPECT_EQ(queue.remove_if(1, [](int task) { return task < 3; }), (std::vector {1, 2})); + EXPECT_EQ(queue.size(), 2); + EXPECT_EQ(queue.pop(), 4); + EXPECT_EQ(queue.pop(), 3); + queue.push(1, 3, 5); + queue.push(2, 0, 6); + EXPECT_EQ(queue.remove_if(1, [](int) { return true; }), (std::vector {5})); + queue.push(1, 0, 7); + EXPECT_EQ(queue.pop(), 6); + EXPECT_EQ(queue.pop(), 7); + EXPECT_TRUE(queue.remove_if(99, [](int) { return true; }).empty()); + EXPECT_TRUE(queue.empty()); +} + +TEST(LoadTaskQueueTest, MoveOnlyTasksAndDeferredDestruction) { + LoadTaskQueue> queue; + queue.push(1, 0, std::make_unique(1)); + queue.push(1, 0, std::make_unique(2)); + auto removed = queue.remove_if(1, [](const auto& task) { return *task == 1; }); + ASSERT_EQ(removed.size(), 1); + EXPECT_EQ(*removed.front(), 1); + EXPECT_EQ(*queue.pop(), 2); + EXPECT_TRUE(queue.empty()); +} + +TEST(LoadTaskQueueTest, BusyHighPriorityLoadDoesNotStarveAnotherLoad) { + LoadTaskQueue queue; + queue.push(1, 0, 10); + queue.push(1, 0, 11); + queue.push(2, 3, 20); + queue.push(2, 3, 21); + EXPECT_EQ(queue.pop(), 10); + queue.push(1, 0, 12); + EXPECT_EQ(queue.pop(), 20); + EXPECT_EQ(queue.pop(), 11); + queue.push(1, 0, 13); + EXPECT_EQ(queue.pop(), 21); + EXPECT_EQ(queue.pop(), 12); + EXPECT_EQ(queue.pop(), 13); + EXPECT_TRUE(queue.empty()); +} + +TEST(LoadTaskQueueTest, CancelMiddleLoadDoesNotInspectOtherLoads) { + LoadTaskQueue queue; + queue.push(1, 3, 10); + queue.push(2, 0, 20); + queue.push(2, 3, 21); + queue.push(3, 0, 30); + int inspected = 0; + auto removed = queue.remove_if(2, [&](int task) { + ++inspected; + EXPECT_TRUE(task == 20 || task == 21); + return true; + }); + EXPECT_EQ(inspected, 2); + EXPECT_EQ(removed, (std::vector {20, 21})); + EXPECT_TRUE(queue.remove_if(2, [&](int) { + ADD_FAILURE() << "an empty load inspected another load's tasks"; + return true; + }).empty()); + queue.push(2, 0, 22); + for (int expected : {10, 30, 22}) { + EXPECT_EQ(queue.pop(), expected); + } + EXPECT_TRUE(queue.empty()); +} + +TEST(LoadTaskQueueTest, EraseHandlePreservesOtherHandlesAndTurns) { + LoadTaskQueue queue; + queue.push(1, 0, 10); + queue.push(1, 0, 11); + auto middle = queue.push(1, 0, 12); + auto last = queue.push(1, 0, 13); + auto other_load = queue.push(2, 0, 20); + queue.push(3, 0, 30); + EXPECT_EQ(queue.pop(), 10); // Rotate load 1 before erasing its own remaining tasks. + queue.erase(middle); + EXPECT_EQ(queue.remove_if(1, [](int task) { return task == 11; }), (std::vector {11})); + queue.erase(last); // remove_if must not invalidate another entry's handle. + queue.erase(other_load); // Remove another load from the ready list. + queue.push(1, 0, 14); + EXPECT_EQ(queue.size(), 2); + EXPECT_EQ(queue.pop(), 30); + EXPECT_EQ(queue.pop(), 14); + EXPECT_TRUE(queue.empty()); +} + +} // namespace doris diff --git a/be/test/util/load_threadpool_test.cpp b/be/test/util/load_threadpool_test.cpp new file mode 100644 index 00000000000000..5a4c26ecbb7598 --- /dev/null +++ b/be/test/util/load_threadpool_test.cpp @@ -0,0 +1,573 @@ +// 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 + +#include +#include +#include +#include +#include +#include + +#include "common/signal_handler.h" +#include "runtime/thread_context.h" +#include "runtime/workload_management/resource_context.h" +#include "storage/delete/calc_delete_bitmap_executor.h" +#include "util/countdown_latch.h" +#include "util/defer_op.h" +#include "util/threadpool.h" + +namespace doris { +using namespace std::chrono_literals; + +TEST(LoadThreadPoolTest, MultipleTokensShareOneLoadTurn) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("load_fifo_test").set_max_threads(1).build(&pool).ok()); + auto flush = pool->new_load_token(1, LoadTaskPriority::LOW, LoadTaskType::PARENT); + auto bitmap = pool->new_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT); + auto dup = pool->new_load_token(2, LoadTaskPriority::LOW, LoadTaskType::PARENT); + CountDownLatch entered(1), release(1); + std::vector order; + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(pool->submit_func([&] { + entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(entered.wait_for(5s)); + EXPECT_TRUE(flush->submit_func([&] { order.push_back(3); }).ok()); + EXPECT_TRUE(bitmap->submit_func([&] { order.push_back(0); }).ok()); + EXPECT_TRUE(dup->submit_func([&] { order.push_back(2); }).ok()); + release.count_down(); + pool->wait(); + EXPECT_EQ(order, (std::vector {0, 2, 3})); +} + +TEST(LoadThreadPoolTest, TokenlessTasksKeepTheirLoadAndPriority) { + class RecordTask : public Runnable { + public: + RecordTask(std::vector* order, int value) : _order(order), _value(value) {} + void run() override { _order->push_back(_value); } + + private: + std::vector* _order; + int _value; + }; + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("load_tokenless_order").set_max_threads(1).build(&pool).ok()); + CountDownLatch entered(1), release(1); + std::vector order; + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(pool->submit_func([&] { + entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(entered.wait_for(5s)); + EXPECT_TRUE( + pool->submit_load(std::make_shared(&order, 13), 1, LoadTaskPriority::LOW) + .ok()); + EXPECT_TRUE( + pool->submit_load(std::make_shared(&order, 23), 2, LoadTaskPriority::LOW) + .ok()); + EXPECT_TRUE(pool->submit_load(std::make_shared(&order, 10), 1, + LoadTaskPriority::HIGHEST) + .ok()); + release.count_down(); + pool->wait(); + EXPECT_EQ(order, (std::vector {10, 23, 13})); +} + +TEST(LoadThreadPoolTest, OneLoadCanUseAllWorkers) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("load_parallel_test").set_max_threads(2).build(&pool).ok()); + auto token = pool->new_load_token(1, LoadTaskPriority::LOW, LoadTaskType::PARENT); + CountDownLatch entered(2), release(1); + Defer unblock = [&] { release.count_down(); }; + for (int i = 0; i < 2; ++i) { + EXPECT_TRUE(token->submit_func([&] { + entered.count_down(); + release.wait(); + }).ok()); + } + EXPECT_TRUE(entered.wait_for(5s)); + release.count_down(); + token->wait(); +} + +TEST(LoadThreadPoolTest, CancelOnlyRemovesItsOwnTasks) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("load_cancel_test").set_max_threads(1).build(&pool).ok()); + auto cancelled = pool->new_load_token(1, LoadTaskPriority::MID, LoadTaskType::LEAF); + auto kept = pool->new_load_token(1, LoadTaskPriority::LOW, LoadTaskType::PARENT); + auto other_load = pool->new_load_token(2, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT); + CountDownLatch entered(1), release(1); + std::vector order; + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(pool->submit_func([&] { + entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(entered.wait_for(5s)); + EXPECT_TRUE(cancelled->submit_func([&] { ADD_FAILURE() << "cancelled task ran"; }).ok()); + EXPECT_TRUE(kept->submit_func([&] { order.push_back(1); }).ok()); + EXPECT_TRUE(other_load->submit_func([&] { order.push_back(2); }).ok()); + cancelled->shutdown(); + cancelled->shutdown(); + EXPECT_FALSE(cancelled->submit_func([] {}).ok()); + release.count_down(); + pool->wait(); + EXPECT_EQ(order, (std::vector {1, 2})); +} + +TEST(LoadThreadPoolTest, NestedBitmapHelpsOnlyOwnTokenWithOneWorker) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("load_nested_test").set_max_threads(1).build(&pool).ok()); + CalcDeleteBitmapExecutor executor; + executor.init("background_bitmap_test", 1, pool.get()); + auto resource_ctx = ResourceContext::create_shared(); + auto request_tracker = + MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, "load_nested_request"); + auto tablet_tracker = + MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, "load_nested_tablet"); + resource_ctx->memory_context()->set_mem_tracker(request_tracker); + TUniqueId task_id; + task_id.hi = 1; + task_id.lo = 2; + resource_ctx->task_controller()->set_task_id(task_id); + SCOPED_ATTACH_TASK(resource_ctx); + auto parent = + executor.create_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT, nullptr); + std::atomic completed = 0; + std::atomic unrelated_ran = false; + auto unrelated = pool->new_load_token(2, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT); + EXPECT_TRUE( + parent->submit_func([&] { + EXPECT_EQ(thread_context()->resource_ctx(), resource_ctx); + EXPECT_EQ(thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker(), + request_tracker.get()); + // Match the cloud tablet handler's tracker switch inside an + // already attached bitmap callback. + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tablet_tracker); + auto check_context = [&] { + EXPECT_TRUE(thread_context()->is_attach_task()); + EXPECT_EQ(thread_context()->resource_ctx(), resource_ctx); + EXPECT_EQ(signal::query_id_hi, task_id.hi); + EXPECT_EQ(signal::query_id_lo, task_id.lo); + EXPECT_EQ(thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker(), + tablet_tracker.get()); + }; + EXPECT_TRUE(unrelated->submit_func([&] { unrelated_ran = true; }).ok()); + auto child = executor.create_load_token(1, LoadTaskPriority::HIGHEST, + LoadTaskType::LEAF, nullptr); + for (int i = 0; i < 2; ++i) { + EXPECT_TRUE(child->submit_func([&] { + check_context(); + EXPECT_TRUE(ThreadPool::is_helping_load_task()); + EXPECT_FALSE(unrelated_ran.load()); + EXPECT_EQ(pool->num_active_threads(), 1); + ++completed; + return Status::OK(); + }).ok()); + EXPECT_EQ(completed.load(), 0); + check_context(); + } + EXPECT_TRUE(child->submit_func([&] { + check_context(); + ++completed; + return Status::InternalError("test bitmap failure"); + }).ok()); + EXPECT_FALSE(child->wait().ok()); + EXPECT_FALSE(unrelated_ran.load()); + EXPECT_FALSE(ThreadPool::is_helping_load_task()); + child.reset(); // No scheduler reference may survive this destruction. + auto next = executor.create_load_token(1, LoadTaskPriority::HIGHEST, + LoadTaskType::LEAF, nullptr); + EXPECT_TRUE(next->submit_func([&] { + check_context(); + ++completed; + return Status::OK(); + }).ok()); + EXPECT_TRUE(next->wait().ok()); + check_context(); + return Status::OK(); + }).ok()); + EXPECT_TRUE(parent->wait().ok()); + EXPECT_EQ(completed.load(), 4); + pool->wait(); + EXPECT_TRUE(unrelated_ran.load()); +} + +TEST(LoadThreadPoolTest, NestedBitmapUsesSpareWorkerAndParentPool) { + std::unique_ptr pool, default_pool; + ASSERT_TRUE(ThreadPoolBuilder("bitmap_parallel_parent").set_max_threads(2).build(&pool).ok()); + ASSERT_TRUE( + ThreadPoolBuilder("bitmap_other_domain").set_max_threads(1).build(&default_pool).ok()); + CalcDeleteBitmapExecutor executor; + executor.init("bitmap_parallel_background", 1, default_pool.get()); + auto resource_ctx = ResourceContext::create_shared(); + auto request_tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "bitmap_parallel_request"); + auto tablet_tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "bitmap_parallel_tablet"); + resource_ctx->memory_context()->set_mem_tracker(request_tracker); + SCOPED_ATTACH_TASK(resource_ctx); + CalcDeleteBitmapToken parent( + pool->new_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT)); + CountDownLatch worker_entered(1), helper_entered(1), release(1); + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE( + parent.submit_func([&] { + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tablet_tracker); + auto child = executor.create_load_token(1, LoadTaskPriority::HIGHEST, + LoadTaskType::LEAF, nullptr); + auto check_context = [&] { + EXPECT_EQ(ThreadPool::current_load_pool(), pool.get()); + EXPECT_EQ(thread_context()->resource_ctx(), resource_ctx); + EXPECT_EQ(thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker(), + tablet_tracker.get()); + }; + EXPECT_TRUE(child->submit_func([&] { + check_context(); + EXPECT_FALSE(ThreadPool::is_helping_load_task()); + worker_entered.count_down(); + release.wait(); + return Status::OK(); + }).ok()); + EXPECT_TRUE(worker_entered.wait_for(5s)); + EXPECT_TRUE(child->submit_func([&] { + check_context(); + EXPECT_TRUE(ThreadPool::is_helping_load_task()); + helper_entered.count_down(); + release.wait(); + return Status::OK(); + }).ok()); + auto st = child->wait(); + check_context(); + EXPECT_FALSE(ThreadPool::is_helping_load_task()); + return st; + }).ok()); + EXPECT_TRUE(helper_entered.wait_for(5s)); + EXPECT_EQ(pool->num_active_threads(), 2); // Helping does not create a third physical worker. + release.count_down(); + EXPECT_TRUE(parent.wait().ok()); +} + +TEST(LoadThreadPoolTest, AllWorkersCanHelpTheirOwnChildren) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("bitmap_saturated").set_max_threads(2).build(&pool).ok()); + CalcDeleteBitmapExecutor executor; + executor.init("bitmap_saturated_background", 1, pool.get()); + SCOPED_ATTACH_TASK( + MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, "bitmap_saturated")); + auto first = + executor.create_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT, nullptr); + auto second = + executor.create_load_token(2, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT, nullptr); + CountDownLatch parents_entered(2), children_entered(2), release(1); + std::atomic completed = 0; + Defer unblock = [&] { release.count_down(); }; + auto run_parent = [&](int64_t load_id) { + parents_entered.count_down(); + EXPECT_TRUE(parents_entered.wait_for(5s)); + auto child = executor.create_load_token(load_id, LoadTaskPriority::HIGHEST, + LoadTaskType::LEAF, nullptr); + EXPECT_TRUE(child->submit_func([&] { + EXPECT_TRUE(ThreadPool::is_helping_load_task()); + children_entered.count_down(); + release.wait(); + ++completed; + return Status::OK(); + }).ok()); + return child->wait(); + }; + EXPECT_TRUE(first->submit_func([&] { return run_parent(1); }).ok()); + EXPECT_TRUE(second->submit_func([&] { return run_parent(2); }).ok()); + EXPECT_TRUE(children_entered.wait_for(5s)); + EXPECT_EQ(pool->num_active_threads(), 2); + release.count_down(); + EXPECT_TRUE(first->wait().ok()); + EXPECT_TRUE(second->wait().ok()); + EXPECT_EQ(completed.load(), 2); +} + +TEST(LoadThreadPoolTest, ParentCanCancelQueuedPublishChildren) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("bitmap_cancel_children").set_max_threads(1).build(&pool).ok()); + CalcDeleteBitmapExecutor executor; + executor.init("bitmap_cancel_background", 1, pool.get()); + SCOPED_ATTACH_TASK(MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "bitmap_cancel_children")); + auto parent = + executor.create_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT, nullptr); + EXPECT_TRUE(parent->submit_func([&] { + auto child = executor.create_load_token(1, LoadTaskPriority::HIGHEST, + LoadTaskType::LEAF, nullptr); + EXPECT_TRUE(child->submit_func([] { + ADD_FAILURE() << "cancelled child ran"; + return Status::OK(); + }).ok()); + child->cancel(); + EXPECT_TRUE(child->wait().is()); + return Status::OK(); + }).ok()); + EXPECT_TRUE(parent->wait().ok()); +} + +TEST(LoadThreadPoolTest, HelpingCallbackExceptionRetiresTask) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("bitmap_help_exception").set_max_threads(1).build(&pool).ok()); + auto parent = pool->new_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT); + EXPECT_TRUE(parent->submit_func([&] { + auto child = pool->new_load_token(1, LoadTaskPriority::HIGHEST, + LoadTaskType::LEAF); + EXPECT_TRUE(child->submit_func([] { + throw std::runtime_error("child failure"); + }).ok()); + EXPECT_THROW( + { + auto st = child->wait_and_help(); + EXPECT_TRUE(st.ok()); + }, + std::runtime_error); + EXPECT_FALSE(ThreadPool::is_helping_load_task()); + EXPECT_EQ(ThreadPool::current_load_pool(), pool.get()); + EXPECT_EQ(child->num_tasks(), 0); + EXPECT_TRUE( + child->wait_and_help() + .ok()); // Failed callback must not leave an active task. + }).ok()); + parent->wait(); + EXPECT_EQ(pool->get_queue_size(), 0); +} + +TEST(LoadThreadPoolTest, HelpRejectsInvalidCallersWithoutChangingTasks) { + std::unique_ptr pool, other_pool; + ASSERT_TRUE(ThreadPoolBuilder("help_contract").set_max_threads(1).build(&pool).ok()); + ASSERT_TRUE(ThreadPoolBuilder("help_other_pool").set_max_threads(1).build(&other_pool).ok()); + // A MID parent proves priority no longer implies a leaf dependency role. + auto parent = pool->new_load_token(1, LoadTaskPriority::MID, LoadTaskType::PARENT); + auto leaf = pool->new_load_token(1, LoadTaskPriority::LOW, LoadTaskType::LEAF); + auto leaf_caller = pool->new_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::LEAF); + auto nonleaf = pool->new_load_token(1, LoadTaskPriority::MID, LoadTaskType::PARENT); + auto ordinary = pool->new_token(ThreadPool::ExecutionMode::CONCURRENT); + auto foreign = other_pool->new_load_token(2, LoadTaskPriority::LOW, LoadTaskType::PARENT); + CountDownLatch entered(1), release(1); + std::atomic completed = 0; + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(pool->submit_func([&] { + EXPECT_TRUE(leaf->wait_and_help().is()); + entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(entered.wait_for(5s)); + EXPECT_TRUE(leaf->submit_func([&] { ++completed; }).ok()); + EXPECT_TRUE(leaf->wait_and_help().is()); + EXPECT_EQ(leaf->num_tasks(), 1); + EXPECT_TRUE(foreign->submit_func([&] { + EXPECT_TRUE(leaf->wait_and_help().is()); + EXPECT_EQ(leaf->num_tasks(), 1); + }) + .ok()); + foreign->wait(); + EXPECT_TRUE(leaf_caller + ->submit_func([&] { + EXPECT_TRUE(leaf->wait_and_help().is()); + EXPECT_TRUE( + leaf_caller->wait_and_help().is()); + EXPECT_EQ(leaf->num_tasks(), 1); + }) + .ok()); + EXPECT_TRUE(parent->submit_func([&] { + EXPECT_TRUE(parent->wait_and_help().is()); + EXPECT_TRUE(nonleaf->wait_and_help().is()); + EXPECT_TRUE(ordinary->wait_and_help().is()); + EXPECT_TRUE(leaf->wait_and_help().ok()); + }).ok()); + release.count_down(); + pool->wait(); + EXPECT_EQ(completed.load(), 1); +} + +TEST(LoadThreadPoolTest, BitmapInvalidHelperWaitCancelsItsQueuedTasks) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("bitmap_invalid_helper").set_max_threads(1).build(&pool).ok()); + SCOPED_ATTACH_TASK(MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "bitmap_invalid_helper")); + CalcDeleteBitmapToken child( + pool->new_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::LEAF), nullptr, true); + CountDownLatch entered(1), release(1); + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(pool->submit_func([&] { + entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(entered.wait_for(5s)); + EXPECT_TRUE(child.submit_func([] { + ADD_FAILURE() << "invalid helper wait left a callback queued"; + return Status::OK(); + }).ok()); + auto st = child.wait(); // An external caller cannot help on a pool worker's behalf. + EXPECT_TRUE(st.is()); + EXPECT_EQ(pool->get_queue_size(), 0); + EXPECT_EQ(child.wait().to_string(), st.to_string()); + EXPECT_EQ(child.submit_func([] { return Status::OK(); }).to_string(), st.to_string()); + release.count_down(); + pool->wait(); +} + +TEST(LoadThreadPoolTest, EnqueueWakesSleepingHelperAndCompletionWakesAllWaiters) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("helper_enqueue_wakeup").set_max_threads(2).build(&pool).ok()); + auto parent = pool->new_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT); + auto child = pool->new_load_token(1, LoadTaskPriority::MID, LoadTaskType::LEAF); + CountDownLatch worker_entered(1), helped(1), release(1); + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(child->submit_func([&] { + worker_entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(worker_entered.wait_for(5s)); + EXPECT_TRUE(parent->submit_func([&] { EXPECT_TRUE(child->wait_and_help().ok()); }).ok()); + bool helper_waiting = false; + auto deadline = std::chrono::steady_clock::now() + 5s; + while (std::chrono::steady_clock::now() < deadline) { + { + std::lock_guard lock(pool->_lock); + helper_waiting = child->_waiting_helpers == 1; + } + if (helper_waiting) { + break; + } + std::this_thread::yield(); + } + EXPECT_TRUE(helper_waiting); + std::thread first_waiter([&] { child->wait(); }); + std::thread second_waiter([&] { child->wait(); }); + EXPECT_TRUE(child->submit_func([&] { + EXPECT_TRUE(ThreadPool::is_helping_load_task()); + helped.count_down(); + }).ok()); + EXPECT_TRUE(helped.wait_for(5s)); // The worker remains blocked; only the helper can execute it. + release.count_down(); + parent->wait(); + first_waiter.join(); + second_waiter.join(); + { + std::lock_guard lock(pool->_lock); + EXPECT_EQ(child->_waiting_helpers, 0); + } + EXPECT_EQ(pool->get_queue_size(), 0); +} + +TEST(LoadThreadPoolTest, CancelledBitmapIsNotReportedAsComplete) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("load_shutdown_test").set_max_threads(1).build(&pool).ok()); + CalcDeleteBitmapExecutor executor; + executor.init("background_shutdown_test", 1, pool.get()); + auto token = executor.create_load_token(1, LoadTaskPriority::MID, LoadTaskType::LEAF, nullptr); + CountDownLatch entered(1), release(1); + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(pool->submit_func([&] { + entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(entered.wait_for(5s)); + EXPECT_TRUE(token->submit_func([] { return Status::OK(); }).ok()); + token->cancel(); + EXPECT_TRUE(token->wait().is()); + release.count_down(); + pool->shutdown(); + EXPECT_FALSE(token->submit_func([] { return Status::OK(); }).ok()); +} + +TEST(LoadThreadPoolTest, BitmapSubmissionFailureSurvivesWait) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("bitmap_capacity_test") + .set_max_threads(1) + .set_max_queue_size(1) + .build(&pool) + .ok()); + CalcDeleteBitmapToken token(pool->new_load_token(1, LoadTaskPriority::MID, LoadTaskType::LEAF)); + SCOPED_ATTACH_TASK(MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "bitmap_capacity_test")); + CountDownLatch entered(1), release(1); + std::atomic completed = 0; + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(pool->submit_func([&] { + entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(entered.wait_for(5s)); + EXPECT_TRUE(token.submit_func([&] { + ++completed; + return Status::OK(); + }).ok()); + auto rejected = token.submit_func([] { + ADD_FAILURE() << "rejected bitmap callback ran"; + return Status::OK(); + }); + EXPECT_TRUE(rejected.is()); + EXPECT_NE(rejected.to_string().find("at capacity"), std::string::npos); + release.count_down(); + EXPECT_EQ(token.wait().to_string(), rejected.to_string()); + EXPECT_EQ(completed.load(), 1); + // The token remains failed even after the queue has drained. + EXPECT_EQ(token.submit_func([] { return Status::OK(); }).to_string(), rejected.to_string()); +} + +TEST(LoadThreadPoolTest, BitmapSubmissionAfterShutdownPreservesReason) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("bitmap_rejected_test").set_max_threads(1).build(&pool).ok()); + CalcDeleteBitmapToken token( + pool->new_load_token(1, LoadTaskPriority::HIGHEST, LoadTaskType::PARENT)); + SCOPED_ATTACH_TASK(MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "bitmap_rejected_test")); + pool->shutdown(); + auto rejected = token.submit_func([] { + ADD_FAILURE() << "shutdown pool accepted bitmap callback"; + return Status::OK(); + }); + EXPECT_TRUE(rejected.is()); + EXPECT_NE(rejected.to_string().find("shut down"), std::string::npos); + EXPECT_EQ(token.wait().to_string(), rejected.to_string()); +} + +TEST(LoadThreadPoolTest, LoadCleanupCanJoinRunningBitmapLeaves) { + for (auto priority : {LoadTaskPriority::MID, LoadTaskPriority::HIGHEST}) { + std::unique_ptr pool; + ASSERT_TRUE(ThreadPoolBuilder("load_cleanup_test").set_max_threads(2).build(&pool).ok()); + auto leaf = pool->new_load_token(1, priority, LoadTaskType::LEAF); + auto parent = pool->new_load_token(1, LoadTaskPriority::LOW, LoadTaskType::PARENT); + CountDownLatch leaf_entered(1), parent_entered(1), release(1); + Defer unblock = [&] { release.count_down(); }; + EXPECT_TRUE(leaf->submit_func([&] { + leaf_entered.count_down(); + release.wait(); + }).ok()); + EXPECT_TRUE(leaf_entered.wait_for(5s)); + EXPECT_TRUE(parent->submit_func([&] { + parent_entered.count_down(); + leaf->shutdown(); + }).ok()); + EXPECT_TRUE(parent_entered.wait_for(5s)); + release.count_down(); + parent->wait(); + EXPECT_FALSE(leaf->submit_func([] {}).ok()); + } +} + +} // namespace doris diff --git a/build-support/check-header-deps.py b/build-support/check-header-deps.py index 109f4b4d8d5c5f..463b1b03acd683 100755 --- a/build-support/check-header-deps.py +++ b/build-support/check-header-deps.py @@ -528,7 +528,7 @@ "gen_cpp/data.pb.h": 220, "io/fs/s3_file_system.h": 109, # the AWS SDK surface "util/brpc_closure.h": 61, - "runtime/workload_group/workload_group.h": 42, # thrift type universe carrier + "runtime/workload_group/workload_group.h": 44, # thrift carrier; +2 local publish test TUs } REVERSE_SLACK = 0.10