[Refactor] Move FragmentContext into ExecRuntime (#74774)

Signed-off-by: alvin-phoenix-ai <alvin.zhao@phoenixdata.ai>

Author: alvin-phoenix-ai

be/AGENTS.md

@@ -233,7 +233,7 @@ Primitive execution contracts, DataSink base contract, runtime-filter infrastruc
 ### ExecRuntime (`execruntime`)
 Execution runtime driver lifecycle and scheduling-adjacent behavior above ComputeEnv and ExecPrimitive without concrete operators, storage, service, connector, cache, or broad Exec coupling.
 - Targets: `ExecRuntime`
-- Allowed internal include prefixes: `compute_env/`, `exec/runtime/`, `exec/pipeline/group_execution/execution_group.h`, `exec/pipeline/group_execution/execution_group_builder.h`, `exec/pipeline/group_execution/execution_group_fwd.h`, `exec/pipeline/adaptive/adaptive_dop_param.h`, `exec/pipeline/pipeline_driver.h`, `exec/pipeline/pipeline_fwd.h`, `exec/pipeline/operator.h`, `exec/pipeline/operator_with_dependency.h`, `exec/pipeline/source_operator.h`, `exec/pipeline/primitives/`, `exec/pipeline/schedule/common.h`, `exec/pipeline/schedule/event_scheduler.h`, `exec/pipeline/schedule/pipeline_driver_observer.h`, `exec/pipeline/schedule/utils.h`, `exec/pipeline/runtime_filter_hub.h`, `exec/pipeline/runtime_filter_core_types.h`, `exec/pipeline/scan/morsel_queue.h`, `exec/pipeline/scan/split_morsel_ticket_checker.h`, `exec/pipeline/scan/ticketed_morsel_queue.h`, `exec/runtime_filter/`, `storage/primitive/`, `exprs/`, `serde/`, `runtime/env/`, `runtime/`, `platform/`, `fs/`, `io/`, `column/`, `types/`, `common/`, `base/`, `gutil/`, `gen_cpp/`
+- Allowed internal include prefixes: `compute_env/`, `exec/runtime/`, `exec/pipeline/fragment_context.h`, `exec/pipeline/group_execution/execution_group.h`, `exec/pipeline/group_execution/execution_group_builder.h`, `exec/pipeline/group_execution/execution_group_fwd.h`, `exec/pipeline/adaptive/adaptive_dop_param.h`, `exec/pipeline/pipeline_driver.h`, `exec/pipeline/pipeline_fwd.h`, `exec/pipeline/operator.h`, `exec/pipeline/operator_with_dependency.h`, `exec/pipeline/source_operator.h`, `exec/pipeline/primitives/`, `exec/pipeline/schedule/common.h`, `exec/pipeline/schedule/event_scheduler.h`, `exec/pipeline/schedule/pipeline_driver_observer.h`, `exec/pipeline/schedule/utils.h`, `exec/pipeline/runtime_filter_hub.h`, `exec/pipeline/runtime_filter_core_types.h`, `exec/pipeline/scan/morsel_queue.h`, `exec/pipeline/scan/split_morsel_ticket_checker.h`, `exec/pipeline/scan/ticketed_morsel_queue.h`, `exec/runtime_filter/`, `storage/primitive/`, `exprs/`, `serde/`, `runtime/env/`, `runtime/`, `platform/`, `fs/`, `io/`, `column/`, `types/`, `common/`, `base/`, `gutil/`, `gen_cpp/`
 - Allowed target deps: `ComputeEnv`, `ExecPrimitive`, `StoragePrimitive`, `ExprCore`, `Serde`, `RuntimeEnv`, `RuntimeCore`, `Platform`, `FSCore`, `IO`, `ChunkCore`, `ColumnCore`, `Types`, `Common`, `Base`, `Gutil`, `StarRocksGen`
 - Core tests: `exec_runtime_test`
 - Remediation: Keep ExecRuntime limited to driver lifecycle and runtime behavior that can be expressed through ComputeEnv and ExecPrimitive contracts; move concrete operators, storage, service, connector, cache, HTTP, and broad Exec integration upward.

be/module_boundary_manifest.json

@@ -542,6 +542,7 @@
       "allowed_include_prefixes": [
         "compute_env/",
         "exec/runtime/",
+        "exec/pipeline/fragment_context.h",
         "exec/pipeline/group_execution/execution_group.h",
         "exec/pipeline/group_execution/execution_group_builder.h",
         "exec/pipeline/group_execution/execution_group_fwd.h",
@@ -579,7 +580,6 @@
       ],
       "forbidden_include_prefixes": ["exec/", "storage/", "service/", "connector/", "cache/", "http/", "agent/", "util/"],
       "forbidden_includes": [
-        "exec/pipeline/fragment_context.h",
         "exec/pipeline/pipeline_driver_executor.h",
         "exec/pipeline/pipeline_driver_queue.h",
         "exec/pipeline/pipeline_driver_poller.h",

be/src/exec/CMakeLists.txt

@@ -159,7 +159,6 @@ set(PIPELINE_RUNTIME_SRCS
     pipeline/pipeline_driver_poller.cpp
     pipeline/audit_statistics_reporter.cpp
     pipeline/exec_state_reporter.cpp
-    pipeline/fragment_context.cpp
     pipeline/fragment_context_cancel.cpp
     pipeline/fragment_context_manager.cpp
     pipeline/query_context.cpp
@@ -507,6 +506,7 @@ ADD_BE_LIB(ExecRuntime
     pipeline/group_execution/execution_group_builder.cpp
     pipeline/pipeline.cpp
     pipeline/pipeline_driver.cpp
+    pipeline/fragment_context.cpp
     pipeline/schedule/event_scheduler.cpp
     pipeline/schedule/pipeline_driver_observer.cpp
 )

be/src/exec/pipeline/fragment_context.cpp

@@ -44,8 +44,8 @@
 #include "exec/runtime/query_runtime_state.h"
 #include "platform/query_timeout_hook.h"
 #include "platform/thrift_rpc_helper.h"
-#include "runtime/exec_env.h"
 #include "runtime/fragment_attachment.h"
+#include "runtime/service_contexts.h"
 
 namespace starrocks::pipeline {
 
@@ -276,9 +276,10 @@ void FragmentContext::set_final_status(const Status& status) {
 
         const bool finished_cancel = detailed_message == "QueryFinished" || detailed_message == "LimitReach";
         if (!s_status.ok() && !finished_cancel) {
-            const auto* executors = workgroup() != nullptr
-                                            ? workgroup()->executors()
-                                            : ExecEnv::GetInstance()->workgroup_manager()->shared_executors();
+            const auto* executors =
+                    workgroup() != nullptr
+                            ? workgroup()->executors()
+                            : execution_services(_runtime_state.get()).workgroup_manager->shared_executors();
             auto* executor = executors->driver_executor();
             auto* query_ctx = _runtime_state->query_ctx();
             if (query_ctx != nullptr) {
@@ -340,10 +341,6 @@ Status FragmentContext::prepare_all_pipelines() {
     return Status::OK();
 }
 
-void FragmentContext::_set_default_workgroup() {
-    set_workgroup(ExecEnv::GetInstance()->workgroup_manager()->get_default_workgroup());
-}
-
 void FragmentContext::set_fragment_attachments(std::vector<std::unique_ptr<FragmentAttachment>>&& attachments) {
     _fragment_attachments = std::move(attachments);
 }

be/src/exec/pipeline/fragment_context.h

@@ -188,7 +188,6 @@ class FragmentContext : public ExecutionGroupLifecycle {
     void init_event_scheduler();
 
 private:
-    void _set_default_workgroup();
     void _close_fragment_attachments();
 
     bool _enable_group_execution = false;

be/src/exec/pipeline/fragment_context_manager.cpp

@@ -44,23 +44,6 @@ const RuntimeServices& runtime_services(const RuntimeState* state) {
 FragmentContextManager::FragmentContextManager(FragmentLifecycleWeakPtr fragment_lifecycle)
         : _fragment_lifecycle(std::move(fragment_lifecycle)) {}
 
-FragmentContext* FragmentContextManager::get_or_register(const TUniqueId& fragment_id) {
-    std::lock_guard<std::mutex> lock(_lock);
-    auto it = _fragment_contexts.find(fragment_id);
-    if (it != _fragment_contexts.end()) {
-        return it->second.get();
-    } else {
-        auto&& ctx = std::make_unique<FragmentContext>();
-        auto* raw_ctx = ctx.get();
-        if (!_fragment_lifecycle.expired()) {
-            raw_ctx->set_fragment_lifecycle(_fragment_lifecycle);
-        }
-        _fragment_contexts.emplace(fragment_id, std::move(ctx));
-        raw_ctx->_set_default_workgroup();
-        return raw_ctx;
-    }
-}
-
 Status FragmentContextManager::register_ctx(const TUniqueId& fragment_id, FragmentContextPtr fragment_ctx) {
     std::lock_guard<std::mutex> lock(_lock);
 

be/src/exec/pipeline/fragment_context_manager.h

@@ -36,7 +36,6 @@ class FragmentContextManager {
     FragmentContextManager& operator=(const FragmentContextManager&) = delete;
     FragmentContextManager& operator=(FragmentContextManager&&) = delete;
 
-    FragmentContext* get_or_register(const TUniqueId& fragment_id);
     FragmentContextPtr get(const TUniqueId& fragment_id);
 
     Status register_ctx(const TUniqueId& fragment_id, FragmentContextPtr fragment_ctx);

be/src/exec/pipeline/group_execution/execution_group.cpp

@@ -91,8 +91,20 @@ void ExecutionGroup::prepare_active_drivers_parallel(const std::shared_ptr<Runti
     DCHECK(pipeline_prepare_pool != nullptr);
 
     for_each_active_driver(_pipelines, [&](const DriverPtr& driver) {
-        // Since prepare is async, hold both the runtime state and driver until the task finishes.
-        bool submitted = pipeline_prepare_pool->try_offer([sync_ctx, driver, runtime_state_holder = state]() {
+        // Capture the runtime state by value to keep it alive across the async prepare, but capture
+        // `driver` by reference -- the worker must NOT own the driver. The reference binds to the
+        // stable element in `pipeline->drivers()`, which the wait barrier below keeps alive until
+        // every task finishes; the worker only touches it inside prepare_local_state().
+        //
+        // Do NOT capture the driver by value (shared_ptr). The waiter wakes as soon as pending_tasks
+        // hits 0 and, on prepare error, tears down the fragment -- destroying each Pipeline and its
+        // operator factories. Some operators (e.g. OlapScanOperator) hold a context whose buffer is
+        // owned by an operator factory (OlapScanContext -> OlapScanContextFactory::_chunk_buffer). If
+        // the worker owned the driver, it could become the last owner and run ~PipelineDriver on the
+        // prepare thread AFTER those factories were freed, closing a dangling buffer (use-after-free).
+        // Keeping the worker non-owning ensures the driver is destroyed only by ~Pipeline, in the safe
+        // _drivers-before-_op_factories order, on the fragment thread.
+        bool submitted = pipeline_prepare_pool->try_offer([sync_ctx, &driver, runtime_state_holder = state]() {
             auto runtime_state = runtime_state_holder.get();
             // make sure mem tracker is instance level
             auto mem_tracker = runtime_state->instance_mem_tracker();

be/test/exec/CMakeLists.txt

@@ -59,6 +59,7 @@ set(EXEC_RUNTIME_TEST_FILES
     pipeline_driver_exec_runtime_test.cpp
     pipeline_event_scheduler_exec_runtime_test.cpp
     pipeline_execution_group_exec_runtime_test.cpp
+    pipeline_fragment_context_exec_runtime_test.cpp
     pipeline_fragment_runtime_state_exec_runtime_test.cpp
     pipeline_query_runtime_state_exec_runtime_test.cpp
 )

be/test/exec/pipeline/pipeline_file_scan_node_test.cpp

@@ -21,6 +21,7 @@
 #include <utility>
 
 #include "base/testutil/assert.h"
+#include "base/uid_util.h"
 #include "base/utility/defer_op.h"
 #include "column/chunk.h"
 #include "column/column_helper.h"
@@ -73,6 +74,9 @@ class PipeLineFileScanNodeTest : public ::testing::Test {
         _exec_env = ExecEnv::GetInstance();
         _exec_env->process_metrics_registry()->root_registry()->set_collect_hook_enabled(true);
 
+        _request.params.query_id = generate_uuid();
+        _request.params.fragment_instance_id = generate_uuid();
+
         const auto& params = _request.params;
         const auto& query_id = params.query_id;
         const auto& fragment_id = params.fragment_instance_id;
@@ -86,15 +90,17 @@ class PipeLineFileScanNodeTest : public ::testing::Test {
         _query_ctx->init_mem_tracker(GlobalEnv::GetInstance()->query_pool_mem_tracker()->limit(),
                                      GlobalEnv::GetInstance()->query_pool_mem_tracker());
 
-        _fragment_ctx = _query_ctx->fragment_mgr()->get_or_register(fragment_id);
+        auto fragment_ctx = std::make_shared<FragmentContext>();
+        _fragment_ctx = fragment_ctx.get();
         _fragment_ctx->set_query_id(query_id);
         _fragment_ctx->set_fragment_instance_id(fragment_id);
+        _fragment_ctx->set_workgroup(_exec_env->workgroup_manager()->get_default_workgroup());
         _fragment_ctx->set_runtime_state(std::make_unique<RuntimeState>(
                 _request.params.query_id, _request.params.fragment_instance_id, _request.query_options,
                 _request.query_globals, &_exec_env->query_execution_services(), _exec_env));
-        _fragment_ctx->set_workgroup(ExecEnv::GetInstance()->workgroup_manager()->get_default_workgroup());
 
         _fragment_future = _fragment_ctx->finish_future();
+        ASSERT_OK(_query_ctx->fragment_mgr()->register_ctx(fragment_id, std::move(fragment_ctx)));
         _runtime_state = _fragment_ctx->runtime_state();
 
         _runtime_state->set_chunk_size(config::vector_chunk_size);