This is an automated email from the ASF dual-hosted git repository.

yiguolei pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/doris.git


The following commit(s) were added to refs/heads/branch-4.1 by this push:
     new c2091067fd7 branch-4.1: [fix](be) Keep paused queries waiting under 
process memory pressure before cancelling #67845 (#67846)
c2091067fd7 is described below

commit c2091067fd78068ff78918882df50596ffba3108
Author: Jerry Hu <[email protected]>
AuthorDate: Wed Nov 11 15:33:08 2026 +0800

    branch-4.1: [fix](be) Keep paused queries waiting under process memory 
pressure before cancelling #67845 (#67846)
    
    ### What problem does this PR solve?
    
    Cherry-picked from #67845
    
    Issue Number: None
    
    Related PR: #67845
    
    Problem Summary:
    
    Backport of #67845 to branch-4.1. The first commit applies cleanly.
    branch-4.1 does not have the paused-query helpers split out by #62485,
    so the prompt-recovery check of the second commit is placed in the
    `PROCESS_MEMORY_EXCEEDED` branch of `handle_paused_queries()` instead of
    `handle_process_memory_exceeded_()`. branch-4.1 also lacks the timeout
    fallback that #62485 added for queries whose workload group uses no more
    than its min memory limit, so that fallback is ported into
    `handle_paused_queries()` as well: once such a query has waited
    `spill_in_paused_queue_timeout_ms` and memory cannot be revoked from any
    other workload group, it is handed to `handle_single_query_()` and
    cancelled like on master. The rest is identical.
    
    When a query is paused because of `PROCESS_MEMORY_EXCEEDED` and it has
    no revocable memory, the query was cancelled about 50-100ms later if the
    process was still above the soft memory limit, even when cancelling it
    did not relieve the process memory pressure. This PR keeps such queries
    paused until they have waited `spill_in_paused_queue_timeout_ms`,
    re-checking the soft memory limit every round so they resume once other
    queries release memory.
    
    ### Release note
    
    Queries paused under process memory pressure without revocable memory
    now wait up to `spill_in_paused_queue_timeout_ms` for memory to be
    released before being cancelled.
    
    ### Check List (For Author)
    
    - Test
        - [x] Unit Test
    - Behavior changed:
        - [x] Yes. Same as #67845.
    - Does this need documentation?
        - [x] No.
    
    https://claude.ai/code/session_014c64WZPFui9AyvVMiseceY
---
 .../workload_group/workload_group_manager.cpp      | 140 +++++----
 .../workload_group/workload_group_manager.h        |  12 +-
 be/src/util/mem_info.h                             |   7 +
 .../workload_group/workload_group_manager_test.cpp | 319 ++++++++++++++++++++-
 4 files changed, 422 insertions(+), 56 deletions(-)

diff --git a/be/src/runtime/workload_group/workload_group_manager.cpp 
b/be/src/runtime/workload_group/workload_group_manager.cpp
index 9c50c3c5fb6..25c4fc92de8 100644
--- a/be/src/runtime/workload_group/workload_group_manager.cpp
+++ b/be/src/runtime/workload_group/workload_group_manager.cpp
@@ -52,7 +52,7 @@ PausedQuery::PausedQuery(std::shared_ptr<ResourceContext> 
resource_ctx, double c
           cache_ratio_(cache_ratio),
           reserve_size_(reserve_size),
           query_id_(print_id(resource_ctx->task_controller()->task_id())) {
-    enqueue_at = std::chrono::system_clock::now();
+    enqueue_at = std::chrono::steady_clock::now();
 }
 
 WorkloadGroupMgr::~WorkloadGroupMgr() = default;
@@ -585,9 +585,51 @@ void WorkloadGroupMgr::handle_paused_queries() {
                     ++query_it;
                     continue;
                 }
+
+                // The process memory pressure may have been relieved by cache 
reclamation or by
+                // other queries that finished. Check it before routing the 
query below, otherwise
+                // a query in a workload group that uses less than its min 
memory limit has to
+                // wait for the timeout.
+                const size_t test_memory_size =
+                        std::max<size_t>(query_it->reserve_size_, 32L * 1024 * 
1024);
+                if 
(!GlobalMemoryArbitrator::is_exceed_soft_mem_limit(test_memory_size)) {
+                    LOG(INFO) << "Query: " << 
print_id(resource_ctx->task_controller()->task_id())
+                              << ", process limit not exceeded now, resume 
this query"
+                              << ", process memory info: "
+                              << 
GlobalMemoryArbitrator::process_memory_used_details_str()
+                              << ", wg info: " << wg->debug_string();
+                    
resource_ctx->task_controller()->set_memory_sufficient(true);
+                    query_it = queries_list.erase(query_it);
+                    continue;
+                }
+
                 // If workload group's memory usage > min memory, then it 
means the workload group use too much memory
                 // in memory contention state. Should just spill
-                if (wg->total_mem_used() > wg->min_memory_limit()) {
+                bool handle_query_now = wg->total_mem_used() > 
wg->min_memory_limit();
+                if (!handle_query_now) {
+                    // Other workload groups many use a lot of memory, should 
revoke memory from other workload groups
+                    // by cancelling their queries.
+                    int64_t revoked_size = revoke_memory_from_other_groups_();
+                    if (revoked_size > 0) {
+                        // Revoke memory from other workload groups will 
cancel some queries, wait them cancel finished
+                        // and then check it again.
+                        revoking_memory_from_other_query_ = true;
+                        return;
+                    }
+
+                    // TODO revoke from memtable
+
+                    // Fallback: no other workload group can release memory. 
Once the process
+                    // reaches the hard memory limit or the query has waited 
too long, let
+                    // handle_single_query_ spill or cancel it to protect the 
process. The hard
+                    // limit is checked here as well, because the query would 
otherwise wait for
+                    // the timeout at the hard limit, and memory gc may be 
disabled.
+                    handle_query_now =
+                            GlobalMemoryArbitrator::is_exceed_hard_mem_limit() 
||
+                            query_it->elapsed_time() > 
config::spill_in_paused_queue_timeout_ms;
+                }
+
+                if (handle_query_now) {
                     bool spill_res = handle_single_query_(
                             resource_ctx, query_it->reserve_size_, 
query_it->elapsed_time(),
                             resource_ctx->task_controller()->paused_reason());
@@ -608,18 +650,6 @@ void WorkloadGroupMgr::handle_paused_queries() {
                     }
                 }
 
-                // Other workload groups many use a lot of memory, should 
revoke memory from other workload groups
-                // by cancelling their queries.
-                int64_t revoked_size = revoke_memory_from_other_groups_();
-                if (revoked_size > 0) {
-                    // Revoke memory from other workload groups will cancel 
some queries, wait them cancel finished
-                    // and then check it again.
-                    revoking_memory_from_other_query_ = true;
-                    return;
-                }
-
-                // TODO revoke from memtable
-
                 ++query_it;
             }
         }
@@ -641,6 +671,9 @@ void WorkloadGroupMgr::handle_paused_queries() {
 // 2. revoke 10% memory of the workload group that exceeded. For example, if 
the workload group exceed 10g,
 //    then revoke 1g memory.
 // 3. After revoke memory, go to the loop and wait for the query to be 
cancelled and check again.
+// Returns the memory that is actually being released by the cancelled 
queries, which is 0 when
+// no query of that workload group could be cancelled (for example, every 
query is too small to
+// be revoked), so that the caller does not wait for memory that will never be 
freed.
 int64_t WorkloadGroupMgr::revoke_memory_from_other_groups_() {
     MonotonicStopWatch watch;
     watch.start();
@@ -675,16 +708,16 @@ int64_t 
WorkloadGroupMgr::revoke_memory_from_other_groups_() {
                   << " less than 128MB, no need to revoke memory";
         return 0;
     }
-    int64_t freed_mem = static_cast<int64_t>((double)max_exceeded_memory * 
0.1);
+    int64_t need_free_mem = static_cast<int64_t>((double)max_exceeded_memory * 
0.1);
     // Revoke 10% of memory from the workload group that exceed most memory
-    max_wg->revoke_memory(freed_mem, "exceed_memory", profile.get());
+    int64_t freed_mem = max_wg->revoke_memory(need_free_mem, "exceed_memory", 
profile.get());
     std::stringstream ss;
     profile->pretty_print(&ss);
     LOG(INFO) << fmt::format(
             "[MemoryGC] process memory not enough, revoke memory from 
workload_group: {}, "
-            "free memory {}. cost(us): {}, details: {}",
-            max_wg->memory_debug_string(), 
PrettyPrinter::print_bytes(freed_mem),
-            watch.elapsed_time() / 1000, ss.str());
+            "need free memory {}, freed memory {}. cost(us): {}, details: {}",
+            max_wg->memory_debug_string(), 
PrettyPrinter::print_bytes(need_free_mem),
+            PrettyPrinter::print_bytes(freed_mem), watch.elapsed_time() / 
1000, ss.str());
     return freed_mem;
 }
 
@@ -796,38 +829,45 @@ bool WorkloadGroupMgr::handle_single_query_(const 
std::shared_ptr<ResourceContex
             return true;
         }
     } else {
-        // Should not consider about process memory. For example, the query's 
limit is 100g, workload
-        // group's memlimit is 10g, process memory is 20g. The query reserve 
will always failed in wg
-        // limit, and process is always have memory, so that it will resume 
and failed reserve again.
-        const size_t test_memory_size = std::max<size_t>(size_to_reserve, 32L 
* 1024 * 1024);
-        if 
(!GlobalMemoryArbitrator::is_exceed_soft_mem_limit(test_memory_size)) {
-            LOG(INFO) << "Query: " << query_id
-                      << ", process limit not exceeded now, resume this query"
-                      << ", process memory info: "
-                      << 
GlobalMemoryArbitrator::process_memory_used_details_str()
-                      << ", wg info: " << wg->debug_string();
-            requestor->task_controller()->set_memory_sufficient(true);
-            return true;
-        } else {
-            // if cannot find any memory to release, then let the query 
continue to run as far as possible
-            // or cancelled by gc if memory is really not enough.
-            Status error_status = Status::MemoryLimitExceeded(
-                    "Query {} process memory is exceeded"
-                    ", and there is no cache now. And could not find task to 
spill, disable "
-                    "reserve memory and resume it. "
-                    "Query memory usage: {}, limit: {}, reserved "
-                    "size: {}, try to reserve: {}, wg info: {}."
-                    " Maybe you should set the workload group's limit to a 
lower value. {}",
-                    query_id, PrettyPrinter::print_bytes(memory_usage),
-                    PrettyPrinter::print_bytes(limit), 
PrettyPrinter::print_bytes(reserved_size),
-                    PrettyPrinter::print_bytes(size_to_reserve), 
wg->memory_debug_string(),
-                    doris::ProcessProfile::instance()
-                            ->memory_profile()
-                            ->process_memory_detail_str());
-            LOG_LONG_STRING(INFO, error_status.to_string());
-            requestor->task_controller()->cancel(error_status);
-            return true;
+        // PROCESS_MEMORY_EXCEEDED. The caller (handle_paused_queries) has 
already resumed
+        // the query if the process is no longer above the soft memory limit, 
so the process
+        // memory is still exceeded here.
+        const bool exceed_hard_mem_limit = 
GlobalMemoryArbitrator::is_exceed_hard_mem_limit();
+        if (time_in_queue < config::spill_in_paused_queue_timeout_ms && 
!exceed_hard_mem_limit) {
+            // The query has no revocable memory, cancelling it will not 
release much memory.
+            // Keep it paused so that it can be resumed once other queries 
release memory,
+            // and cancel it only after it has waited for 
`spill_in_paused_queue_timeout_ms`.
+            // Stop waiting once the process reaches the hard memory limit, so 
that the
+            // protection does not depend on memory gc, which may be disabled.
+            LOG_EVERY_T(INFO, 1) << "Query: " << query_id
+                                 << " process memory is exceeded, and could 
not find task to "
+                                    "spill, keep it paused. Waited "
+                                 << time_in_queue
+                                 << " ms, timeout: " << 
config::spill_in_paused_queue_timeout_ms
+                                 << " ms, process memory info: "
+                                 << 
GlobalMemoryArbitrator::process_memory_used_details_str()
+                                 << ", wg info: " << wg->debug_string();
+            return false;
         }
+
+        // Waited long enough or the process reached the hard memory limit, 
and still
+        // could not find any memory to release, cancel the query to protect 
the process.
+        Status error_status = Status::MemoryLimitExceeded(
+                "Query {} process memory is exceeded"
+                ", and there is no cache now. And could not find task to spill 
after "
+                "waiting {} ms in paused queue (timeout: {} ms, exceed hard 
limit: {}), "
+                "try to cancel query. "
+                "Query memory usage: {}, limit: {}, reserved "
+                "size: {}, try to reserve: {}, wg info: {}."
+                " Maybe you should set the workload group's limit to a lower 
value. {}",
+                query_id, time_in_queue, 
config::spill_in_paused_queue_timeout_ms,
+                exceed_hard_mem_limit, 
PrettyPrinter::print_bytes(memory_usage),
+                PrettyPrinter::print_bytes(limit), 
PrettyPrinter::print_bytes(reserved_size),
+                PrettyPrinter::print_bytes(size_to_reserve), 
wg->memory_debug_string(),
+                
doris::ProcessProfile::instance()->memory_profile()->process_memory_detail_str());
+        LOG_LONG_STRING(INFO, error_status.to_string());
+        requestor->task_controller()->cancel(error_status);
+        return true;
     }
 }
 
diff --git a/be/src/runtime/workload_group/workload_group_manager.h 
b/be/src/runtime/workload_group/workload_group_manager.h
index 5b2b1160d60..c1e85267969 100644
--- a/be/src/runtime/workload_group/workload_group_manager.h
+++ b/be/src/runtime/workload_group/workload_group_manager.h
@@ -18,6 +18,7 @@
 
 #include <stdint.h>
 
+#include <chrono>
 #include <shared_mutex>
 #include <unordered_map>
 
@@ -39,7 +40,7 @@ public:
     // Use weak ptr to save resource ctx, to make sure if the query is 
cancelled
     // the resource will be released
     std::weak_ptr<ResourceContext> resource_ctx_;
-    std::chrono::system_clock::time_point enqueue_at;
+    std::chrono::steady_clock::time_point enqueue_at;
     size_t last_mem_usage {0};
     double cache_ratio_ {0.0};
     int64_t reserve_size_ {0};
@@ -48,7 +49,7 @@ public:
                 int64_t reserve_size);
 
     int64_t elapsed_time() const {
-        auto now = std::chrono::system_clock::now();
+        auto now = std::chrono::steady_clock::now();
         return std::chrono::duration_cast<std::chrono::milliseconds>(now - 
enqueue_at).count();
     }
 
@@ -102,6 +103,13 @@ private:
 
     WorkloadGroupPtr get_or_create_workload_group(const WorkloadGroupInfo& 
workload_group_info);
 
+    // Attempt to resolve a single paused query: spill if it has revocable 
memory, resume
+    // if it is under its limit, otherwise cancel it. For 
PROCESS_MEMORY_EXCEEDED with no
+    // revocable memory, keep the query paused until it has waited
+    // spill_in_paused_queue_timeout_ms or the process reaches the hard memory 
limit, then
+    // cancel it. Returns true if the query was acted upon 
(spilled/cancelled/resumed),
+    // false if it should keep waiting (still has running tasks, or is within 
the
+    // process-memory grace period).
     bool handle_single_query_(const std::shared_ptr<ResourceContext>& 
requestor,
                               size_t size_to_reserve, int64_t time_in_queue, 
Status paused_reason);
     int64_t revoke_memory_from_other_groups_();
diff --git a/be/src/util/mem_info.h b/be/src/util/mem_info.h
index 113b3352c25..3a3ef3cecc0 100644
--- a/be/src/util/mem_info.h
+++ b/be/src/util/mem_info.h
@@ -87,6 +87,13 @@ public:
     static void set_mem_limit_for_test(int64_t mem_limit) {
         _s_mem_limit.store(mem_limit, std::memory_order_relaxed);
     }
+    static void set_soft_mem_limit_for_test(int64_t soft_mem_limit) {
+        _s_soft_mem_limit.store(soft_mem_limit, std::memory_order_relaxed);
+    }
+    // Returns the previous value so that the test can restore it.
+    static int64_t set_sys_mem_available_for_test(int64_t sys_mem_available) {
+        return _s_sys_mem_available.exchange(sys_mem_available, 
std::memory_order_relaxed);
+    }
 #endif
     static inline std::string mem_limit_str() {
         DCHECK(_s_initialized);
diff --git a/be/test/runtime/workload_group/workload_group_manager_test.cpp 
b/be/test/runtime/workload_group/workload_group_manager_test.cpp
index 8a10c619bfd..e1f085cff25 100644
--- a/be/test/runtime/workload_group/workload_group_manager_test.cpp
+++ b/be/test/runtime/workload_group/workload_group_manager_test.cpp
@@ -27,9 +27,12 @@
 #include <chrono>
 #include <cmath>
 #include <filesystem>
+#include <limits>
 #include <memory>
 #include <sstream>
+#include <string>
 #include <tuple>
+#include <vector>
 
 #include "common/config.h"
 #include "common/status.h"
@@ -155,6 +158,26 @@ private:
         }
     }
 
+    size_t _paused_queries_count(const WorkloadGroupPtr& wg) {
+        std::unique_lock<std::mutex> lock(_wg_manager->_paused_queries_lock);
+        if (!_wg_manager->_paused_queries_list.contains(wg)) {
+            return 0;
+        }
+        return _wg_manager->_paused_queries_list[wg].size();
+    }
+
+    // Let the only paused query of `wg` exceed 
`spill_in_paused_queue_timeout_ms` by backdating
+    // its enqueue time instead of sleeping, the timer is monotonic.
+    void _expire_paused_query(const WorkloadGroupPtr& wg) {
+        std::unique_lock<std::mutex> lock(_wg_manager->_paused_queries_lock);
+        auto& queries = _wg_manager->_paused_queries_list[wg];
+        ASSERT_EQ(queries.size(), 1);
+        auto node = queries.extract(queries.begin());
+        node.value().enqueue_at -=
+                
std::chrono::milliseconds(config::spill_in_paused_queue_timeout_ms + 1);
+        queries.insert(std::move(node));
+    }
+
     std::unique_ptr<WorkloadGroupMgr> _wg_manager;
     std::string _test_dir;
     const int64_t _spill_in_paused_queue_timeout_ms = 
config::spill_in_paused_queue_timeout_ms;
@@ -665,7 +688,9 @@ TEST_F(WorkloadGroupManagerTest, ProcessMemoryNotEnough) {
     wg2->refresh_memory_usage();
     EXPECT_EQ(wg2->total_mem_used(), 1024 * 1024 * 410);
     EXPECT_EQ(wg2->min_memory_limit(), 1024 * 1024 * 100);
-    EXPECT_EQ(31 * 1024 * 1024, 
_wg_manager->revoke_memory_from_other_groups_());
+    // WG2 exceeds 310MB, so 31MB should be revoked. The memory of the 
cancelled query23 is
+    // what is actually being released.
+    EXPECT_EQ(300 * 1024 * 1024, 
_wg_manager->revoke_memory_from_other_groups_());
     ASSERT_FALSE(query_context21->is_cancelled());
     ASSERT_FALSE(query_context22->is_cancelled());
     ASSERT_TRUE(query_context23->is_cancelled());
@@ -691,15 +716,17 @@ TEST_F(WorkloadGroupManagerTest, ProcessMemoryNotEnough) {
     wg3->refresh_memory_usage();
     EXPECT_EQ(wg3->total_mem_used(), 1024 * 1024 * 500); // WG3 exceed 400MB
 
-    EXPECT_EQ(40 * 1024 * 1024, 
_wg_manager->revoke_memory_from_other_groups_());
+    // WG3 exceeds most, 40MB should be revoked, and query31 (500MB) is 
cancelled for it.
+    EXPECT_EQ(500 * 1024 * 1024, 
_wg_manager->revoke_memory_from_other_groups_());
 
     wg1->refresh_memory_usage();
     wg2->refresh_memory_usage();
     wg3->refresh_memory_usage();
 
     ASSERT_TRUE(query_context31->is_cancelled());
-    // query31 is still in wg3, so that it is not cancel again.
-    EXPECT_EQ(40 * 1024 * 1024, 
_wg_manager->revoke_memory_from_other_groups_());
+    // query31 is still in wg3 and was cancelled recently, so that it is not 
cancelled again,
+    // but its memory is still counted as being released.
+    EXPECT_EQ(500 * 1024 * 1024, 
_wg_manager->revoke_memory_from_other_groups_());
     ASSERT_FALSE(query_context11->is_cancelled());
     ASSERT_FALSE(query_context21->is_cancelled());
     ASSERT_FALSE(query_context22->is_cancelled());
@@ -1360,4 +1387,288 @@ TEST_F(WorkloadGroupManagerTest, 
AdaptiveFlushRegistrationSurvivesIdChangeAndReu
     controller->adjust_once();
 }
 
+// Helpers for the process memory exceeded cases. The workload group memory 
limit is derived from
+// the process memory limit, so the process memory limit is set before the 
workload group is
+// created, and only the soft limit is lowered afterwards to simulate process 
memory pressure.
+static constexpr int64_t kProcessMemLimitForWg = 1024L * 1024 * 1000;
+static constexpr int64_t kLargeProcessMemLimit = 
std::numeric_limits<int64_t>::max() / 2;
+static constexpr int64_t kWgMinMemLimit = 1024L * 1024 * 100;
+
+// GlobalMemoryArbitrator::is_exceed_soft_mem_limit() / 
is_exceed_hard_mem_limit() also compare
+// the system available memory, which MemInfo::init() reads from the host, 
against the water
+// marks. Pin it to a large value while a case runs, so that the expected 
results do not depend on
+// the memory state of the machine that runs the test, and restore all limits 
afterwards.
+class ScopedProcessMemLimits {
+public:
+    ScopedProcessMemLimits()
+            : _mem_limit(MemInfo::mem_limit()),
+              _soft_mem_limit(MemInfo::soft_mem_limit()),
+              
_sys_mem_available(MemInfo::set_sys_mem_available_for_test(kLargeProcessMemLimit))
 {
+        MemInfo::set_mem_limit_for_test(kProcessMemLimitForWg);
+    }
+    ~ScopedProcessMemLimits() {
+        MemInfo::set_mem_limit_for_test(_mem_limit);
+        MemInfo::set_soft_mem_limit_for_test(_soft_mem_limit);
+        MemInfo::set_sys_mem_available_for_test(_sys_mem_available);
+    }
+
+    // Process soft memory limit is exceeded, hard memory limit is not.
+    static void exceed_soft_limit() {
+        MemInfo::set_mem_limit_for_test(kLargeProcessMemLimit);
+        MemInfo::set_soft_mem_limit_for_test(1);
+    }
+    // Both the soft and the hard process memory limits are exceeded.
+    static void exceed_hard_limit() {
+        MemInfo::set_mem_limit_for_test(0);
+        MemInfo::set_soft_mem_limit_for_test(0);
+    }
+    // The process memory pressure is relieved.
+    static void relieve() {
+        MemInfo::set_mem_limit_for_test(kLargeProcessMemLimit);
+        MemInfo::set_soft_mem_limit_for_test(kLargeProcessMemLimit);
+    }
+
+private:
+    const int64_t _mem_limit;
+    const int64_t _soft_mem_limit;
+    const int64_t _sys_mem_available;
+};
+
+// Creates a workload group whose min memory limit is 100MB (10% of 1000MB).
+static WorkloadGroupPtr create_process_pressure_wg(WorkloadGroupMgr* 
wg_manager, uint64_t id = 1) {
+    WorkloadGroupInfo wg_info {.id = id,
+                               .memory_limit = kProcessMemLimitForWg,
+                               .min_memory_percent = 10,
+                               .max_memory_percent = 100};
+    auto wg = wg_manager->get_or_create_workload_group(wg_info);
+    EXPECT_EQ(wg->min_memory_limit(), kWgMinMemLimit);
+    return wg;
+}
+
+static void assert_cancelled_for_process_memory(const 
std::shared_ptr<QueryContext>& query,
+                                                bool exceed_hard_limit) {
+    ASSERT_TRUE(query->is_cancelled());
+    ASSERT_TRUE(query->exec_status().is<ErrorCode::MEM_LIMIT_EXCEEDED>())
+            << query->exec_status().to_string();
+    const std::string expected =
+            std::string("exceed hard limit: ") + (exceed_hard_limit ? "true" : 
"false");
+    ASSERT_NE(query->exec_status().to_string().find(expected), 
std::string::npos)
+            << query->exec_status().to_string();
+}
+
+// When the process memory is exceeded and the paused query has no revocable 
memory, the query
+// should be kept paused instead of being cancelled immediately, so that it 
can be resumed once
+// other queries release memory.
+TEST_F(WorkloadGroupManagerTest, 
process_mem_exceeded_keeps_paused_and_resumes) {
+    ScopedProcessMemLimits process_mem_limits;
+    auto wg = create_process_pressure_wg(_wg_manager.get());
+
+    auto query = _generate_on_query(wg);
+    // Let the workload group use more than its min memory limit, so that the 
paused query is
+    // handled by handle_single_query_ directly instead of waiting for other 
workload groups.
+    query->query_mem_tracker()->consume(1024L * 1024 * 128);
+    Defer release_memory {[&]() { query->query_mem_tracker()->consume(-1024L * 
1024 * 128); }};
+    wg->refresh_memory_usage();
+    ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+    ScopedProcessMemLimits::exceed_soft_limit();
+    _wg_manager->add_paused_query(query->resource_ctx(), 1024L,
+                                  
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+
+    config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+    for (int i = 0; i < 3; ++i) {
+        _wg_manager->handle_paused_queries();
+        ASSERT_FALSE(query->is_cancelled());
+        ASSERT_EQ(_paused_queries_count(wg), 1);
+    }
+    ASSERT_TRUE(query->resource_ctx()
+                        ->task_controller()
+                        ->paused_reason()
+                        .is<ErrorCode::PROCESS_MEMORY_EXCEEDED>());
+
+    // Process memory pressure is relieved, the query should be resumed.
+    ScopedProcessMemLimits::relieve();
+    _wg_manager->handle_paused_queries();
+    ASSERT_FALSE(query->is_cancelled());
+    
ASSERT_TRUE(query->resource_ctx()->task_controller()->paused_reason().ok());
+    ASSERT_EQ(_paused_queries_count(wg), 0);
+}
+
+// When the process memory is still exceeded after the paused query has waited 
for
+// `spill_in_paused_queue_timeout_ms`, the query should be cancelled.
+TEST_F(WorkloadGroupManagerTest, process_mem_exceeded_cancels_after_timeout) {
+    ScopedProcessMemLimits process_mem_limits;
+    auto wg = create_process_pressure_wg(_wg_manager.get());
+
+    auto query = _generate_on_query(wg);
+    query->query_mem_tracker()->consume(1024L * 1024 * 128);
+    Defer release_memory {[&]() { query->query_mem_tracker()->consume(-1024L * 
1024 * 128); }};
+    wg->refresh_memory_usage();
+    ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+    ScopedProcessMemLimits::exceed_soft_limit();
+    _wg_manager->add_paused_query(query->resource_ctx(), 1024L,
+                                  
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+
+    config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+    _wg_manager->handle_paused_queries();
+    ASSERT_FALSE(query->is_cancelled());
+    _expire_paused_query(wg);
+
+    _wg_manager->handle_paused_queries();
+    assert_cancelled_for_process_memory(query, false);
+}
+
+// A paused query whose workload group uses no more than its min memory limit 
is not handled by
+// handle_single_query_. It should still be resumed as soon as the process 
memory pressure is
+// relieved, instead of waiting for the timeout.
+TEST_F(WorkloadGroupManagerTest, 
process_mem_exceeded_below_min_memory_resumes) {
+    ScopedProcessMemLimits process_mem_limits;
+    auto wg = create_process_pressure_wg(_wg_manager.get());
+
+    auto query = _generate_on_query(wg);
+    query->query_mem_tracker()->consume(1024L * 4);
+    Defer release_memory {[&]() { query->query_mem_tracker()->consume(-1024L * 
4); }};
+    wg->refresh_memory_usage();
+    ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+    ScopedProcessMemLimits::exceed_soft_limit();
+    _wg_manager->add_paused_query(query->resource_ctx(), 1024L,
+                                  
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+
+    config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+    _wg_manager->handle_paused_queries();
+    ASSERT_FALSE(query->is_cancelled());
+    ASSERT_EQ(_paused_queries_count(wg), 1);
+
+    ScopedProcessMemLimits::relieve();
+    _wg_manager->handle_paused_queries();
+    ASSERT_FALSE(query->is_cancelled());
+    
ASSERT_TRUE(query->resource_ctx()->task_controller()->paused_reason().ok());
+    ASSERT_EQ(_paused_queries_count(wg), 0);
+}
+
+// A paused query whose workload group uses no more than its min memory limit 
should still be
+// cancelled once it has waited for `spill_in_paused_queue_timeout_ms` under 
process memory
+// pressure.
+TEST_F(WorkloadGroupManagerTest, 
process_mem_exceeded_below_min_memory_cancels_after_timeout) {
+    ScopedProcessMemLimits process_mem_limits;
+    auto wg = create_process_pressure_wg(_wg_manager.get());
+
+    auto query = _generate_on_query(wg);
+    query->query_mem_tracker()->consume(1024L * 4);
+    Defer release_memory {[&]() { query->query_mem_tracker()->consume(-1024L * 
4); }};
+    wg->refresh_memory_usage();
+    ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+    ScopedProcessMemLimits::exceed_soft_limit();
+    _wg_manager->add_paused_query(query->resource_ctx(), 1024L,
+                                  
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+
+    config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+    _wg_manager->handle_paused_queries();
+    ASSERT_FALSE(query->is_cancelled());
+    _expire_paused_query(wg);
+
+    _wg_manager->handle_paused_queries();
+    assert_cancelled_for_process_memory(query, false);
+}
+
+// A paused query whose workload group uses no more than its min memory limit 
should also be
+// cancelled immediately at the process hard limit when no other workload 
group can release
+// memory. This path must not rely on the memory GC daemon, which may be 
disabled.
+TEST_F(WorkloadGroupManagerTest, 
process_mem_exceeded_below_min_memory_cancels_at_hard_limit) {
+    ScopedProcessMemLimits process_mem_limits;
+    auto wg = create_process_pressure_wg(_wg_manager.get());
+
+    auto query = _generate_on_query(wg);
+    query->query_mem_tracker()->consume(1024L * 4);
+    Defer release_memory {[&]() { query->query_mem_tracker()->consume(-1024L * 
4); }};
+    wg->refresh_memory_usage();
+    ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+    ScopedProcessMemLimits::exceed_hard_limit();
+    _wg_manager->add_paused_query(query->resource_ctx(), 1024L,
+                                  
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+
+    config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+    _wg_manager->handle_paused_queries();
+    assert_cancelled_for_process_memory(query, true);
+}
+
+// Another workload group exceeds its min memory limit by more than 128MB, but 
all of its queries
+// are too small to be revoked, so revoking memory from it releases nothing. 
The paused query must
+// not wait for that memory forever: it is kept paused until the timeout and 
then cancelled, while
+// the queries of the other workload group are left untouched.
+TEST_F(WorkloadGroupManagerTest,
+       process_mem_exceeded_other_group_not_reclaimable_cancels_after_timeout) 
{
+    ScopedProcessMemLimits process_mem_limits;
+    auto wg = create_process_pressure_wg(_wg_manager.get(), 1);
+    auto other_wg = create_process_pressure_wg(_wg_manager.get(), 2);
+
+    auto query = _generate_on_query(wg);
+    query->query_mem_tracker()->consume(1024L * 4);
+    Defer release_memory {[&]() { query->query_mem_tracker()->consume(-1024L * 
4); }};
+    wg->refresh_memory_usage();
+    ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+    // 8 queries of 30MB each: the workload group exceeds its min memory limit 
by 140MB, but
+    // every query is smaller than SMALL_MEMORY_TASK (32MB) and is excluded 
from revocation.
+    constexpr int64_t kSmallQueryMem = 1024L * 1024 * 30;
+    std::vector<std::shared_ptr<QueryContext>> other_queries;
+    for (int i = 0; i < 8; ++i) {
+        auto other_query = _generate_on_query(other_wg);
+        other_query->query_mem_tracker()->consume(kSmallQueryMem);
+        other_queries.push_back(other_query);
+    }
+    Defer release_other_memory {[&]() {
+        for (auto& other_query : other_queries) {
+            other_query->query_mem_tracker()->consume(-kSmallQueryMem);
+        }
+    }};
+    other_wg->refresh_memory_usage();
+    ASSERT_GT(other_wg->total_mem_used() - other_wg->min_memory_limit(), 128L 
* 1024 * 1024);
+    ASSERT_EQ(_wg_manager->revoke_memory_from_other_groups_(), 0);
+
+    ScopedProcessMemLimits::exceed_soft_limit();
+    _wg_manager->add_paused_query(query->resource_ctx(), 1024L,
+                                  
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+
+    config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+    for (int i = 0; i < 3; ++i) {
+        _wg_manager->handle_paused_queries();
+        ASSERT_FALSE(query->is_cancelled());
+        ASSERT_FALSE(_wg_manager->revoking_memory_from_other_query_);
+        ASSERT_EQ(_paused_queries_count(wg), 1);
+    }
+    _expire_paused_query(wg);
+
+    _wg_manager->handle_paused_queries();
+    assert_cancelled_for_process_memory(query, false);
+    for (auto& other_query : other_queries) {
+        ASSERT_FALSE(other_query->is_cancelled());
+    }
+}
+
+// When the process reaches the hard memory limit, the paused query should be 
cancelled without
+// waiting for the timeout, so that the protection does not depend on memory 
gc.
+TEST_F(WorkloadGroupManagerTest, process_mem_exceeded_cancels_at_hard_limit) {
+    ScopedProcessMemLimits process_mem_limits;
+    auto wg = create_process_pressure_wg(_wg_manager.get());
+
+    auto query = _generate_on_query(wg);
+    query->query_mem_tracker()->consume(1024L * 1024 * 128);
+    Defer release_memory {[&]() { query->query_mem_tracker()->consume(-1024L * 
1024 * 128); }};
+    wg->refresh_memory_usage();
+    ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+    ScopedProcessMemLimits::exceed_hard_limit();
+    _wg_manager->add_paused_query(query->resource_ctx(), 1024L,
+                                  
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+
+    config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+    _wg_manager->handle_paused_queries();
+    assert_cancelled_for_process_memory(query, true);
+}
+
 } // namespace doris


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to