// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

#pragma once

#include <brpc/closure_guard.h>

#include <gen_cpp/FrontendService_types.h>

#include <gen_cpp/QueryPlanExtra_types.h>

#include <gen_cpp/Types_types.h>

#include <gen_cpp/types.pb.h>

#include <cstdint>

#include <functional>

#include <iosfwd>

#include <map>

#include <memory>

#include <mutex>

#include <set>

#include <string>

#include <unordered_map>

#include <vector>

#include "common/be_mock_util.h"

#include "common/metrics/metrics.h"

#include "common/status.h"

#include "exec/runtime_filter/runtime_filter_mgr.h"

#include "runtime/query_context.h"

#include "service/http/rest_monitor_iface.h"

#include "util/countdown_latch.h"

#include "util/hash_util.hpp" // IWYU pragma: keep

namespace butil {

class IOBufAsZeroCopyInputStream;

}

namespace doris {

#include "common/compile_check_begin.h"

extern bvar::Adder<uint64_t> g_fragment_executing_count;

extern bvar::Status<uint64_t> g_fragment_last_active_time;

class PipelineFragmentContext;

class QueryContext;

class ExecEnv;

class ThreadPool;

class PExecPlanFragmentStartRequest;

class PMergeFilterRequest;

class RuntimeProfile;

class RuntimeState;

class TPipelineFragmentParams;

class TPipelineInstanceParams;

class TScanColumnDesc;

class TScanOpenParams;

class Thread;

class WorkloadQueryInfo;

std::string to_load_error_http_path(const std::string& file_name);

template <typename Key, typename Value, typename ValueType>

class ConcurrentContextMap {

public:

    using ApplyFunction = std::function<Status(phmap::flat_hash_map<Key, Value>&)>;

    ConcurrentContextMap();

    Value find(const Key& query_id);

    void insert(const Key& query_id, std::shared_ptr<ValueType>);

    void clear();

    bool erase(const Key& query_id);

    size_t num_items() const {

        size_t n = 0;

        for (auto& pair : _internal_map) {

            std::shared_lock lock(*pair.first);

            auto& map = pair.second;

            n += map.size();

        }

        return n;

    }

_ZNK5doris20ConcurrentContextMapINS_9TUniqueIdESt8weak_ptrINS_12QueryContextEES3_E9num_itemsEv

Line

Count

Source

81

698

    size_t num_items() const {

82

698

        size_t n = 0;

83

89.3k

        for (auto& pair : _internal_map) {

84

89.3k

            std::shared_lock lock(*pair.first);

85

89.3k

            auto& map = pair.second;

86

89.3k

            n += map.size();

87

89.3k

        }

88

698

        return n;

89

698

    }

_ZNK5doris20ConcurrentContextMapISt4pairINS_9TUniqueIdEiESt10shared_ptrINS_23PipelineFragmentContextEES5_E9num_itemsEv

Line

Count

Source

81

72

    size_t num_items() const {

82

72

        size_t n = 0;

83

9.21k

        for (auto& pair : _internal_map) {

84

9.21k

            std::shared_lock lock(*pair.first);

85

9.21k

            auto& map = pair.second;

86

9.21k

            n += map.size();

87

9.21k

        }

88

72

        return n;

89

72

    }

    void apply(ApplyFunction&& function) {

        for (auto& pair : _internal_map) {

            // TODO: Now only the cancel worker do the GC the _query_ctx_map. each query must

            // do erase the finish query unless in _query_ctx_map. Rethink the logic is ok

            std::unique_lock lock(*pair.first);

            static_cast<void>(function(pair.second));

        }

    }

_ZN5doris20ConcurrentContextMapISt4pairINS_9TUniqueIdEiESt10shared_ptrINS_23PipelineFragmentContextEES5_E5applyEOSt8functionIFNS_6StatusERN5phmap13flat_hash_mapIS3_S6_NSA_4HashIS3_EENSA_7EqualToIS3_EESaIS1_IKS3_S6_EEEEEE

Line

Count

Source

90

6.22k

    void apply(ApplyFunction&& function) {

91

796k

        for (auto& pair : _internal_map) {

92

            // TODO: Now only the cancel worker do the GC the _query_ctx_map. each query must

93

            // do erase the finish query unless in _query_ctx_map. Rethink the logic is ok

94

796k

            std::unique_lock lock(*pair.first);

95

796k

            static_cast<void>(function(pair.second));

96

796k

        }

97

6.22k

    }

_ZN5doris20ConcurrentContextMapINS_9TUniqueIdESt8weak_ptrINS_12QueryContextEES3_E5applyEOSt8functionIFNS_6StatusERN5phmap13flat_hash_mapIS1_S4_NS8_4HashIS1_EENS8_7EqualToIS1_EESaISt4pairIKS1_S4_EEEEEE

Line

Count

Source

90

24.4k

    void apply(ApplyFunction&& function) {

91

3.12M

        for (auto& pair : _internal_map) {

92

            // TODO: Now only the cancel worker do the GC the _query_ctx_map. each query must

93

            // do erase the finish query unless in _query_ctx_map. Rethink the logic is ok

94

3.12M

            std::unique_lock lock(*pair.first);

95

3.12M

            static_cast<void>(function(pair.second));

96

3.12M

        }

97

24.4k

    }

    Status apply_if_not_exists(const Key& query_id, std::shared_ptr<ValueType>& query_ctx,

                               ApplyFunction&& function);

private:

    // The lock should only be used to protect the structures in fragment manager. Has to be

    // used in a very small scope because it may dead lock. For example, if the _lock is used

    // in prepare stage, the call path is  prepare --> expr prepare --> may call allocator

    // when allocate failed, allocator may call query_is_cancelled, query is callced will also

    // call _lock, so that there is dead lock.

    std::vector<std::pair<std::unique_ptr<std::shared_mutex>, phmap::flat_hash_map<Key, Value>>>

            _internal_map;

};

// This class used to manage all the fragment execute in this instance

class FragmentMgr : public RestMonitorIface {

public:

    using FinishCallback = std::function<void(RuntimeState*, Status*)>;

    FragmentMgr(ExecEnv* exec_env);

    ~FragmentMgr() override;

    void stop();

    // execute one plan fragment

    Status exec_plan_fragment(const TPipelineFragmentParams& params, const QuerySource query_type,

                              const TPipelineFragmentParamsList& parent);

    void remove_pipeline_context(std::pair<TUniqueId, int> key);

    void remove_query_context(const TUniqueId& key);

    // `is_prepare_success` is used by invoker to ensure callback can be handle correctly (eg. stream_load_executor)

    Status exec_plan_fragment(const TPipelineFragmentParams& params, const QuerySource query_type,

                              const FinishCallback& cb, const TPipelineFragmentParamsList& parent,

                              std::shared_ptr<bool> is_prepare_success = nullptr);

    Status start_query_execution(const PExecPlanFragmentStartRequest* request);

    Status trigger_pipeline_context_report(const ReportStatusRequest,

                                           std::shared_ptr<PipelineFragmentContext>&&);

    // Can be used in both version.

    MOCK_FUNCTION void cancel_query(const TUniqueId query_id, const Status reason);

    void cancel_worker();

    void debug(std::stringstream& ss) override;

    // input: TQueryPlanInfo fragment_instance_id

    // output: selected_columns

    // execute external query, all query info are packed in TScanOpenParams

    Status exec_external_plan_fragment(const TScanOpenParams& params,

                                       const TQueryPlanInfo& t_query_plan_info,

                                       const TUniqueId& query_id,

                                       const TUniqueId& fragment_instance_id,

                                       std::vector<TScanColumnDesc>* selected_columns);

    Status apply_filterv2(const PPublishFilterRequestV2* request,

                          butil::IOBufAsZeroCopyInputStream* attach_data);

    Status merge_filter(const PMergeFilterRequest* request,

                        butil::IOBufAsZeroCopyInputStream* attach_data);

    Status send_filter_size(const PSendFilterSizeRequest* request);

    Status sync_filter_size(const PSyncFilterSizeRequest* request);

    std::string to_http_path(const std::string& file_name);

    void coordinator_callback(const ReportStatusRequest& req);

    ThreadPool* get_thread_pool() { return _thread_pool.get(); }

    // When fragment mgr is going to stop, the _stop_background_threads_latch is set to 0

    // and other module that use fragment mgr's thread pool should get this signal and exit.

    bool shutting_down() { return _stop_background_threads_latch.count() == 0; }

    int32_t running_query_num() { return cast_set<int32_t>(_query_ctx_map.num_items()); }

    std::string dump_pipeline_tasks(int64_t duration = 0);

    std::string dump_pipeline_tasks(TUniqueId& query_id);

    void get_runtime_query_info(std::vector<std::weak_ptr<ResourceContext>>* _resource_ctx_list);

    Status get_realtime_exec_status(const TUniqueId& query_id,

                                    TReportExecStatusParams* exec_status);

    // get the query statistics of with a given query id

    Status get_query_statistics(const TUniqueId& query_id, TQueryStatistics* query_stats);

    std::shared_ptr<QueryContext> get_query_ctx(const TUniqueId& query_id);

    Status transmit_rec_cte_block(const TUniqueId& query_id, const TUniqueId& instance_id,

                                  int node_id,

                                  const google::protobuf::RepeatedPtrField<PBlock>& pblocks,

                                  bool eos);

    Status rerun_fragment(const std::shared_ptr<brpc::ClosureGuard>& guard,

                          const TUniqueId& query_id, int fragment,

                          PRerunFragmentParams_Opcode stage);

    Status reset_global_rf(const TUniqueId& query_id,

                           const google::protobuf::RepeatedField<int32_t>& filter_ids);

private:

    struct BrpcItem {

        TNetworkAddress network_address;

        std::vector<std::weak_ptr<QueryContext>> queries;

    };

    Status _get_or_create_query_ctx(const TPipelineFragmentParams& params,

                                    const TPipelineFragmentParamsList& parent,

                                    QuerySource query_type,

                                    std::shared_ptr<QueryContext>& query_ctx);

    void _collect_timeout_queries_and_brpc_items(

            std::vector<TUniqueId>& queries_timeout,

            std::unordered_map<std::shared_ptr<PBackendService_Stub>, BrpcItem>&

                    brpc_stub_with_queries,

            timespec now);

    void _collect_invalid_queries(

            std::vector<TUniqueId>& queries_lost_coordinator,

            std::vector<TUniqueId>& queries_pipeline_task_leak,

            const std::map<int64_t, std::unordered_set<TUniqueId>>& running_queries_on_all_fes,

            const std::map<TNetworkAddress, FrontendInfo>& running_fes,

            timespec check_invalid_query_last_timestamp);

    void _check_brpc_available(const std::shared_ptr<PBackendService_Stub>& brpc_stub,

                               const BrpcItem& brpc_item);

    // This is input params

    ExecEnv* _exec_env = nullptr;

    // (QueryID, FragmentID) -> PipelineFragmentContext

    ConcurrentContextMap<std::pair<TUniqueId, int>, std::shared_ptr<PipelineFragmentContext>,

                         PipelineFragmentContext>

            _pipeline_map;

    // Saved params and callback for rerunnable (recursive CTE) fragments.

    // Only populated when need_notify_close == true during exec_plan_fragment.

    // Lifecycle: created in exec_plan_fragment(), used in rerun_fragment(rebuild)

    // to recreate PFC with fresh state, cleaned up in remove_query_context().

    struct RerunableFragmentInfo {

        // Runtime filter IDs registered by the old PFC, collected during wait_for_destroy.

        // These are deregistered from the RuntimeFilterMgr before the new PFC is created.

        std::set<int> deregister_runtime_filter_ids;

        // Original params from FE, used to recreate the PFC each round.

        TPipelineFragmentParams params;

        TPipelineFragmentParamsList parent;

        FinishCallback finish_callback;

        // Hold query_ctx to prevent it from being destroyed while rerunnable fragments exist.

        std::shared_ptr<QueryContext> query_ctx;

        // Monotonically increasing stage counter, stamps runtime filter RPCs.

        uint32_t stage = 0;

    };

    std::mutex _rerunnable_params_lock;

    std::map<std::pair<TUniqueId, int>, RerunableFragmentInfo> _rerunnable_params_map;

    // query id -> QueryContext

    ConcurrentContextMap<TUniqueId, std::weak_ptr<QueryContext>, QueryContext> _query_ctx_map;

    // keep query ctx do not delete immediately to make rf coordinator merge filter work well after query eos

    ConcurrentContextMap<TUniqueId, std::shared_ptr<QueryContext>, QueryContext>

            _query_ctx_map_delay_delete;

    CountDownLatch _stop_background_threads_latch;

    std::shared_ptr<Thread> _cancel_thread;

    // This pool is used as global async task pool

    std::unique_ptr<ThreadPool> _thread_pool;

    std::shared_ptr<MetricEntity> _entity;

    UIntGauge* timeout_canceled_fragment_count = nullptr;

};

uint64_t get_fragment_executing_count();

uint64_t get_fragment_last_active_time();

#include "common/compile_check_end.h"

} // namespace doris