This is an automated email from the ASF dual-hosted git repository.
mrhhsg pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new 14209286e9d [fix](be) Keep paused queries waiting under process memory
pressure before cancelling (#67845)
14209286e9d is described below
commit 14209286e9d5f5cdf24b169df4aa20515900ebfd
Author: Jerry Hu <[email protected]>
AuthorDate: Tue Oct 13 17:33:50 2026 +0800
[fix](be) Keep paused queries waiting under process memory pressure before
cancelling (#67845)
### What problem does this PR solve?
Issue Number: None
Related PR: #62485
Problem Summary:
When a query is paused because of `PROCESS_MEMORY_EXCEEDED` and it has
no
revocable memory, `WorkloadGroupMgr::handle_single_query_()` cancelled
it
immediately if the process was still above the soft memory limit. The
memory maintenance thread runs every 50ms by default, so a query that
just
finished spilling everything it could was cancelled about 50-100ms
later,
even when it only used a few dozen MiB and cancelling it did not relieve
the process memory pressure at all.
The timeout fallback added by #62485 in
`handle_process_memory_exceeded_()`
only applies when the workload group uses less than its min memory
limit.
With the default `min_memory_percent = 0` that branch is never reached,
and the query goes straight to `handle_single_query_()`, which cancels
it
at once.
This PR keeps a query paused by process memory pressure waiting until it
has
waited `spill_in_paused_queue_timeout_ms`, re-checking the pressure
every
round so the query resumes as soon as other queries release memory. The
query
is cancelled after the timeout, or immediately once the process reaches
the
hard memory limit. In detail:
- `handle_process_memory_exceeded_()` checks the pressure before routing
the
query, so a query whose workload group is at or below its minimum also
resumes promptly. The check tests the recorded reservation size, which
is the
same predicate the failed reservation used, instead of a 32 MiB probe
that
could keep a small request paused until the timeout.
- `revoke_memory_from_other_groups_()` returns the memory actually
revoked
from the overcommitted peer workload group. A peer whose queries are all
too
small to be cancelled frees nothing, so the paused query no longer waits
for a
release that never comes (and is no longer resumed, restarting its
timer,
without any memory being freed); it falls through to the
hard-limit/timeout
fallback instead.
- The process-pressure decision of `handle_single_query_()` moved to
`resolve_process_memory_exceeded_query_()`. It re-evaluates the
soft-limit
predicate right before cancelling, so pressure relieved while the
manager
inspected the pipeline or revoked memory from other groups resumes the
query.
It spills when the query has revocable memory and no running task. Below
the
hard limit a running task postpones the decision until it yields; at the
hard
limit the query is cancelled at once because spilling is unsafe while a
task
runs, so the protection does not depend on `disable_memory_gc`.
- `add_paused_query()` records the largest reservation pending for a
query
whose tasks failed several reservations, and the wait is bounded from
the
first reservation that failed since the last successful one rather than
from
the latest pause: `TaskController` keeps that start, a retry that fails
again
before any reservation succeeded continues the same wait, and
`PipelineTask::_try_to_reserve_memory()` ends it when a reservation
succeeds.
- The pause timer uses a monotonic clock, and the cancellation message
no longer
claims to "disable reserve memory and resume" the query it cancels.
### 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
- [ ] Regression test
- [x] Unit Test
- [ ] Manual test (add detailed scripts or steps below)
- [ ] No need to test or manual test. Explain why:
- Behavior changed:
- [ ] No.
- [x] Yes. Queries paused due to process memory exceeded are kept paused
until `spill_in_paused_queue_timeout_ms` instead of being cancelled
immediately when no memory can be revoked; they are cancelled at once at
the process hard limit, even while one of their tasks is running.
- Does this need documentation?
- [x] No.
- [ ] Yes.
---
be/src/exec/pipeline/pipeline_task.cpp | 4 +
.../workload_group/workload_group_manager.cpp | 342 ++++++---
.../workload_group/workload_group_manager.h | 70 +-
.../runtime/workload_management/task_controller.h | 15 +
be/src/util/mem_info.h | 7 +
be/test/exec/pipeline/pipeline_task_test.cpp | 22 +
.../workload_group/workload_group_manager_test.cpp | 844 ++++++++++++++++++++-
be/test/testutil/mock/mock_query_task_controller.h | 55 ++
8 files changed, 1251 insertions(+), 108 deletions(-)
diff --git a/be/src/exec/pipeline/pipeline_task.cpp
b/be/src/exec/pipeline/pipeline_task.cpp
index 1bffe2e764a..aa5d392c19e 100644
--- a/be/src/exec/pipeline/pipeline_task.cpp
+++ b/be/src/exec/pipeline/pipeline_task.cpp
@@ -870,6 +870,10 @@ bool PipelineTask::_try_to_reserve_memory(const size_t
reserve_size, OperatorBas
_spilling = true;
return false;
}
+ // The query made progress: a later reservation that fails for process
memory starts a new
+ // bounded wait instead of continuing the one that this reservation may
have been retried
+ // from.
+
_state->get_query_ctx()->resource_ctx()->task_controller()->end_process_memory_wait();
return true;
}
diff --git a/be/src/runtime/workload_group/workload_group_manager.cpp
b/be/src/runtime/workload_group/workload_group_manager.cpp
index ed4ce5bd0d0..c6dd3f6bb7f 100644
--- a/be/src/runtime/workload_group/workload_group_manager.cpp
+++ b/be/src/runtime/workload_group/workload_group_manager.cpp
@@ -23,6 +23,7 @@
#include <memory>
#include <mutex>
#include <unordered_map>
+#include <vector>
#include "common/config.h"
#include "common/status.h"
@@ -46,13 +47,12 @@ const static std::string INTERNAL_NORMAL_WG_NAME = "normal";
const static uint64_t INTERNAL_NORMAL_WG_ID = 1;
PausedQuery::PausedQuery(std::shared_ptr<ResourceContext> resource_ctx, double
cache_ratio,
- int64_t reserve_size)
+ int64_t reserve_size, int64_t enqueue_at_ms)
: resource_ctx_(resource_ctx),
+ enqueue_at(enqueue_at_ms),
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();
-}
+ query_id_(print_id(resource_ctx->task_controller()->task_id())) {}
WorkloadGroupMgr::~WorkloadGroupMgr() = default;
@@ -288,18 +288,51 @@ void WorkloadGroupMgr::add_paused_query(const
std::shared_ptr<ResourceContext>&
resource_ctx->task_controller()->update_paused_reason(status);
resource_ctx->task_controller()->set_low_memory_mode(true);
resource_ctx->task_controller()->set_memory_sufficient(false);
+ // A query paused for process memory is resumed once its recorded
reservation fits and every
+ // blocked task then retries its own request. A retry that fails again
(another task took the
+ // memory first, or a sibling that fit on its own still does not fit next
to the others)
+ // pauses the query again: it continues the wait it was resumed from, so
that the wait stays
+ // bounded by spill_in_paused_queue_timeout_ms. The wait ends with the
first reservation that
+ // succeeds, see PipelineTask::_try_to_reserve_memory().
+ const int64_t now_ms = MonotonicMillis();
+ const int64_t enqueue_at =
+ status.is<ErrorCode::PROCESS_MEMORY_EXCEEDED>()
+ ?
resource_ctx->task_controller()->start_process_memory_wait(now_ms)
+ : now_ms;
std::lock_guard<std::mutex> lock(_paused_queries_lock);
auto wg = resource_ctx->workload_group();
- auto&& [it, inserted] = _paused_queries_list[wg].emplace(
+ auto& queries_list = _paused_queries_list[wg];
+ auto&& [it, inserted] = queries_list.emplace(
resource_ctx,
doris::GlobalMemoryArbitrator::last_affected_cache_capacity_adjust_weighted,
- reserve_size);
+ reserve_size, enqueue_at);
// Check if this is an invalid reserve, for example, if the reserve size
is too large, larger than the query limit
// if hard limit is enabled, then not need enable other queries hard limit.
if (inserted) {
LOG(INFO) << "Insert one new paused query: "
<< resource_ctx->task_controller()->debug_string()
- << ", workload group: " << wg->debug_string();
+ << ", workload group: " << wg->debug_string() << ", waited "
<< it->elapsed_time()
+ << " ms";
+ } else if (reserve_size > it->reserve_size_) {
+ // Another task of the same query already failed a smaller
reservation. The query is
+ // resumed as a whole and every blocked task then retries its own
request, so record the
+ // largest pending one: the query is woken up once that one fits,
which is when the
+ // tasks can reserve one after another (a task releases its
reservation after each
+ // block, so the pending requests do not need to fit at the same
time). Waking it up
+ // for a smaller one would fail the largest again right away.
+ LOG(INFO) << "Query: " <<
print_id(resource_ctx->task_controller()->task_id())
+ << " is already paused, raise the recorded reservation from "
+ << PrettyPrinter::print_bytes(it->reserve_size_) << " to "
+ << PrettyPrinter::print_bytes(reserve_size);
+ auto node = queries_list.extract(it);
+ node.value().reserve_size_ = reserve_size;
+ queries_list.insert(std::move(node));
+ } else {
+ LOG(INFO) << "Query: " <<
print_id(resource_ctx->task_controller()->task_id())
+ << " is already paused, keep the recorded reservation "
+ << PrettyPrinter::print_bytes(it->reserve_size_)
+ << " which is not smaller than the pending "
+ << PrettyPrinter::print_bytes(reserve_size);
}
}
@@ -602,6 +635,16 @@ bool WorkloadGroupMgr::handle_process_memory_exceeded_(
return false;
}
+ // 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.
+ // Test the recorded reservation itself: this is the predicate the failed
reservation used,
+ // so the query is resumed exactly when its request fits now.
+ if
(!GlobalMemoryArbitrator::is_exceed_soft_mem_limit(query_it->reserve_size_)) {
+ resume_process_paused_query_(wg, queries_list, query_it, resource_ctx);
+ return false;
+ }
+
// 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()) {
@@ -609,8 +652,11 @@ bool WorkloadGroupMgr::handle_process_memory_exceeded_(
}
// 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_();
+ // by cancelling their queries. Only the memory actually revoked counts: a
workload group that
+ // exceeds its min memory but has nothing cancellable releases nothing
(the next exceeding
+ // workload group is tried instead), and waiting for it would resume this
query without any
+ // memory being freed.
+ int64_t revoked_size =
revoke_memory_from_other_groups_(query_it->reserve_size_);
if (revoked_size > 0) {
// Revoke memory from other workload groups will cancel some queries,
wait them cancel finished
// and then check it again.
@@ -618,19 +664,39 @@ bool WorkloadGroupMgr::handle_process_memory_exceeded_(
return true;
}
+ // While the other workload groups were scanned, other queries or the
cache may have released
+ // enough memory for the recorded reservation (the walk stops as soon as
it fits). Resume the
+ // query instead of cancelling it below from the stale observation.
+ if
(!GlobalMemoryArbitrator::is_exceed_soft_mem_limit(query_it->reserve_size_)) {
+ resume_process_paused_query_(wg, queries_list, query_it, resource_ctx);
+ return false;
+ }
+
// TODO revoke from memtable
- // Fallback: if we have waited too long and cannot revoke from anywhere,
- // cancel the query or disable reserve memory to let it proceed.
- if (query_it->elapsed_time() > config::spill_in_paused_queue_timeout_ms) {
- // Cannot spill (no revocable memory), cannot revoke from other WGs,
- // and process memory is still exceeded. Cancel the query to protect
the system.
+ // Fallback: if the process reaches the hard limit or we have waited too
long and cannot
+ // revoke from anywhere, let this query spill or cancel it to protect the
process.
+ if (GlobalMemoryArbitrator::is_exceed_hard_mem_limit() ||
+ query_it->elapsed_time() > config::spill_in_paused_queue_timeout_ms) {
+ // No other workload group can release memory, and process memory is
still exceeded.
return release_query_memory_(queries_list, query_it, resource_ctx,
true);
}
++query_it;
return false;
}
+void WorkloadGroupMgr::resume_process_paused_query_(
+ const WorkloadGroupPtr& wg, PausedQuerySet& queries_list,
PausedQueryIterator& query_it,
+ const std::shared_ptr<ResourceContext>& resource_ctx) {
+ 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);
+}
+
bool WorkloadGroupMgr::release_query_memory_(PausedQuerySet& queries_list,
PausedQueryIterator& query_it,
const
std::shared_ptr<ResourceContext>& resource_ctx,
@@ -655,19 +721,25 @@ bool
WorkloadGroupMgr::release_query_memory_(PausedQuerySet& queries_list,
return stop_after_release;
}
-// Find the workload group that could revoke lot of memory:
-// 1. workload group = max(total used memory - min memory that should reserved
for it)
+// Find the workload groups that could revoke lot of memory:
+// 1. order the workload groups by (total used memory - min memory that should
reserved for it),
+// starting with the one that exceeds most.
// 2. revoke 10% memory of the workload group that exceeded. For example, if
the workload group exceed 10g,
-// then revoke 1g memory.
+// then revoke 1g memory. A workload group may release nothing although it
exceeds its min
+// memory (for example, every query is too small to be cancelled), then try
the next one.
// 3. After revoke memory, go to the loop and wait for the query to be
cancelled and check again.
-int64_t WorkloadGroupMgr::revoke_memory_from_other_groups_() {
+// `reserve_size` is the reservation of the paused query this walk serves: the
walk stops once it
+// fits, nothing needs to be revoked for it any more.
+// Returns the memory actually revoked, which is 0 when no query of any of
these workload groups
+// could be cancelled.
+int64_t WorkloadGroupMgr::revoke_memory_from_other_groups_(int64_t
reserve_size) {
MonotonicStopWatch watch;
watch.start();
std::unique_ptr<RuntimeProfile> profile =
std::make_unique<RuntimeProfile>("RevokeMemoryFromOtherGroups");
- WorkloadGroupPtr max_wg = nullptr;
- int64_t max_exceeded_memory = 0;
+ // (exceeded memory, workload group), the most exceeded first.
+ std::vector<std::pair<int64_t, WorkloadGroupPtr>> exceeded_wgs;
{
std::shared_lock<std::shared_mutex> r_lock(_group_mutex);
for (auto& workload_group : _workload_groups) {
@@ -678,32 +750,76 @@ int64_t
WorkloadGroupMgr::revoke_memory_from_other_groups_() {
// then not revoke memory from it.
continue;
}
- if (total_used_memory - min_memory_limit > max_exceeded_memory) {
- max_wg = workload_group.second;
- max_exceeded_memory = total_used_memory - min_memory_limit;
- }
+ exceeded_wgs.emplace_back(total_used_memory - min_memory_limit,
workload_group.second);
+ }
+ }
+ std::sort(exceeded_wgs.begin(), exceeded_wgs.end(),
+ [](const auto& lhs, const auto& rhs) { return lhs.first >
rhs.first; });
+
+ int64_t freed_mem = 0;
+ size_t tried_wgs = 0;
+ for (const auto& [snapshot_exceeded_memory, wg] : exceeded_wgs) {
+ if (snapshot_exceeded_memory < 1 << 27) {
+ // The remaining workload groups exceed even less.
+ LOG(INFO) << "The workload group that exceed most memory among the
untried ones is :"
+ << wg->memory_debug_string() << ", exceeded_memory: "
+ << PrettyPrinter::print(snapshot_exceeded_memory,
TUnit::BYTES)
+ << " less than 128MB, no need to revoke memory";
+ break;
+ }
+ // Scanning the earlier workload groups took time (a cancellation in
flight is waited for
+ // up to revoke_memory_max_tolerance_ms), during which other queries
or the cache may have
+ // released memory. The caller checked the reservation before the walk
only: stop as soon
+ // as it fits, so that no further query is cancelled for pressure that
is gone.
+ if (tried_wgs > 0 &&
!GlobalMemoryArbitrator::is_exceed_soft_mem_limit(reserve_size)) {
+ LOG(INFO) << "Process memory pressure is relieved after trying "
<< tried_wgs
+ << " workload group(s), the reservation "
+ << PrettyPrinter::print_bytes(reserve_size)
+ << " fits now, stop revoking memory from other workload
groups";
+ break;
+ }
+ // The snapshot above was taken before the earlier workload groups
were scanned, during
+ // which this one may have released memory (its queries finished or
were cancelled).
+ // total_mem_used() is the usage cached by the last maintenance
refresh, so refresh it and
+ // decide from the current usage: the min memory stays reserved for a
workload group that
+ // no longer exceeds it, and the amount to revoke is what it exceeds
by now.
+ const int64_t exceeded_memory = wg->refresh_memory_usage() -
wg->min_memory_limit();
+ if (exceeded_memory < 1 << 27) {
+ LOG(INFO) << "The workload group " << wg->memory_debug_string()
+ << " exceeded its min memory by "
+ << PrettyPrinter::print(snapshot_exceeded_memory,
TUnit::BYTES)
+ << " when the walk started, now by "
+ << PrettyPrinter::print(exceeded_memory, TUnit::BYTES)
+ << " less than 128MB, skip it";
+ continue;
}
+ auto need_free_mem = static_cast<int64_t>((double)exceeded_memory *
0.1);
+ // Revoke 10% of memory from the workload group that exceed most memory
+ freed_mem = wg->revoke_memory(need_free_mem, "exceed_memory",
profile.get());
+ ++tried_wgs;
+ LOG(INFO) << fmt::format(
+ "[MemoryGC] process memory not enough, revoke memory from
workload_group: {}, "
+ "need free memory {}, freed memory {}. cost(us): {}",
+ wg->memory_debug_string(),
PrettyPrinter::print_bytes(need_free_mem),
+ PrettyPrinter::print_bytes(freed_mem), watch.elapsed_time() /
1000);
+ if (freed_mem > 0) {
+ break;
+ }
+ // Nothing could be cancelled in this workload group, try the next one
that exceeds
+ // its min memory.
+ }
+ if (tried_wgs > 0) {
+ // Every tried workload group appended its own child to the profile,
so print it once
+ // after the walk instead of once per group: this runs under
_paused_queries_lock on
+ // every maintenance round while the pressure lasts.
+ std::stringstream ss;
+ profile->pretty_print(&ss);
+ LOG(INFO) << fmt::format(
+ "[MemoryGC] process memory not enough, tried to revoke memory
from {} workload "
+ "group(s), freed memory {}. cost(us): {}, details: {}",
+ tried_wgs, PrettyPrinter::print_bytes(freed_mem),
watch.elapsed_time() / 1000,
+ ss.str());
}
- if (max_wg == nullptr) {
- return 0;
- }
- if (max_exceeded_memory < 1 << 27) {
- LOG(INFO) << "The workload group that exceed most memory is :"
- << max_wg->memory_debug_string() << ", max_exceeded_memory: "
- << PrettyPrinter::print(max_exceeded_memory, TUnit::BYTES)
- << " less than 128MB, no need to revoke memory";
- return 0;
- }
- int64_t freed_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());
- 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());
return freed_mem;
}
@@ -726,20 +842,23 @@ bool WorkloadGroupMgr::handle_single_query_(const
std::shared_ptr<ResourceContex
const auto query_id = print_id(requestor->task_controller()->task_id());
requestor->task_controller()->get_revocable_info(&revocable_size,
&memory_usage,
&has_running_task);
+ if (!paused_reason.is<ErrorCode::QUERY_MEMORY_EXCEEDED>() &&
+ !paused_reason.is<ErrorCode::WORKLOAD_GROUP_MEMORY_EXCEEDED>()) {
+ // PROCESS_MEMORY_EXCEEDED is resolved separately: the process hard
limit and the
+ // timeout must be enforced even when a running task prevents
spilling, which the
+ // early exits below skip.
+ return resolve_process_memory_exceeded_query_(requestor,
size_to_reserve, time_in_queue,
+ memory_usage,
has_running_task);
+ }
if (has_running_task) {
- LOG(INFO) << "Query: " <<
print_id(requestor->task_controller()->task_id())
- << " is paused, but still has running task, skip it.";
+ LOG(INFO) << "Query: " << query_id << " is paused, but still has
running task, skip it.";
return false;
}
const auto wg = requestor->workload_group();
auto revocable_tasks = requestor->task_controller()->get_revocable_tasks();
if (!revocable_tasks.empty()) {
- SCOPED_ATTACH_TASK(requestor);
- auto status = requestor->task_controller()->revoke_memory();
- if (!status.ok()) {
- requestor->task_controller()->cancel(status);
- }
+ spill_query_(requestor);
return true;
}
const auto limit = requestor->memory_context()->mem_limit();
@@ -781,7 +900,10 @@ bool WorkloadGroupMgr::handle_single_query_(const
std::shared_ptr<ResourceContex
requestor->task_controller()->set_memory_sufficient(true);
return true;
}
- } else if (paused_reason.is<ErrorCode::WORKLOAD_GROUP_MEMORY_EXCEEDED>()) {
+ }
+
+ DCHECK(paused_reason.is<ErrorCode::WORKLOAD_GROUP_MEMORY_EXCEEDED>());
+ {
bool exceed_low_watermark = false;
bool exceed_high_watermark = false;
wg->check_mem_used(size_to_reserve, &exceed_low_watermark,
&exceed_high_watermark);
@@ -823,42 +945,92 @@ bool WorkloadGroupMgr::handle_single_query_(const
std::shared_ptr<ResourceContex
requestor->task_controller()->cancel(error_status);
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;
- }
}
}
+void WorkloadGroupMgr::spill_query_(const std::shared_ptr<ResourceContext>&
requestor) {
+ SCOPED_ATTACH_TASK(requestor);
+ auto status = requestor->task_controller()->revoke_memory();
+ if (!status.ok()) {
+ requestor->task_controller()->cancel(status);
+ }
+}
+
+bool WorkloadGroupMgr::resolve_process_memory_exceeded_query_(
+ const std::shared_ptr<ResourceContext>& requestor, size_t
size_to_reserve,
+ int64_t time_in_queue, size_t memory_usage, bool has_running_task) {
+ const auto query_id = print_id(requestor->task_controller()->task_id());
+ const auto wg = requestor->workload_group();
+
+ // Spilling needs every task of the query to be idle, so the revocable
tasks are only
+ // collected while none is running. A running task does not postpone the
cancellation
+ // below: cancelling is safe while the task runs.
+ std::vector<PipelineTask*> revocable_tasks;
+ if (!has_running_task) {
+ revocable_tasks = requestor->task_controller()->get_revocable_tasks();
+ }
+
+ // The caller observed the pressure before it inspected the pipeline tasks
and tried to
+ // revoke memory from other workload groups, during which other queries or
the cache may have
+ // released memory. Re-check the recorded reservation after the last
inspection and right
+ // before acting, so that the query is never spilled or cancelled from a
stale observation.
+ if (!GlobalMemoryArbitrator::is_exceed_soft_mem_limit(size_to_reserve)) {
+ 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;
+ }
+
+ if (!revocable_tasks.empty()) {
+ // The query is resumed by the spill callback and reserves memory
again. If that
+ // still fails, it comes back here without revocable memory.
+ spill_query_(requestor);
+ return true;
+ }
+
+ const bool exceed_hard_mem_limit =
GlobalMemoryArbitrator::is_exceed_hard_mem_limit();
+ if (!exceed_hard_mem_limit && time_in_queue <
config::spill_in_paused_queue_timeout_ms) {
+ // The query has no revocable memory (or a running task prevents
spilling), 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, has running task: " << has_running_task
+ << ", 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:
{}, has running "
+ "task: {}), 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, has_running_task,
PrettyPrinter::print_bytes(memory_usage),
+
PrettyPrinter::print_bytes(requestor->memory_context()->mem_limit()),
+
PrettyPrinter::print_bytes(requestor->memory_context()->reserved_consumption()),
+ 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;
+}
+
void WorkloadGroupMgr::update_queries_limit_(WorkloadGroupPtr wg, bool
enable_hard_limit) {
auto wg_mem_limit = wg->memory_limit();
auto all_resource_ctxs = wg->resource_ctxs();
diff --git a/be/src/runtime/workload_group/workload_group_manager.h
b/be/src/runtime/workload_group/workload_group_manager.h
index f4c4504ad08..ea379594260 100644
--- a/be/src/runtime/workload_group/workload_group_manager.h
+++ b/be/src/runtime/workload_group/workload_group_manager.h
@@ -25,6 +25,7 @@
#include "common/be_mock_util.h"
#include "runtime/workload_group/workload_group.h"
+#include "util/time.h"
namespace doris {
@@ -41,18 +42,17 @@ 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;
+ // Monotonic ms. For a query paused for process memory this is the start
of its bounded wait,
+ // which may predate this entry: see
TaskController::start_process_memory_wait().
+ int64_t enqueue_at {0};
size_t last_mem_usage {0};
double cache_ratio_ {0.0};
int64_t reserve_size_ {0};
PausedQuery(std::shared_ptr<ResourceContext> resource_ctx_, double
cache_ratio,
- int64_t reserve_size);
+ int64_t reserve_size, int64_t enqueue_at_ms);
- int64_t elapsed_time() const {
- auto now = std::chrono::system_clock::now();
- return std::chrono::duration_cast<std::chrono::milliseconds>(now -
enqueue_at).count();
- }
+ int64_t elapsed_time() const { return MonotonicMillis() - enqueue_at; }
std::string query_id() const { return query_id_; }
@@ -194,16 +194,28 @@ private:
// Handle a query paused due to PROCESS_MEMORY_EXCEEDED. The logic is:
// 1. If any recently cancelled query exists globally, skip (wait for
// process-level memory release).
- // 2. If the WG's usage exceeds its min_memory_limit, spill via
+ // 2. If the process is no longer above the soft memory limit, resume the
+ // query immediately.
+ // 3. If the WG's usage exceeds its min_memory_limit, spill via
// release_query_memory_(stop_after_release=true) — one spill per
round.
- // 3. Otherwise, try to revoke memory from other overcommitted WGs by
- // cancelling their largest queries.
+ // 4. Otherwise, try to revoke memory from other overcommitted WGs by
+ // cancelling their largest queries. Only memory actually revoked
counts.
+ // 5. If nothing could be revoked, re-check the recorded reservation:
the walk may have
+ // taken a while and stops as soon as the reservation fits, then
resume the query.
+ // 6. If the process has reached its hard limit or the query has waited
for
+ // spill_in_paused_queue_timeout_ms, spill or cancel it via
release_query_memory_.
// Returns true if the caller should stop processing further queries/WGs.
bool handle_process_memory_exceeded_(const WorkloadGroupPtr& wg,
PausedQuerySet& queries_list,
PausedQueryIterator& query_it,
const
std::shared_ptr<ResourceContext>& resource_ctx,
const RecentlyCancelledQueries&
recently_cancelled);
+ // Resume a query paused due to PROCESS_MEMORY_EXCEEDED whose recorded
reservation fits now,
+ // and erase it from the paused list (the iterator is advanced).
+ void resume_process_paused_query_(const WorkloadGroupPtr& wg,
PausedQuerySet& queries_list,
+ PausedQueryIterator& query_it,
+ const std::shared_ptr<ResourceContext>&
resource_ctx);
+
// Common helper: attempt to release memory for a single paused query by
calling
// handle_single_query_() (which triggers spill or cancel). On success,
erases the
// query from the paused list. If the query was cancelled (not spilled),
sets
@@ -217,16 +229,42 @@ private:
// Attempt to resolve a single paused query: if revocable memory exists,
trigger
// spill; if under limit, resume; if no memory can be freed, cancel the
query or
- // disable reserve memory and resume. Returns true if the query was acted
upon
- // (spilled/cancelled/resumed), false if it should keep waiting (e.g.,
still has
- // running tasks).
+ // disable reserve memory and resume. PROCESS_MEMORY_EXCEEDED is delegated
to
+ // resolve_process_memory_exceeded_query_. 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);
- // Find the most overcommitted workload group (usage - min_memory_limit is
- // largest) and cancel its biggest query to reclaim ~10% of the excess
memory.
- // Returns the amount of memory expected to be freed, or 0 if no WG
qualifies.
- int64_t revoke_memory_from_other_groups_();
+ // Resolve a query paused due to PROCESS_MEMORY_EXCEEDED, called by
handle_single_query_:
+ // 1. Collect the revocable tasks if no task is running, then re-check
the recorded
+ // reservation against the soft limit and resume the query if it fits
now, the
+ // caller's observation may be stale.
+ // 2. Spill if revocable tasks were collected.
+ // 3. Otherwise 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. A
running task only
+ // prevents spilling, not the cancellation, so the bounded wait and
the hard-limit
+ // protection hold for it too and do not depend on memory gc.
+ // Return value as handle_single_query_.
+ bool resolve_process_memory_exceeded_query_(const
std::shared_ptr<ResourceContext>& requestor,
+ size_t size_to_reserve,
int64_t time_in_queue,
+ size_t memory_usage, bool
has_running_task);
+
+ // Spill the revocable tasks of the query, the query is resumed by the
spill callback.
+ // The query is cancelled if the spill could not be started.
+ void spill_query_(const std::shared_ptr<ResourceContext>& requestor);
+
+ // Walk the overcommitted workload groups (usage - min_memory_limit,
largest first)
+ // and cancel the biggest queries of the first one that releases memory,
to reclaim
+ // ~10% of its excess memory. A WG whose queries are all too small to be
cancelled
+ // releases nothing, so the next one is tried. Each WG is re-checked
against its
+ // refreshed usage right before it is revoked from, so one that fell back
within its
+ // min memory while the earlier ones were scanned keeps that memory
reserved. The walk
+ // also stops before the next WG once `reserve_size` (the reservation of
the paused
+ // query it serves) fits under the process soft limit.
+ // Returns the amount of memory actually revoked, or 0 if no WG qualifies
or no
+ // query of any of them could be cancelled.
+ int64_t revoke_memory_from_other_groups_(int64_t reserve_size);
// Recalculate and apply per-query memory limits for all queries in a
workload
// group based on the slot memory policy (NONE/FIXED/DYNAMIC). When
diff --git a/be/src/runtime/workload_management/task_controller.h
b/be/src/runtime/workload_management/task_controller.h
index ed7b3e8e3c8..d38d44ec607 100644
--- a/be/src/runtime/workload_management/task_controller.h
+++ b/be/src/runtime/workload_management/task_controller.h
@@ -97,6 +97,20 @@ public:
void reset_paused_reason() { paused_reason_.reset(); }
Status paused_reason() { return paused_reason_.status(); }
void add_paused_count() { paused_count_.fetch_add(1); }
+ // The wait of a task paused for process memory is bounded from the first
reservation that
+ // failed since its last successful one, not from its latest pause. The
task is resumed to
+ // retry its pending reservations, and a retry that fails again before any
of them succeeded
+ // continues the same wait instead of starting a new one. Returns the
start of the wait
+ // (monotonic ms), starting it at `now_ms` when none is pending.
+ int64_t start_process_memory_wait(int64_t now_ms) {
+ int64_t expected = 0;
+ return process_memory_wait_start_ms_.compare_exchange_strong(expected,
now_ms) ? now_ms
+
: expected;
+ }
+ // A reservation succeeded: the task made progress, the next failed
reservation starts a new
+ // wait.
+ void end_process_memory_wait() { process_memory_wait_start_ms_ = 0; }
+ int64_t process_memory_wait_start_ms() const { return
process_memory_wait_start_ms_; }
/* memory status action
*/
@@ -146,6 +160,7 @@ protected:
*/
AtomicStatus paused_reason_;
std::atomic<int64_t> paused_count_ = 0;
+ std::atomic<int64_t> process_memory_wait_start_ms_ = 0;
/* memory status property
*/
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/exec/pipeline/pipeline_task_test.cpp
b/be/test/exec/pipeline/pipeline_task_test.cpp
index 34f8df73810..b91c4c6917e 100644
--- a/be/test/exec/pipeline/pipeline_task_test.cpp
+++ b/be/test/exec/pipeline/pipeline_task_test.cpp
@@ -1427,6 +1427,28 @@ TEST_F(PipelineTaskTest,
TEST_RESERVE_MEMORY_FAIL_SPILLABLE) {
EXPECT_TRUE(
((MockWorkloadGroupMgr*)ExecEnv::GetInstance()->_workload_group_manager)->_paused);
}
+ {
+ // The reservations succeed now. The query made progress: the wait
that a reservation
+ // failed for process memory may have started is ended, and the task
is not paused.
+
((MockThreadMemTrackerMgr*)thread_context()->thread_mem_tracker_mgr.get())
+ ->_test_low_memory = false;
+
_query_ctx->resource_ctx()->task_controller()->start_process_memory_wait(1);
+
EXPECT_EQ(_query_ctx->resource_ctx()->task_controller()->process_memory_wait_start_ms(),
1);
+ task->_spilling = false;
+
((MockWorkloadGroupMgr*)ExecEnv::GetInstance()->_workload_group_manager)->_paused
= false;
+ EXPECT_EQ(task->_exec_state, PipelineTask::State::RUNNABLE);
+ bool done = false;
+ EXPECT_TRUE(task->execute(&done).ok());
+ EXPECT_FALSE(task->_eos);
+ EXPECT_FALSE(done);
+ EXPECT_FALSE(task->_spilling);
+
EXPECT_TRUE(_query_ctx->resource_ctx()->task_controller()->is_enable_reserve_memory());
+ EXPECT_FALSE(
+
((MockWorkloadGroupMgr*)ExecEnv::GetInstance()->_workload_group_manager)->_paused);
+
EXPECT_EQ(_query_ctx->resource_ctx()->task_controller()->process_memory_wait_start_ms(),
0);
+
((MockThreadMemTrackerMgr*)thread_context()->thread_mem_tracker_mgr.get())
+ ->_test_low_memory = true;
+ }
{
// Reserve failed and paused.
task->_operators.front()->cast<DummyOperator>()._disable_reserve_mem =
true;
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 03a9cedae3d..db047257af5 100644
--- a/be/test/runtime/workload_group/workload_group_manager_test.cpp
+++ b/be/test/runtime/workload_group/workload_group_manager_test.cpp
@@ -17,6 +17,7 @@
#include "runtime/workload_group/workload_group_manager.h"
+#include <fmt/format.h>
#include <gen_cpp/BackendService_types.h>
#include <gen_cpp/PaloInternalService_types.h>
#include <gen_cpp/Types_types.h>
@@ -27,9 +28,14 @@
#include <chrono>
#include <cmath>
#include <filesystem>
+#include <functional>
+#include <limits>
#include <memory>
#include <sstream>
+#include <string>
#include <tuple>
+#include <utility>
+#include <vector>
#include "common/config.h"
#include "common/status.h"
@@ -38,6 +44,7 @@
#include "exec/spill/spill_file_manager.h"
#include "load/memtable/memtable_memory_limiter.h"
#include "runtime/exec_env.h"
+#include "runtime/memory/global_memory_arbitrator.h"
#include "runtime/query_context.h"
#include "runtime/runtime_query_statistics_mgr.h"
#include "runtime/thread_context.h"
@@ -52,10 +59,27 @@
namespace doris {
+// Process memory limit for the PROCESS_MEMORY_EXCEEDED cases. The workload
group memory limits
+// are derived from the process memory limit whenever their usage is
refreshed, so it is set before
+// the workload group is created and kept; process memory pressure is
simulated by the soft memory
+// limit, the memory growth since the last refresh and the system available
memory instead.
+static constexpr int64_t kProcessMemLimitForWg = 1024L * 1024 * 1000;
+// System available memory far above the warning water mark, so that only the
process memory
+// limits decide the process memory predicates unless a case lowers it on
purpose.
+static constexpr int64_t kLargeSysMemAvailable =
std::numeric_limits<int64_t>::max() / 2;
+static constexpr int64_t kProcessPausedReserveSize = 1024L;
+
class WorkloadGroupManagerTest : public testing::Test {
public:
protected:
void SetUp() override {
+ _original_mem_limit = MemInfo::mem_limit();
+ _original_soft_mem_limit = MemInfo::soft_mem_limit();
+ // The process memory predicates also depend on the system available
memory, which is
+ // read from the runner otherwise: a runner below the water marks
reports pressure before
+ // a case sets any process memory limit.
+ _original_sys_mem_available =
+ MemInfo::set_sys_mem_available_for_test(kLargeSysMemAvailable);
_wg_manager = std::make_unique<WorkloadGroupMgr>();
// generate a unique test directory to avoid conflicts between
parallel runs
std::ostringstream _oss;
@@ -95,6 +119,14 @@ protected:
doris::ExecEnv::GetInstance()->set_memtable_memory_limiter(new
MemTableMemoryLimiter());
}
void TearDown() override {
+ for (auto& [query, memory] : _consumed_memory) {
+ query->query_mem_tracker()->consume(-memory);
+ }
+ _consumed_memory.clear();
+ MemInfo::set_mem_limit_for_test(_original_mem_limit);
+ MemInfo::set_soft_mem_limit_for_test(_original_soft_mem_limit);
+ MemInfo::set_sys_mem_available_for_test(_original_sys_mem_available);
+ GlobalMemoryArbitrator::reset_refresh_interval_memory_growth();
_wg_manager.reset();
ExecEnv::GetInstance()->spill_file_mgr()->stop();
SAFE_DELETE(ExecEnv::GetInstance()->_spill_file_mgr);
@@ -153,9 +185,151 @@ private:
}
}
+ // Helpers for the PROCESS_MEMORY_EXCEEDED cases.
+
+ // A workload group whose min memory limit is 100 MiB, so that a query
consuming more than
+ // that routes into handle_single_query_ directly, and a smaller one is
routed through the
+ // other workload groups first.
+ std::shared_ptr<WorkloadGroup> _create_wg_with_min_memory(uint64_t id) {
+ MemInfo::set_mem_limit_for_test(kProcessMemLimitForWg);
+ 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(), 1024L * 1024 * 100);
+ return wg;
+ }
+
+ // A query on `wg` that consumes `memory` bytes until TearDown.
+ std::shared_ptr<QueryContext>
_create_query_with_memory(std::shared_ptr<WorkloadGroup>& wg,
+ int64_t memory) {
+ auto query = _generate_on_query(wg);
+ query->query_mem_tracker()->consume(memory);
+ _consumed_memory.emplace_back(query, memory);
+ wg->refresh_memory_usage();
+ return query;
+ }
+
+ // Release `memory` bytes of `query` without refreshing the usage of its
workload group:
+ // total_mem_used() keeps the value cached by the last refresh, as it does
in production
+ // until the next maintenance round or revocation refreshes it.
+ void _release_query_memory(const std::shared_ptr<QueryContext>& query,
int64_t memory) {
+ query->query_mem_tracker()->consume(-memory);
+ _consumed_memory.emplace_back(query, -memory);
+ }
+
+ // The process memory predicates are controlled from both sides in the
following helpers:
+ // the process memory usage is the vm_rss of PerfCounters (never refreshed
in a unit test)
+ // plus the memory growth since the last refresh, which the helpers set;
the system available
+ // memory is set in SetUp far above the water marks. The process memory
limit is kept at
+ // kProcessMemLimitForWg, which refresh_memory_usage() derives the
workload group limits from.
+
+ // Process soft memory limit is exceeded, hard memory limit is not.
+ static void _exceed_process_soft_mem_limit() {
+ GlobalMemoryArbitrator::reset_refresh_interval_memory_growth();
+ MemInfo::set_mem_limit_for_test(kProcessMemLimitForWg);
+ MemInfo::set_soft_mem_limit_for_test(1);
+
ASSERT_TRUE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(kProcessPausedReserveSize));
+ ASSERT_FALSE(GlobalMemoryArbitrator::is_exceed_hard_mem_limit());
+ }
+
+ // Both the soft and the hard process memory limits are exceeded: the
process has grown by
+ // its whole memory limit since the last refresh.
+ static void _exceed_process_hard_mem_limit() {
+ MemInfo::set_mem_limit_for_test(kProcessMemLimitForWg);
+ MemInfo::set_soft_mem_limit_for_test(kProcessMemLimitForWg);
+ GlobalMemoryArbitrator::refresh_interval_memory_growth =
kProcessMemLimitForWg;
+ ASSERT_TRUE(GlobalMemoryArbitrator::is_exceed_hard_mem_limit());
+ }
+
+ // Neither process memory limit is exceeded: the process memory limit has
its whole headroom
+ // and the system available memory is far above the warning water mark.
+ static void _relieve_process_mem_limit() {
+ GlobalMemoryArbitrator::reset_refresh_interval_memory_growth();
+ MemInfo::set_mem_limit_for_test(kProcessMemLimitForWg);
+ MemInfo::set_soft_mem_limit_for_test(kProcessMemLimitForWg);
+ MemInfo::set_sys_mem_available_for_test(kLargeSysMemAvailable);
+
ASSERT_FALSE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(kProcessPausedReserveSize));
+ }
+
+ void _pause_for_process_memory(const std::shared_ptr<QueryContext>& query)
{
+ _wg_manager->add_paused_query(query->resource_ctx(),
kProcessPausedReserveSize,
+
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+ }
+
+ size_t _paused_query_count(const std::shared_ptr<WorkloadGroup>& 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();
+ }
+
+ // The reservation recorded for the single paused query of `wg`.
+ int64_t _recorded_reserve_size(const std::shared_ptr<WorkloadGroup>& wg) {
+ std::unique_lock<std::mutex> lock(_wg_manager->_paused_queries_lock);
+ auto& queries = _wg_manager->_paused_queries_list[wg];
+ EXPECT_EQ(queries.size(), 1);
+ return queries.begin()->reserve_size_;
+ }
+
+ // Backdate the wait of the single paused query of `wg` by `ms` instead of
sleeping, the timer
+ // is monotonic: both its entry and the wait start that its task
controller carries across a
+ // resume are moved back.
+ void _backdate_paused_query(const std::shared_ptr<WorkloadGroup>& wg,
int64_t ms) {
+ 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 -= ms;
+ auto resource_ctx = node.value().resource_ctx_.lock();
+ ASSERT_TRUE(resource_ctx != nullptr);
+ resource_ctx->task_controller()->end_process_memory_wait();
+ ASSERT_EQ(
+
resource_ctx->task_controller()->start_process_memory_wait(node.value().enqueue_at),
+ node.value().enqueue_at);
+ queries.insert(std::move(node));
+ }
+
+ // How long the single paused query of `wg` has waited.
+ int64_t _paused_elapsed_time(const std::shared_ptr<WorkloadGroup>& wg) {
+ std::unique_lock<std::mutex> lock(_wg_manager->_paused_queries_lock);
+ auto& queries = _wg_manager->_paused_queries_list[wg];
+ EXPECT_EQ(queries.size(), 1);
+ return queries.begin()->elapsed_time();
+ }
+
+ static void _assert_still_paused(const std::shared_ptr<QueryContext>&
query) {
+ ASSERT_FALSE(query->is_cancelled()) <<
query->exec_status().to_string();
+ ASSERT_TRUE(query->resource_ctx()
+ ->task_controller()
+ ->paused_reason()
+ .is<ErrorCode::PROCESS_MEMORY_EXCEEDED>());
+ }
+
+ static void _assert_resumed(const std::shared_ptr<QueryContext>& query) {
+ ASSERT_FALSE(query->is_cancelled()) <<
query->exec_status().to_string();
+
ASSERT_TRUE(query->resource_ctx()->task_controller()->paused_reason().ok());
+ }
+
+ static void _assert_cancelled_by_process_memory(const
std::shared_ptr<QueryContext>& query,
+ bool exceed_hard_limit) {
+ ASSERT_TRUE(query->is_cancelled());
+ const auto status = query->exec_status().to_string();
+ ASSERT_TRUE(query->exec_status().is<ErrorCode::MEM_LIMIT_EXCEEDED>())
<< status;
+ ASSERT_NE(status.find(fmt::format("exceed hard limit: {}",
exceed_hard_limit)),
+ std::string::npos)
+ << status;
+ }
+
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;
+ int64_t _original_mem_limit {0};
+ int64_t _original_soft_mem_limit {0};
+ int64_t _original_sys_mem_available {0};
+ std::vector<std::pair<std::shared_ptr<QueryContext>, int64_t>>
_consumed_memory;
};
TEST_F(WorkloadGroupManagerTest, get_or_create_workload_group) {
@@ -631,7 +805,7 @@ TEST_F(WorkloadGroupManagerTest, ProcessMemoryNotEnough) {
wg3->refresh_memory_usage();
// There is no query in workload groups, so that revoke memory will return 0
- EXPECT_EQ(0, _wg_manager->revoke_memory_from_other_groups_());
+ EXPECT_EQ(0,
_wg_manager->revoke_memory_from_other_groups_(kProcessPausedReserveSize));
// If exceed memory less than 128MB, then not revoke
auto query_context21 = _generate_on_query(wg2);
@@ -641,7 +815,7 @@ TEST_F(WorkloadGroupManagerTest, ProcessMemoryNotEnough) {
EXPECT_EQ(wg2->total_mem_used(), 1024 * 1024 * 50);
EXPECT_EQ(wg2->min_memory_limit(), 1024 * 1024 * 100);
// There is not workload group's memory usage > it's min memory limit.
- EXPECT_EQ(0, _wg_manager->revoke_memory_from_other_groups_());
+ EXPECT_EQ(0,
_wg_manager->revoke_memory_from_other_groups_(kProcessPausedReserveSize));
ASSERT_FALSE(query_context21->is_cancelled());
// Add another query that use a lot of memory
@@ -652,7 +826,7 @@ TEST_F(WorkloadGroupManagerTest, ProcessMemoryNotEnough) {
EXPECT_EQ(wg2->total_mem_used(), 1024 * 1024 * 110);
EXPECT_EQ(wg2->min_memory_limit(), 1024 * 1024 * 100);
// Could not revoke larger than 128MB, not revoke.
- EXPECT_EQ(0, _wg_manager->revoke_memory_from_other_groups_());
+ EXPECT_EQ(0,
_wg_manager->revoke_memory_from_other_groups_(kProcessPausedReserveSize));
ASSERT_FALSE(query_context21->is_cancelled());
ASSERT_FALSE(query_context22->is_cancelled());
@@ -663,7 +837,10 @@ 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 its min memory by 310MB, so 31MB should be revoked. The
largest query (300MB)
+ // is cancelled and its whole memory is reported as revoked.
+ EXPECT_EQ(300L * 1024 * 1024,
+
_wg_manager->revoke_memory_from_other_groups_(kProcessPausedReserveSize));
ASSERT_FALSE(query_context21->is_cancelled());
ASSERT_FALSE(query_context22->is_cancelled());
ASSERT_TRUE(query_context23->is_cancelled());
@@ -689,15 +866,19 @@ 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 the cancelled query31
frees 500MB.
+ EXPECT_EQ(500L * 1024 * 1024,
+
_wg_manager->revoke_memory_from_other_groups_(kProcessPausedReserveSize));
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, so that it is not cancel again. It was
cancelled just now and is
+ // still releasing memory, so its memory is counted as revoked again.
+ EXPECT_EQ(500L * 1024 * 1024,
+
_wg_manager->revoke_memory_from_other_groups_(kProcessPausedReserveSize));
ASSERT_FALSE(query_context11->is_cancelled());
ASSERT_FALSE(query_context21->is_cancelled());
ASSERT_FALSE(query_context22->is_cancelled());
@@ -1358,4 +1539,653 @@ TEST_F(WorkloadGroupManagerTest,
AdaptiveFlushRegistrationSurvivesIdChangeAndReu
controller->adjust_once();
}
+// 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) {
+ auto wg = _create_wg_with_min_memory(1);
+ // 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.
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ for (int i = 0; i < 3; ++i) {
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ ASSERT_EQ(_paused_query_count(wg), 1);
+ }
+
+ // Process memory pressure is relieved, the query should be resumed.
+ _relieve_process_mem_limit();
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ ASSERT_EQ(_paused_query_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) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ _backdate_paused_query(wg, config::spill_in_paused_queue_timeout_ms + 1);
+
+ _wg_manager->handle_paused_queries();
+ _assert_cancelled_by_process_memory(query, false);
+}
+
+// The recovery check must test the recorded reservation, not a larger probe:
after other
+// queries release some memory, a small reservation fits under the process
soft limit while a
+// 32 MiB probe does not, and the query must be resumed instead of waiting for
the timeout.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_resumes_when_reservation_fits_soft_limit) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+
+ // Leave 16 MiB headroom below the soft limit: more than the recorded
reservation, less
+ // than a 32 MiB probe.
+
MemInfo::set_soft_mem_limit_for_test(GlobalMemoryArbitrator::process_memory_usage()
+
+ 1024L * 1024 * 16);
+
ASSERT_FALSE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(kProcessPausedReserveSize));
+ ASSERT_TRUE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(1024L * 1024
* 32));
+ _backdate_paused_query(wg, config::spill_in_paused_queue_timeout_ms + 1);
+
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ ASSERT_EQ(_paused_query_count(wg), 0);
+}
+
+// Same as above for the system available memory boundary: the recorded
reservation keeps the
+// system available memory above the warning water mark while a 32 MiB probe
does not.
+TEST_F(WorkloadGroupManagerTest,
+ process_mem_exceeded_resumes_when_reservation_fits_sys_mem_available) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+
+ // The process limits have their whole headroom, only the system available
memory is close
+ // to the warning water mark: 16 MiB above it.
+ _relieve_process_mem_limit();
+ ASSERT_GT(MemInfo::sys_mem_available_warning_water_mark(), 0);
+
MemInfo::set_sys_mem_available_for_test(MemInfo::sys_mem_available_warning_water_mark()
+
+ 1024L * 1024 * 16);
+
ASSERT_FALSE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(kProcessPausedReserveSize));
+ ASSERT_TRUE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(1024L * 1024
* 32));
+ _backdate_paused_query(wg, config::spill_in_paused_queue_timeout_ms + 1);
+
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ ASSERT_EQ(_paused_query_count(wg), 0);
+}
+
+// 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) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 4);
+ ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ ASSERT_EQ(_paused_query_count(wg), 1);
+
+ _relieve_process_mem_limit();
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ ASSERT_EQ(_paused_query_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) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 4);
+ ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ _backdate_paused_query(wg, config::spill_in_paused_queue_timeout_ms + 1);
+
+ _wg_manager->handle_paused_queries();
+ _assert_cancelled_by_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) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 4);
+ ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_hard_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_cancelled_by_process_memory(query, true);
+}
+
+// Another workload group exceeds its min memory limit by more than 128 MiB,
but every query in
+// it is too small to be cancelled, so revoking memory from it frees nothing.
The paused query
+// must not be treated as if memory had been revoked (which would resume it
without any memory
+// being freed), it keeps waiting below the hard limit and is cancelled at the
hard limit.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_below_min_memory_with_non_reclaimable_peer) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 4);
+ ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+ auto peer_wg = _create_wg_with_min_memory(2);
+ std::vector<std::shared_ptr<QueryContext>> peer_queries;
+ for (int i = 0; i < 10; ++i) {
+ // Not larger than SMALL_MEMORY_TASK (32 MiB), so memory reclamation
skips it.
+ peer_queries.push_back(_create_query_with_memory(peer_wg, 1024L * 1024
* 30));
+ }
+ ASSERT_GT(peer_wg->total_mem_used(), peer_wg->min_memory_limit() + (1 <<
27));
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ for (int i = 0; i < 3; ++i) {
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ ASSERT_EQ(_paused_query_count(wg), 1);
+ ASSERT_FALSE(_wg_manager->revoking_memory_from_other_query_);
+ }
+
+ _exceed_process_hard_mem_limit();
+ _wg_manager->handle_paused_queries();
+ _assert_cancelled_by_process_memory(query, true);
+ for (const auto& peer_query : peer_queries) {
+ ASSERT_FALSE(peer_query->is_cancelled());
+ }
+}
+
+// Two workload groups exceed their min memory limit by more than 128 MiB. The
one that exceeds
+// most only holds queries too small to be cancelled, so revoking memory from
it frees nothing;
+// the other one holds a single cancellable query. The paused query must not
fall through to the
+// hard-limit fallback while the second workload group can still release
memory: its query is
+// cancelled and the paused query waits for the release.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_below_min_memory_tries_next_peer) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 4);
+ ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+ auto small_peer_wg = _create_wg_with_min_memory(2);
+ std::vector<std::shared_ptr<QueryContext>> small_peer_queries;
+ for (int i = 0; i < 10; ++i) {
+ // Not larger than SMALL_MEMORY_TASK (32 MiB), so memory reclamation
skips it.
+ small_peer_queries.push_back(_create_query_with_memory(small_peer_wg,
1024L * 1024 * 30));
+ }
+ auto large_peer_wg = _create_wg_with_min_memory(3);
+ auto large_peer_query = _create_query_with_memory(large_peer_wg, 1024L *
1024 * 250);
+ ASSERT_GT(large_peer_wg->total_mem_used(),
large_peer_wg->min_memory_limit() + (1 << 27));
+ // The non-reclaimable workload group exceeds its min memory most, so it
is tried first.
+ ASSERT_GT(small_peer_wg->total_mem_used() -
small_peer_wg->min_memory_limit(),
+ large_peer_wg->total_mem_used() -
large_peer_wg->min_memory_limit());
+
+ _exceed_process_hard_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ ASSERT_TRUE(large_peer_query->is_cancelled());
+ for (const auto& small_peer_query : small_peer_queries) {
+ ASSERT_FALSE(small_peer_query->is_cancelled());
+ }
+ _assert_still_paused(query);
+ ASSERT_EQ(_paused_query_count(wg), 1);
+ ASSERT_TRUE(_wg_manager->revoking_memory_from_other_query_);
+}
+
+// Revoking memory from another workload group cancels one of its queries,
which holds its
+// memory until the cancellation completes. While that cancellation is in
flight (within
+// `revoke_memory_max_tolerance_ms`), memory reclamation keeps reporting its
memory as revoked,
+// so the paused query waits for the release instead of being cancelled, also
at the hard limit.
+// Once the cancellation has taken longer than the tolerance, the peer no
longer counts,
+// nothing is revoked and the paused query falls through to the hard-limit
fallback.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_below_min_memory_waits_for_cancelling_peer) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 4);
+ ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+ auto peer_wg = _create_wg_with_min_memory(2);
+ auto peer_query = _create_query_with_memory(peer_wg, 1024L * 1024 * 300);
+ auto* peer_controller = _install_mock_query_task_controller(peer_query);
+ ASSERT_GT(peer_wg->total_mem_used(), peer_wg->min_memory_limit() + (1 <<
27));
+
+ _exceed_process_hard_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ // The peer query is cancelled, the paused query waits for its memory to
be released.
+ _wg_manager->handle_paused_queries();
+ ASSERT_TRUE(peer_query->is_cancelled());
+ _assert_still_paused(query);
+ ASSERT_TRUE(_wg_manager->revoking_memory_from_other_query_);
+
+ // The cancelled peer is not in the paused list, so the next round resumes
the paused
+ // query, whose reservation fails again and pauses it again.
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ ASSERT_FALSE(_wg_manager->revoking_memory_from_other_query_);
+ _pause_for_process_memory(query);
+
+ // The peer still holds its memory and its cancellation is in flight: it
is reported as
+ // revoked memory again, and the paused query keeps waiting for it.
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ ASSERT_TRUE(_wg_manager->revoking_memory_from_other_query_);
+
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ _pause_for_process_memory(query);
+
+ // The cancellation has taken longer than the tolerance: the peer no
longer counts, nothing
+ // is revoked and the paused query is cancelled at the hard limit.
+ peer_controller->set_cancelled_time(peer_controller->cancelled_time() -
+ config::revoke_memory_max_tolerance_ms
- 1);
+ _wg_manager->handle_paused_queries();
+ _assert_cancelled_by_process_memory(query, true);
+}
+
+// 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) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_hard_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_cancelled_by_process_memory(query, true);
+}
+
+// A query with a running task cannot be spilled, so it keeps waiting for the
task to yield.
+// The wait is still bounded: at the timeout the query is cancelled,
cancelling is safe while
+// the task runs.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_running_task_cancels_after_timeout) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+ auto* mock_controller = _install_mock_query_task_controller(query);
+ mock_controller->has_running_task_ = true;
+ // Revocable memory is not spilled while a task runs.
+ mock_controller->has_revocable_task_ = true;
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ for (int i = 0; i < 3; ++i) {
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ ASSERT_EQ(_paused_query_count(wg), 1);
+ ASSERT_EQ(mock_controller->revoke_memory_calls_, 0);
+ }
+
+ _backdate_paused_query(wg, config::spill_in_paused_queue_timeout_ms + 1);
+ _wg_manager->handle_paused_queries();
+ ASSERT_EQ(mock_controller->revoke_memory_calls_, 0);
+ _assert_cancelled_by_process_memory(query, false);
+ const auto status = query->exec_status().to_string();
+ ASSERT_NE(status.find("has running task: true"), std::string::npos) <<
status;
+}
+
+// At the hard limit a running task must not postpone the protection: spilling
is unsafe while
+// the task runs, so the query is cancelled at once.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_running_task_cancels_at_hard_limit) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+ auto* mock_controller = _install_mock_query_task_controller(query);
+ mock_controller->has_running_task_ = true;
+
+ _exceed_process_hard_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_cancelled_by_process_memory(query, true);
+ const auto status = query->exec_status().to_string();
+ ASSERT_NE(status.find("has running task: true"), std::string::npos) <<
status;
+}
+
+// A query with revocable memory is spilled rather than cancelled, also at the
hard limit. If
+// the spill resumes the query without freeing memory and it is paused again
without revocable
+// memory, the next round cancels it at the hard limit.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_spills_revocable_tasks_at_hard_limit) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+ auto* mock_controller = _install_mock_query_task_controller(query);
+ mock_controller->has_revocable_task_ = true;
+
+ _exceed_process_hard_mem_limit();
+ _pause_for_process_memory(query);
+ ASSERT_FALSE(query->get_memory_sufficient_dependency()->ready());
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ ASSERT_EQ(mock_controller->revoke_memory_calls_, 1);
+ // The spill completed and resumed the query: its memory dependency is
ready again, the
+ // paused reason is cleared and the entry is gone, so its tasks retry
their reservations.
+ _assert_resumed(query);
+ ASSERT_TRUE(query->get_memory_sufficient_dependency()->ready());
+ ASSERT_EQ(_paused_query_count(wg), 0);
+
+ // The spill freed nothing and the process is still at its hard limit, so
the retried
+ // reservation fails and the query is paused again, now without revocable
memory.
+ _pause_for_process_memory(query);
+ ASSERT_FALSE(query->get_memory_sufficient_dependency()->ready());
+ ASSERT_FALSE(mock_controller->has_revocable_task_);
+ _wg_manager->handle_paused_queries();
+ ASSERT_EQ(mock_controller->revoke_memory_calls_, 1);
+ _assert_cancelled_by_process_memory(query, true);
+}
+
+// The process memory pressure observed at the beginning of the round may be
relieved while the
+// manager inspects the pipeline tasks, up to the last scan for revocable
tasks. The timeout
+// decision must re-check the pressure after that scan and resume the query
instead of
+// cancelling it from the stale observation.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_recovery_before_timeout_decision_resumes) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+ auto* mock_controller = _install_mock_query_task_controller(query);
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ _backdate_paused_query(wg, config::spill_in_paused_queue_timeout_ms + 1);
+
+ // The pressure is relieved after the round has observed it, routed the
query and
+ // inspected its pipeline tasks for the last time.
+ mock_controller->on_get_revocable_tasks_ = []() {
_relieve_process_mem_limit(); };
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ ASSERT_EQ(_paused_query_count(wg), 0);
+}
+
+// Two tasks of one query fail process reservations of different sizes before
the next
+// maintenance round. The query is resumed as a whole, so the recorded
reservation must be the
+// largest pending one: once the smaller one fits, the query stays paused
(keeping its timer)
+// until the larger one fits too. The pending requests do not have to fit at
the same time: a
+// task releases its reservation after each block, so the tasks can reserve
one after another
+// once the largest request fits.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_keeps_largest_pending_reservation) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+ ASSERT_EQ(_recorded_reserve_size(wg), kProcessPausedReserveSize);
+ // A sibling task fails a larger reservation: the entry is shared and
keeps the larger size.
+ _wg_manager->add_paused_query(query->resource_ctx(), 1024L * 1024 * 64,
+
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+ ASSERT_EQ(_recorded_reserve_size(wg), 1024L * 1024 * 64);
+ // A later smaller failure does not lower it.
+ _wg_manager->add_paused_query(query->resource_ctx(), 1024L * 1024 * 4,
+
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+ ASSERT_EQ(_recorded_reserve_size(wg), 1024L * 1024 * 64);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+
+ // 16 MiB headroom: the 1 KiB and 4 MiB reservations fit, the 64 MiB one
does not. The
+ // query must not be woken up only to fail the 64 MiB reservation again.
+
MemInfo::set_soft_mem_limit_for_test(GlobalMemoryArbitrator::process_memory_usage()
+
+ 1024L * 1024 * 16);
+ ASSERT_FALSE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(1024L * 1024
* 4));
+ ASSERT_TRUE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(1024L * 1024
* 64));
+ for (int i = 0; i < 3; ++i) {
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ ASSERT_EQ(_paused_query_count(wg), 1);
+ }
+
+ // 65 MiB headroom: the largest pending reservation fits, although not all
of them do at the
+ // same time. The query is resumed; the tasks reserve one after another.
+
MemInfo::set_soft_mem_limit_for_test(GlobalMemoryArbitrator::process_memory_usage()
+
+ 1024L * 1024 * 65);
+ ASSERT_TRUE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(1024L * 1024
* 68));
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ ASSERT_EQ(_paused_query_count(wg), 0);
+}
+
+// Same two pending reservations near the timeout: the query keeps the entry
it was paused with
+// instead of being woken up by the smaller reservation and starting a fresh
wait, so the bounded
+// wait still ends with the cancellation.
+TEST_F(WorkloadGroupManagerTest,
+ process_mem_exceeded_largest_pending_reservation_cancels_after_timeout)
{
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+ _wg_manager->add_paused_query(query->resource_ctx(), 1024L * 1024 * 64,
+
Status::Error(ErrorCode::PROCESS_MEMORY_EXCEEDED, "test"));
+ ASSERT_EQ(_recorded_reserve_size(wg), 1024L * 1024 * 64);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+
+
MemInfo::set_soft_mem_limit_for_test(GlobalMemoryArbitrator::process_memory_usage()
+
+ 1024L * 1024 * 16);
+
ASSERT_FALSE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(kProcessPausedReserveSize));
+ ASSERT_TRUE(GlobalMemoryArbitrator::is_exceed_soft_mem_limit(1024L * 1024
* 64));
+ _backdate_paused_query(wg, config::spill_in_paused_queue_timeout_ms + 1);
+
+ _wg_manager->handle_paused_queries();
+ _assert_cancelled_by_process_memory(query, false);
+}
+
+// A query resumed because its recorded reservation fits is paused again when
the retry fails
+// (another query took the memory first, or a smaller sibling does not fit
next to the request
+// that was resumed for). The new entry continues the wait the query was
resumed from instead of
+// starting a fresh one, so the bounded wait still ends with the cancellation.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_failed_retry_continues_the_wait) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ _backdate_paused_query(wg, config::spill_in_paused_queue_timeout_ms + 1);
+
+ // The pressure is relieved and the query is resumed to retry its
reservation.
+ _relieve_process_mem_limit();
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ ASSERT_EQ(_paused_query_count(wg), 0);
+
+ // The retry fails before any reservation of the query succeeded: the
query has been waiting
+ // since its first failure and is cancelled at the timeout instead of
waiting another one.
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+ ASSERT_GT(_paused_elapsed_time(wg),
config::spill_in_paused_queue_timeout_ms);
+ _wg_manager->handle_paused_queries();
+ _assert_cancelled_by_process_memory(query, false);
+}
+
+// A successful reservation ends the wait: the query made progress, so a later
failure starts a
+// new bounded wait.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_successful_reservation_starts_new_wait) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 1024 * 128);
+ ASSERT_GT(wg->total_mem_used(), wg->min_memory_limit());
+
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ _backdate_paused_query(wg, config::spill_in_paused_queue_timeout_ms + 1);
+
+ _relieve_process_mem_limit();
+ _wg_manager->handle_paused_queries();
+ _assert_resumed(query);
+ ASSERT_EQ(_paused_query_count(wg), 0);
+
+ // The retried reservation succeeds
(PipelineTask::_try_to_reserve_memory), a later one fails.
+ query->resource_ctx()->task_controller()->end_process_memory_wait();
+ _exceed_process_soft_mem_limit();
+ _pause_for_process_memory(query);
+ ASSERT_LT(_paused_elapsed_time(wg),
config::spill_in_paused_queue_timeout_ms);
+ for (int i = 0; i < 3; ++i) {
+ _wg_manager->handle_paused_queries();
+ _assert_still_paused(query);
+ ASSERT_EQ(_paused_query_count(wg), 1);
+ }
+}
+
+// Two workload groups exceed their min memory limit by more than 128 MiB when
the walk starts.
+// The first one holds only a query whose cancellation has already taken
longer than
+// `revoke_memory_max_tolerance_ms`, so it releases nothing; while it is
scanned, the second one
+// drops back within its min memory. The second one must then be skipped from
its current usage
+// instead of having its remaining query cancelled for the excess it had when
the walk started.
+TEST_F(WorkloadGroupManagerTest,
process_mem_exceeded_below_min_memory_skips_peer_within_min) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 4);
+ ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+ auto first_peer_wg = _create_wg_with_min_memory(2);
+ auto first_peer_query = _create_query_with_memory(first_peer_wg, 1024L *
1024 * 300);
+ auto* first_peer_controller =
_install_mock_query_task_controller(first_peer_query);
+ first_peer_controller->cancel(Status::InternalError("memory gc cancel"));
+
first_peer_controller->set_cancelled_time(first_peer_controller->cancelled_time()
-
+
config::revoke_memory_max_tolerance_ms - 1);
+
+ auto second_peer_wg = _create_wg_with_min_memory(3);
+ auto second_peer_query = _create_query_with_memory(second_peer_wg, 1024L *
1024 * 250);
+ ASSERT_GT(second_peer_wg->total_mem_used(),
second_peer_wg->min_memory_limit() + (1 << 27));
+ // The first peer exceeds its min memory most, so it is scanned first.
+ ASSERT_GT(first_peer_wg->total_mem_used() -
first_peer_wg->min_memory_limit(),
+ second_peer_wg->total_mem_used() -
second_peer_wg->min_memory_limit());
+
+ // While the first peer is scanned, the second peer's query releases 180
MiB and the
+ // workload group drops to 70 MiB, within its 100 MiB min memory; its
query is still larger
+ // than SMALL_MEMORY_TASK, so it could be cancelled. Nothing refreshes the
workload group
+ // usage meanwhile, so its cached usage still shows the 250 MiB of the
snapshot.
+ bool released = false;
+ first_peer_controller->on_is_cancelled_ = [&]() {
+ if (released) {
+ return;
+ }
+ released = true;
+ _release_query_memory(second_peer_query, 1024L * 1024 * 180);
+ ASSERT_GT(second_peer_wg->total_mem_used(),
second_peer_wg->min_memory_limit());
+ };
+
+ _exceed_process_hard_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ ASSERT_TRUE(released);
+ ASSERT_FALSE(second_peer_query->is_cancelled());
+ // The walk refreshed the second peer's usage before deciding on it.
+ ASSERT_EQ(second_peer_wg->total_mem_used(), 1024L * 1024 * 70);
+ ASSERT_EQ(second_peer_wg->min_memory_limit(), 1024L * 1024 * 100);
+ // Nothing could be revoked, the paused query falls through to the
hard-limit fallback.
+ _assert_cancelled_by_process_memory(query, true);
+}
+
+// Two workload groups exceed their min memory limit by more than 128 MiB. The
first one holds
+// only a query whose cancellation has already taken longer than
`revoke_memory_max_tolerance_ms`,
+// so it releases nothing; while it is scanned, the process memory pressure is
relieved (other
+// queries or the cache released memory). The second one must not have its
query cancelled for
+// pressure that is gone, and the paused query must be resumed instead of
being cancelled at the
+// hard limit it was observed under before the walk.
+TEST_F(WorkloadGroupManagerTest,
+
process_mem_exceeded_below_min_memory_resumes_when_relieved_during_peer_scan) {
+ auto wg = _create_wg_with_min_memory(1);
+ auto query = _create_query_with_memory(wg, 1024L * 4);
+ ASSERT_LE(wg->total_mem_used(), wg->min_memory_limit());
+
+ auto first_peer_wg = _create_wg_with_min_memory(2);
+ auto first_peer_query = _create_query_with_memory(first_peer_wg, 1024L *
1024 * 300);
+ auto* first_peer_controller =
_install_mock_query_task_controller(first_peer_query);
+ first_peer_controller->cancel(Status::InternalError("memory gc cancel"));
+
first_peer_controller->set_cancelled_time(first_peer_controller->cancelled_time()
-
+
config::revoke_memory_max_tolerance_ms - 1);
+
+ auto second_peer_wg = _create_wg_with_min_memory(3);
+ auto second_peer_query = _create_query_with_memory(second_peer_wg, 1024L *
1024 * 250);
+ ASSERT_GT(second_peer_wg->total_mem_used(),
second_peer_wg->min_memory_limit() + (1 << 27));
+ // The first peer exceeds its min memory most, so it is scanned first.
+ ASSERT_GT(first_peer_wg->total_mem_used() -
first_peer_wg->min_memory_limit(),
+ second_peer_wg->total_mem_used() -
second_peer_wg->min_memory_limit());
+
+ bool relieved = false;
+ first_peer_controller->on_is_cancelled_ = [&]() {
+ if (relieved) {
+ return;
+ }
+ relieved = true;
+ _relieve_process_mem_limit();
+ };
+
+ _exceed_process_hard_mem_limit();
+ _pause_for_process_memory(query);
+
+ config::spill_in_paused_queue_timeout_ms = 60 * 1000;
+ _wg_manager->handle_paused_queries();
+ ASSERT_TRUE(relieved);
+ ASSERT_FALSE(second_peer_query->is_cancelled());
+ ASSERT_FALSE(_wg_manager->revoking_memory_from_other_query_);
+ _assert_resumed(query);
+ ASSERT_EQ(_paused_query_count(wg), 0);
+}
+
} // namespace doris
diff --git a/be/test/testutil/mock/mock_query_task_controller.h
b/be/test/testutil/mock/mock_query_task_controller.h
index a9c6886aada..96c4a2d1fae 100644
--- a/be/test/testutil/mock/mock_query_task_controller.h
+++ b/be/test/testutil/mock/mock_query_task_controller.h
@@ -16,6 +16,9 @@
// under the License.
#pragma once
+#include <functional>
+#include <vector>
+
#include "runtime/workload_management/query_task_controller.h"
namespace doris {
@@ -35,6 +38,58 @@ struct MockQueryTaskController : public QueryTaskController {
}
void set_cancelled_time(int64_t ctime) { cancelled_time_ = ctime; }
+
+ // Memory reclamation asks every candidate query whether it is cancelled
while it scans a
+ // workload group, so this is where a test can let another query release
memory during that
+ // scan.
+ bool is_cancelled() const override {
+ if (on_is_cancelled_) {
+ on_is_cancelled_();
+ }
+ return QueryTaskController::is_cancelled();
+ }
+
+ // Pipeline state seen by WorkloadGroupMgr::handle_single_query_. The
query context of a
+ // unit test has no fragments, so the real implementation never reports a
running or a
+ // revocable task; these knobs simulate them.
+ // NOLINTNEXTLINE(readability-make-member-function-const): overrides a
non-const virtual.
+ void get_revocable_info(size_t* revocable_size, size_t* memory_usage,
+ bool* has_running_task) override {
+ QueryTaskController::get_revocable_info(revocable_size, memory_usage,
has_running_task);
+ *has_running_task = has_running_task_;
+ }
+
+ // The manager only checks whether the list is empty before it calls
revoke_memory(), which
+ // is mocked below, so a placeholder entry is enough to stand for a
revocable task.
+ std::vector<PipelineTask*> get_revocable_tasks() override {
+ if (on_get_revocable_tasks_) {
+ on_get_revocable_tasks_();
+ }
+ return has_revocable_task_ ? std::vector<PipelineTask*> {nullptr}
+ : std::vector<PipelineTask*> {};
+ }
+
+ // The spill of a unit test completes at once: like the spill callback of
the real
+ // implementation, resume the query so that its blocked tasks retry their
reservations. The
+ // spilled tasks have nothing left to revoke.
+ Status revoke_memory() override {
+ ++revoke_memory_calls_;
+ RETURN_IF_ERROR(revoke_memory_status_);
+ has_revocable_task_ = false;
+ set_memory_sufficient(true);
+ return Status::OK();
+ }
+
+ bool has_running_task_ {false};
+ bool has_revocable_task_ {false};
+ Status revoke_memory_status_ {Status::OK()};
+ int revoke_memory_calls_ {0};
+ // Runs inside get_revocable_tasks(), the last inspection of the pipeline
tasks after the
+ // manager has observed the memory pressure and before it decides what to
do with the query.
+ std::function<void()> on_get_revocable_tasks_;
+ // Runs inside is_cancelled(), while memory reclamation scans the workload
group of this
+ // query.
+ std::function<void()> on_is_cancelled_;
};
} // namespace doris
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
