[workloadgroup](memory) flush memtable when memory is not enough (apache#54642)

### What problem does this PR solve?
1. flush memtable when memory is not enough
2. cancel the query that use more memory than min_memory_percent

Author: yiguolei

be/src/common/config.cpp

@@ -1391,6 +1391,8 @@ DEFINE_Int32(spill_io_thread_pool_queue_size, "102400");
 // paused query in queue timeout(ms) will be resumed or canceled
 DEFINE_Int64(spill_in_paused_queue_timeout_ms, "60000");
 
+DEFINE_Int64(wait_cancel_release_memory_ms, "5000");
+
 DEFINE_mBool(check_segment_when_build_rowset_meta, "false");
 
 DEFINE_mBool(force_azure_blob_global_endpoint, "false");

be/src/common/config.h

@@ -1447,6 +1447,7 @@ DECLARE_mInt32(spill_gc_work_time_ms);
 DECLARE_Int32(spill_io_thread_pool_thread_num);
 DECLARE_Int32(spill_io_thread_pool_queue_size);
 DECLARE_Int64(spill_in_paused_queue_timeout_ms);
+DECLARE_Int64(wait_cancel_release_memory_ms);
 
 DECLARE_mBool(check_segment_when_build_rowset_meta);
 

be/src/common/daemon.cpp

@@ -348,7 +348,7 @@ void Daemon::memory_maintenance_thread() {
 
         // step 6. Refresh weighted memory ratio of workload groups.
         doris::ExecEnv::GetInstance()->workload_group_mgr()->do_sweep();
-        doris::ExecEnv::GetInstance()->workload_group_mgr()->refresh_wg_weighted_memory_limit();
+        doris::ExecEnv::GetInstance()->workload_group_mgr()->refresh_workload_group_memory_state();
 
         // step 7: handle paused queries(caused by memory insufficient)
         doris::ExecEnv::GetInstance()->workload_group_mgr()->handle_paused_queries();

be/src/exec/schema_scanner/schema_workload_group_resource_usage_scanner.cpp

@@ -38,7 +38,6 @@ std::vector<SchemaScanner::ColumnDesc> SchemaBackendWorkloadGroupResourceUsage::
         {"CPU_USAGE_PERCENT", TYPE_DOUBLE, sizeof(double), false},
         {"LOCAL_SCAN_BYTES_PER_SECOND", TYPE_BIGINT, sizeof(int64_t), false},
         {"REMOTE_SCAN_BYTES_PER_SECOND", TYPE_BIGINT, sizeof(int64_t), false},
-        {"WRITE_BUFFER_USAGE_BYTES", TYPE_BIGINT, sizeof(int64_t), false},
 };
 
 SchemaBackendWorkloadGroupResourceUsage::SchemaBackendWorkloadGroupResourceUsage()

be/src/olap/memtable_memory_limiter.cpp

@@ -119,43 +119,11 @@ int64_t MemTableMemoryLimiter::_need_flush() {
     return need_flush - _queue_mem_usage - _flush_mem_usage;
 }
 
-void MemTableMemoryLimiter::handle_workload_group_memtable_flush(
-        WorkloadGroupPtr wg, std::function<bool()> cancel_check) {
-    // It means some query is pending on here to flush memtable and to continue running.
-    // So that should wait here.
-    // Wait at most 3s, because this code is not aware cancel flag. If the load task is cancelled
-    // Should releae memory quickly.
-    using namespace std::chrono_literals;
-    int32_t max_sleep_times = 30;
-    int32_t sleep_times = max_sleep_times;
-    MonotonicStopWatch timer;
-    timer.start();
-    while (wg != nullptr && wg->enable_write_buffer_limit() && wg->exceed_write_buffer_limit() &&
-           sleep_times > 0) {
-        if (cancel_check && cancel_check()) {
-            LOG(INFO) << "cancelled when waiting for memtable flush, wg: "
-                      << (wg == nullptr ? "null" : wg->debug_string());
-            return;
-        }
-        std::this_thread::sleep_for(100ms);
-        --sleep_times;
-    }
-    if (sleep_times < max_sleep_times) {
-        timer.stop();
-        VLOG_DEBUG << "handle_workload_group_memtable_flush waited "
-                   << PrettyPrinter::print(timer.elapsed_time(), TUnit::TIME_NS)
-                   << ", wg: " << wg->debug_string();
-    }
-    // Check process memory again.
-    _handle_memtable_flush(wg, cancel_check);
-}
-
-void MemTableMemoryLimiter::_handle_memtable_flush(WorkloadGroupPtr wg,
-                                                   std::function<bool()> cancel_check) {
+void MemTableMemoryLimiter::handle_memtable_flush(std::function<bool()> cancel_check) {
     // Check the soft limit.
     DCHECK(_load_soft_mem_limit > 0);
     do {
-        DBUG_EXECUTE_IF("MemTableMemoryLimiter._handle_memtable_flush.limit_reached", {
+        DBUG_EXECUTE_IF("MemTableMemoryLimiter.handle_memtable_flush.limit_reached", {
             LOG(INFO) << "debug memtable limit reached";
             break;
         });
@@ -176,8 +144,7 @@ void MemTableMemoryLimiter::_handle_memtable_flush(WorkloadGroupPtr wg,
             }
         }
         if (cancel_check && cancel_check()) {
-            LOG(INFO) << "cancelled when waiting for memtable flush, wg: "
-                      << (wg == nullptr ? "null" : wg->debug_string());
+            LOG(INFO) << "cancelled when waiting for memtable flush";
             return;
         }
         first = false;
@@ -192,15 +159,14 @@ void MemTableMemoryLimiter::_handle_memtable_flush(WorkloadGroupPtr wg,
                       << ", active: " << PrettyPrinter::print_bytes(_active_mem_usage)
                       << ", queue: " << PrettyPrinter::print_bytes(_queue_mem_usage)
                       << ", flush: " << PrettyPrinter::print_bytes(_flush_mem_usage)
-                      << ", need flush: " << PrettyPrinter::print_bytes(need_flush)
-                      << ", wg: " << (wg ? wg->debug_string() : "null");
+                      << ", need flush: " << PrettyPrinter::print_bytes(need_flush);
             if (VLOG_DEBUG_IS_ON) {
                 auto log_str = doris::ProcessProfile::instance()
                                        ->memory_profile()
                                        ->process_memory_detail_str();
                 LOG_LONG_STRING(INFO, log_str);
             }
-            _flush_active_memtables(0, need_flush);
+            _flush_active_memtables(need_flush);
         }
     } while (_hard_limit_reached() && !_load_usage_low());
     g_memtable_memory_limit_waiting_threads << -1;
@@ -216,37 +182,11 @@ void MemTableMemoryLimiter::_handle_memtable_flush(WorkloadGroupPtr wg,
                   << ", memtable writers num: " << _writers.size()
                   << ", active: " << PrettyPrinter::print_bytes(_active_mem_usage)
                   << ", queue: " << PrettyPrinter::print_bytes(_queue_mem_usage)
-                  << ", flush: " << PrettyPrinter::print_bytes(_flush_mem_usage)
-                  << ", wg: " << (wg ? wg->debug_string() : "null.");
-    }
-}
-
-int64_t MemTableMemoryLimiter::flush_workload_group_memtables(uint64_t wg_id, int64_t need_flush) {
-    std::unique_lock<std::mutex> l(_lock);
-    return _flush_active_memtables(wg_id, need_flush);
-}
-
-void MemTableMemoryLimiter::get_workload_group_memtable_usage(uint64_t wg_id, int64_t* active_bytes,
-                                                              int64_t* queue_bytes,
-                                                              int64_t* flush_bytes) {
-    std::unique_lock<std::mutex> l(_lock);
-    *active_bytes = 0;
-    *queue_bytes = 0;
-    *flush_bytes = 0;
-    for (auto it = _writers.begin(); it != _writers.end(); ++it) {
-        if (auto writer = it->lock()) {
-            // If wg id is specified, but wg id not match, then not need flush
-            if (writer->workload_group_id() != wg_id) {
-                continue;
-            }
-            *active_bytes += writer->active_memtable_mem_consumption();
-            *queue_bytes += writer->mem_consumption(MemType::WRITE_FINISHED);
-            *flush_bytes += writer->mem_consumption(MemType::FLUSH);
-        }
+                  << ", flush: " << PrettyPrinter::print_bytes(_flush_mem_usage);
     }
 }
 
-int64_t MemTableMemoryLimiter::_flush_active_memtables(uint64_t wg_id, int64_t need_flush) {
+int64_t MemTableMemoryLimiter::_flush_active_memtables(int64_t need_flush) {
     if (need_flush <= 0) {
         return 0;
     }
@@ -278,10 +218,7 @@ int64_t MemTableMemoryLimiter::_flush_active_memtables(uint64_t wg_id, int64_t n
         if (w == nullptr) {
             continue;
         }
-        // If wg id is specified, but wg id not match, then not need flush
-        if (wg_id != 0 && w->workload_group_id() != wg_id) {
-            continue;
-        }
+
         int64_t mem = w->active_memtable_mem_consumption();
         if (mem < sort_mem * 0.9) {
             // if the memtable writer just got flushed, don't flush it again

be/src/olap/memtable_memory_limiter.h

@@ -40,13 +40,11 @@ class MemTableMemoryLimiter {
 
     Status init(int64_t process_mem_limit);
 
-    void handle_workload_group_memtable_flush(WorkloadGroupPtr wg,
-                                              std::function<bool()> cancel_check);
-
-    int64_t flush_workload_group_memtables(uint64_t wg_id, int64_t need_flush_bytes);
-
-    void get_workload_group_memtable_usage(uint64_t wg_id, int64_t* active_bytes,
-                                           int64_t* queue_bytes, int64_t* flush_bytes);
+    // check if the total mem consumption exceeds limit.
+    // If yes, it will flush memtable to try to reduce memory consumption.
+    // Every write operation will call this API to check if need flush memtable OR hang
+    // when memory is not available.
+    void handle_memtable_flush(std::function<bool()> cancel_check);
 
     void register_writer(std::weak_ptr<MemTableWriter> writer);
 
@@ -57,22 +55,15 @@ class MemTableMemoryLimiter {
     int64_t mem_usage() const { return _mem_usage; }
 
 private:
-    // check if the total mem consumption exceeds limit.
-    // If yes, it will flush memtable to try to reduce memory consumption.
-    // Every write operation will call this API to check if need flush memtable OR hang
-    // when memory is not available.
-    void _handle_memtable_flush(WorkloadGroupPtr wg, std::function<bool()> cancel_check);
-
     static inline int64_t _sys_avail_mem_less_than_warning_water_mark();
     static inline int64_t _process_used_mem_more_than_soft_mem_limit();
 
     bool _soft_limit_reached();
     bool _hard_limit_reached();
     bool _load_usage_low();
     int64_t _need_flush();
-    int64_t _flush_active_memtables(uint64_t wg_id, int64_t need_flush);
+    int64_t _flush_active_memtables(int64_t need_flush);
     void _refresh_mem_tracker();
-
     std::mutex _lock;
     std::condition_variable _hard_limit_end_cond;
     int64_t _mem_usage = 0;

be/src/runtime/load_channel_mgr.cpp

@@ -152,8 +152,8 @@ Status LoadChannelMgr::add_batch(const PTabletWriterAddBlockRequest& request,
         // If this is a high priority load task, do not handle this.
         // because this may block for a while, which may lead to rpc timeout.
         SCOPED_TIMER(channel->get_handle_mem_limit_timer());
-        ExecEnv::GetInstance()->memtable_memory_limiter()->handle_workload_group_memtable_flush(
-                channel->workload_group(), [channel]() { return channel->is_cancelled(); });
+        ExecEnv::GetInstance()->memtable_memory_limiter()->handle_memtable_flush(
+                [channel]() { return channel->is_cancelled(); });
         if (channel->is_cancelled()) {
             return Status::Cancelled("LoadChannel has been cancelled: {}.", load_id.to_string());
         }

be/src/runtime/memory/global_memory_arbitrator.cpp

@@ -62,7 +62,6 @@ std::atomic<double> GlobalMemoryArbitrator::last_periodic_refreshed_cache_capaci
 std::atomic<double> GlobalMemoryArbitrator::last_memory_exceeded_cache_capacity_adjust_weighted {1};
 // The value that take affect
 std::atomic<double> GlobalMemoryArbitrator::last_affected_cache_capacity_adjust_weighted {1};
-std::atomic<bool> GlobalMemoryArbitrator::any_workload_group_exceed_limit {false};
 std::mutex GlobalMemoryArbitrator::memtable_memory_refresh_lock;
 std::condition_variable GlobalMemoryArbitrator::memtable_memory_refresh_cv;
 std::atomic<bool> GlobalMemoryArbitrator::memtable_memory_refresh_notify {false};

be/src/runtime/memory/global_memory_arbitrator.h

@@ -166,7 +166,6 @@ class GlobalMemoryArbitrator {
     static std::atomic<double> last_memory_exceeded_cache_capacity_adjust_weighted;
     // The value that take affect
     static std::atomic<double> last_affected_cache_capacity_adjust_weighted;
-    static std::atomic<bool> any_workload_group_exceed_limit;
 
     static void notify_cache_adjust_capacity() {
         cache_adjust_capacity_notify.store(true, std::memory_order_relaxed);

be/src/runtime/memory/mem_tracker_limiter.h

@@ -220,8 +220,6 @@ class MemTrackerLimiter final {
     static void make_top_consumption_tasks_tracker_profile(RuntimeProfile* profile, int top_num);
     static void make_all_tasks_tracker_profile(RuntimeProfile* profile);
 
-    int64_t write_buffer_size() const { return _write_tracker->consumption(); }
-
     std::shared_ptr<MemTrackerLimiter> write_tracker() { return _write_tracker; }
 
     void print_log_usage(const std::string& msg);