[improvement](be) Name the task of every resource context still alive at graceful exit

be/src/runtime/workload_management/query_task_controller.h

@@ -53,6 +53,8 @@ class QueryTaskController : public TaskController {
     void inc_finished_task_num();
     int get_total_task_num() const;
     int get_finished_task_num() const;
+    // Whether the QueryContext this controller belongs to has not been destroyed yet.
+    bool is_query_context_alive() const { return !query_ctx_.expired(); }
 
 protected:
     QueryTaskController(const std::shared_ptr<QueryContext>& query_ctx) : query_ctx_(query_ctx) {}

be/src/runtime/workload_management/resource_context.cpp

@@ -17,16 +17,152 @@
 
 #include "runtime/workload_management/resource_context.h"
 
+#include <fmt/format.h>
 #include <gen_cpp/FrontendService_types.h>
 #include <gen_cpp/data.pb.h>
 #include <glog/logging.h>
 
+#include <cstdint>
+#include <iostream>
+#include <mutex>
+#include <thread>
+#include <unordered_set>
+#include <vector>
+
+#include "common/check.h"
+#include "common/config.h"
+#include "runtime/memory/mem_tracker_limiter.h"
 #include "runtime/workload_group/workload_group.h"
 #include "runtime/workload_management/query_task_controller.h"
+#include "util/pretty_printer.h"
 #include "util/time.h"
+#include "util/uid_util.h"
 
 namespace doris {
 
+namespace {
+
+// The resource contexts created while memleak-check mode is on. Their addresses are kept
+// disguised (bitwise inverted), so that LeakSanitizer does not take the registry for a
+// reference: a context that nothing else references must still be reported as leaked.
+struct LiveResourceContexts {
+    std::mutex lock;
+    std::unordered_set<uintptr_t> disguised_contexts;
+};
+
+LiveResourceContexts& live_resource_contexts() {
+    // Never destroyed, so that a context released during static destruction can still
+    // unregister itself.
+    static auto* live = new LiveResourceContexts();
+    return *live;
+}
+
+uintptr_t disguise(const ResourceContext* context) {
+    return ~reinterpret_cast<uintptr_t>(context);
+}
+
+ResourceContext* undisguise(uintptr_t disguised_context) {
+    return reinterpret_cast<ResourceContext*>(~disguised_context);
+}
+
+} // namespace
+
+void ResourceContext::_track_if_leak_check_enabled() {
+    if (!config::enable_graceful_exit_check) {
+        return;
+    }
+    auto& live = live_resource_contexts();
+    std::lock_guard<std::mutex> l(live.lock);
+    live.disguised_contexts.insert(disguise(this));
+    _tracked = true;
+}
+
+ResourceContext::~ResourceContext() {
+    if (_tracked) {
+        auto& live = live_resource_contexts();
+        std::lock_guard<std::mutex> l(live.lock);
+        DORIS_CHECK(live.disguised_contexts.erase(disguise(this)) == 1);
+    }
+}
+
+std::string ResourceContext::describe_live_contexts(size_t max_contexts) {
+    // Meant for shutdown, when no context is being created any more: a context that is still
+    // inside std::make_shared has not been handed its weak_from_this() yet.
+    //
+    // Describe the contexts after the registry is unlocked, keeping each one alive meanwhile
+    // with a reference of our own. A context whose last reference is being released right now
+    // cannot be referenced any more and is left out.
+    std::vector<std::shared_ptr<ResourceContext>> contexts;
+    {
+        auto& live = live_resource_contexts();
+        std::lock_guard<std::mutex> l(live.lock);
+        contexts.reserve(live.disguised_contexts.size());
+        for (uintptr_t disguised_context : live.disguised_contexts) {
+            if (auto reference = undisguise(disguised_context)->weak_from_this().lock()) {
+                contexts.push_back(std::move(reference));
+            }
+        }
+    }
+    if (contexts.empty()) {
+        return "";
+    }
+    fmt::memory_buffer out;
+    fmt::format_to(std::back_inserter(out),
+                   "{} resource contexts are still alive. LeakSanitizer reports a leaked context "
+                   "as the memory block that holds it, so match the addresses below against the "
+                   "ranges it prints under 'Objects leaked above'.",
+                   contexts.size());
+    for (size_t i = 0; i < contexts.size() && i < max_contexts; ++i) {
+        // The reference taken above is not counted.
+        fmt::format_to(std::back_inserter(out), "\n  {}",
+                       contexts[i]->_describe(contexts[i].use_count() - 1));
+    }
+    if (contexts.size() > max_contexts) {
+        fmt::format_to(std::back_inserter(out), "\n  ... and {} more",
+                       contexts.size() - max_contexts);
+    }
+    return fmt::to_string(out);
+}
+
+std::string ResourceContext::log_live_contexts() {
+    std::string live_contexts;
+    std::thread describer([&live_contexts]() { live_contexts = describe_live_contexts(100); });
+    describer.join();
+    if (live_contexts.empty()) {
+        LOG(INFO) << "No resource context is alive after shutdown";
+    } else {
+        LOG(WARNING) << live_contexts;
+        std::cerr << live_contexts << '\n';
+    }
+    return live_contexts;
+}
+
+std::string ResourceContext::_describe(long references_held_elsewhere) const {
+    std::string mem_tracker = "none";
+    if (auto tracker = memory_context_->mem_tracker()) {
+        mem_tracker = fmt::format("{} (consumption={}, peak={})", tracker->label(),
+                                  PrettyPrinter::print_bytes(tracker->consumption()),
+                                  PrettyPrinter::print_bytes(tracker->peak_consumption()));
+    }
+    std::string workload_group = "none";
+    if (_workload_group != nullptr) {
+        workload_group = fmt::format("{}({})", _workload_group->name(), _workload_group->id());
+    }
+    // A query's context outlives its QueryContext when something else still holds it.
+    std::string query_context;
+    if (auto* query_task_controller = dynamic_cast<QueryTaskController*>(task_controller_.get())) {
+        query_context = query_task_controller->is_query_context_alive()
+                                ? ", query_context=alive"
+                                : ", query_context=destroyed";
+    }
+    return fmt::format(
+            "{}: references={}, task_id={}, finished={}, running_ms={}, mem_tracker={}, "
+            "workload_group={}{}",
+            fmt::ptr(this), references_held_elsewhere, print_id(task_controller_->task_id()),
+            task_controller_->is_finished(), task_controller_->running_time(), mem_tracker,
+            workload_group, query_context);
+}
+
 void ResourceContext::to_thrift_query_statistics(TQueryStatistics* statistics) const {
     DCHECK(statistics != nullptr);
     statistics->__set_scan_rows(io_context()->scan_rows());

be/src/runtime/workload_management/resource_context.h

@@ -17,7 +17,9 @@
 
 #pragma once
 
+#include <cstddef>
 #include <memory>
+#include <string>
 
 #include "common/factory_creator.h"
 #include "common/multi_version.h"
@@ -51,8 +53,9 @@ class ResourceContext : public std::enable_shared_from_this<ResourceContext> {
         memory_context_->set_resource_ctx(this);
         io_context_->set_resource_ctx(this);
         task_controller_->set_resource_ctx(this);
+        _track_if_leak_check_enabled();
     }
-    ~ResourceContext() = default;
+    ~ResourceContext();
 
     // Only return the raw pointer to the caller, so that the caller should not save it to other variables.
     CPUContext* cpu_context() const { return cpu_context_.get(); }
@@ -99,7 +102,27 @@ class ResourceContext : public std::enable_shared_from_this<ResourceContext> {
         resource_profile_.set(std::move(resource_profile));
     }
 
+    // Memleak-check mode only (config::enable_graceful_exit_check, where doris_main destroys
+    // everything on exit so that LeakSanitizer can check the heap): returns one line per
+    // resource context that is still alive, at most `max_contexts` of them, or an empty string
+    // if there is none. LeakSanitizer reports a context that was never released only by its
+    // allocation stack, which is the same for every query; the task id printed here tells which
+    // query or load it belonged to. Call it at shutdown, when no context is being created.
+    static std::string describe_live_contexts(size_t max_contexts);
+
+    // For the end of a graceful exit in memleak-check mode: logs describe_live_contexts() to
+    // be.WARNING and to stderr (be.out, where the LeakSanitizer report follows), or one INFO line
+    // if no context is alive, and returns what it logged. The description is built on a thread
+    // of its own and the caller never touches a context: a pointer to a leaked context left on
+    // the caller's stack would be found by LeakSanitizer, which would then take the context for
+    // referenced and stop reporting it, while the stack of a finished thread is not scanned.
+    static std::string log_live_contexts();
+
 private:
+    // Registers this context for describe_live_contexts() when memleak-check mode is on.
+    void _track_if_leak_check_enabled();
+    std::string _describe(long references_held_elsewhere) const;
+
     // The controller's init value is nullptr, it means the resource context will ignore this controller.
     std::unique_ptr<CPUContext> cpu_context_ = nullptr;
     std::unique_ptr<MemoryContext> memory_context_ = nullptr;
@@ -108,6 +131,8 @@ class ResourceContext : public std::enable_shared_from_this<ResourceContext> {
 
     WorkloadGroupPtr _workload_group = nullptr;
     MultiVersion<RuntimeProfile> resource_profile_;
+    // Whether this context is registered for describe_live_contexts().
+    bool _tracked = false;
 };
 
 } // namespace doris

be/src/service/doris_main.cpp

@@ -76,6 +76,7 @@
 #include "load/stream_load/stream_load_recorder_manager.h"
 #include "runtime/exec_env.h"
 #include "runtime/user_function_cache.h"
+#include "runtime/workload_management/resource_context.h"
 #include "service/arrow_flight/flight_sql_service.h"
 #include "service/backend_options.h"
 #include "service/backend_service.h"
@@ -863,6 +864,11 @@ int main(int argc, char** argv) {
     service.reset();
     LOG(INFO) << "Backend Service stopped";
     exec_env->destroy();
+    // Everything above should have released the resource context of every query and load. One
+    // that is still alive is held either by a thread that has not finished or by a reference
+    // that is never released; LeakSanitizer reports the latter after main() returns, but cannot
+    // tell which task it belonged to, so name the task here.
+    doris::ResourceContext::log_live_contexts();
     LOG(INFO) << "All service stopped, doris main exited.";
     return 0;
 }