// 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 <gen_cpp/PaloInternalService_types.h>

#include <atomic>

#include <condition_variable>

#include <cstdint>

#include <future>

#include <memory>

#include <mutex>

#include <set>

#include <shared_mutex>

#include <thread>

#include <unordered_map>

#include <utility>

#include "common/status.h"

#include "core/block/block.h"

#include "exec/sink/writer/vwal_writer.h"

#include "load/group_commit/wal/wal_manager.h"

#include "runtime/exec_env.h"

#include "util/threadpool.h"

namespace doris {

class ExecEnv;

class GroupCommitMgr;

class TUniqueId;

class RuntimeState;

class Dependency;

struct BlockData {

BlockData(const std::shared_ptr<Block>& block) : block(block), block_bytes(block->bytes()) {};

std::shared_ptr<Block> block;

size_t block_bytes;

};

class LoadBlockQueue {

public:

LoadBlockQueue(const UniqueId& load_instance_id, std::string& label, int64_t txn_id,

int64_t schema_version, int64_t index_size,

std::shared_ptr<std::atomic_size_t> all_block_queues_bytes,

bool wait_internal_group_commit_finish, int64_t group_commit_interval_ms,

int64_t group_commit_data_bytes)

: load_instance_id(load_instance_id),

label(label),

txn_id(txn_id),

schema_version(schema_version),

index_size(index_size),

wait_internal_group_commit_finish(wait_internal_group_commit_finish),

_group_commit_interval_ms(group_commit_interval_ms),

_start_time(std::chrono::steady_clock::now()),

_last_print_time(_start_time),

_group_commit_data_bytes(group_commit_data_bytes),

_all_block_queues_bytes(all_block_queues_bytes) {};

Status add_block(RuntimeState* runtime_state, std::shared_ptr<Block> block, bool write_wal,

UniqueId& load_id);

Status get_block(RuntimeState* runtime_state, Block* block, bool* find_block, bool* eos,

std::shared_ptr<Dependency> get_block_dep);

bool contain_load_id(const UniqueId& load_id);

Status add_load_id(const UniqueId& load_id, const std::shared_ptr<Dependency> put_block_dep);

Status remove_load_id(const UniqueId& load_id);

void cancel(const Status& st);

bool need_commit() const { return _need_commit.load(); }

Status create_wal(int64_t db_id, int64_t tb_id, int64_t wal_id, const std::string& import_label,

WalManager* wal_manager, std::vector<TSlotDescriptor>& slot_desc,

int be_exe_version);

Status close_wal();

bool has_enough_wal_disk_space(size_t estimated_wal_bytes);

void append_dependency(std::shared_ptr<Dependency> finish_dep);

void append_read_dependency(std::shared_ptr<Dependency> read_dep);

int64_t get_group_commit_interval_ms() { return _group_commit_interval_ms; };

std::string debug_string() const {

fmt::memory_buffer debug_string_buffer;

fmt::format_to(

debug_string_buffer,

"load_instance_id={}, label={}, txn_id={}, "

"wait_internal_group_commit_finish={}, data_size_condition={}, "

"group_commit_load_count={}, process_finish={}, _need_commit={}, schema_version={}",

load_instance_id.to_string(), label, txn_id, wait_internal_group_commit_finish,

data_size_condition, group_commit_load_count, process_finish.load(),

_need_commit.load(), schema_version);

return fmt::to_string(debug_string_buffer);

}

UniqueId load_instance_id;

std::string label;

int64_t txn_id;

int64_t schema_version;

int64_t index_size;

bool wait_internal_group_commit_finish = false;

bool data_size_condition = false;

// counts of load in one group commit

std::atomic_size_t group_commit_load_count = 0;

// only used by fault injection (debug point) to reproduce that multiple loads

// reuse one group commit plan and share a load_id

std::atomic<int> _debug_add_block_seq = 0;

// the execute status of this internal group commit

std::mutex mutex;

std::atomic<bool> process_finish = false;

Status status = Status::OK();

std::vector<std::shared_ptr<Dependency>> dependencies;

private:

void _cancel_without_lock(const Status& st);

std::string _get_load_ids();

// the set of load ids of all blocks in this queue

std::map<UniqueId, std::shared_ptr<Dependency>> _load_ids_to_write_dep;

std::vector<std::shared_ptr<Dependency>> _read_deps;

std::list<BlockData> _block_queue;

// wal

std::string _wal_base_path;

std::shared_ptr<VWalWriter> _v_wal_writer;

// commit

std::atomic<bool> _need_commit = false;

// commit by time interval, can be changed by 'ALTER TABLE my_table SET ("group_commit_interval_ms"="1000");'

int64_t _group_commit_interval_ms;

std::chrono::steady_clock::time_point _start_time;

std::chrono::steady_clock::time_point _last_print_time;

// commit by data size

int64_t _group_commit_data_bytes;

int64_t _data_bytes = 0;

// memory back pressure, memory consumption of all tables' load block queues

std::shared_ptr<std::atomic_size_t> _all_block_queues_bytes;

std::condition_variable _get_cond;

};

class GroupCommitTable {

public:

GroupCommitTable(ExecEnv* exec_env, doris::ThreadPool* thread_pool, int64_t db_id,

int64_t table_id, std::shared_ptr<std::atomic_size_t> all_block_queue_bytes,

GroupCommitMgr* group_commit_mgr)

: _exec_env(exec_env),

_thread_pool(thread_pool),

_all_block_queues_bytes(all_block_queue_bytes),

_db_id(db_id),

_table_id(table_id),

_group_commit_mgr(group_commit_mgr) {};

Status get_first_block_load_queue(int64_t table_id, int64_t base_schema_version,

int64_t index_size, const UniqueId& load_id,

std::shared_ptr<LoadBlockQueue>& load_block_queue,

int be_exe_version,

std::shared_ptr<Dependency> create_plan_dep,

std::shared_ptr<Dependency> put_block_dep);

Status get_load_block_queue(const TUniqueId& instance_id,

std::shared_ptr<LoadBlockQueue>& load_block_queue,

std::shared_ptr<Dependency> get_block_dep);

void remove_load_id(const UniqueId& load_id);

Status submit_create_group_commit_load();

private:

Status _submit_create_group_commit_load();

Status _create_group_commit_load(int be_exe_version,

const std::shared_ptr<MemTrackerLimiter>& mem_tracker,

std::shared_ptr<LoadBlockQueue>& created_load_block_queue);

Status _exec_plan_fragment(int64_t db_id, int64_t table_id, const std::string& label,

int64_t txn_id, const TPipelineFragmentParams& pipeline_params);

Status _finish_group_commit_load(int64_t db_id, int64_t table_id, const std::string& label,

int64_t txn_id, const TUniqueId& instance_id, Status& status,

RuntimeState* state);

ExecEnv* _exec_env = nullptr;

ThreadPool* _thread_pool = nullptr;

// memory consumption of all tables' load block queues, used for memory back pressure.

std::shared_ptr<std::atomic_size_t> _all_block_queues_bytes;

int64_t _db_id;

int64_t _table_id;

GroupCommitMgr* _group_commit_mgr = nullptr;

std::mutex _lock;

// fragment_instance_id to load_block_queue

std::unordered_map<UniqueId, std::shared_ptr<LoadBlockQueue>> _load_block_queues;

bool _is_creating_plan_fragment = false;

// user_load_id -> <create_plan_dep, put_block_dep, base_schema_version, index_size>

std::unordered_map<UniqueId, std::tuple<std::shared_ptr<Dependency>,

std::shared_ptr<Dependency>, int64_t, int64_t>>

_create_plan_deps;

std::string _create_plan_failed_reason;

int _create_plan_be_exe_version = 0;

int64_t _create_plan_start_time_ms = 0;

};

class GroupCommitMgr {

public:

GroupCommitMgr(ExecEnv* exec_env);

virtual ~GroupCommitMgr();

void stop();

// used when init group_commit_scan_node

Status get_load_block_queue(int64_t table_id, const TUniqueId& instance_id,

std::shared_ptr<LoadBlockQueue>& load_block_queue,

std::shared_ptr<Dependency> get_block_dep);

Status get_first_block_load_queue(int64_t db_id, int64_t table_id, int64_t base_schema_version,

int64_t index_size, const UniqueId& load_id,

std::shared_ptr<LoadBlockQueue>& load_block_queue,

int be_exe_version,

std::shared_ptr<Dependency> create_plan_dep,

std::shared_ptr<Dependency> put_block_dep);

void remove_load_id(int64_t table_id, const UniqueId& load_id);

std::promise<Status> debug_promise;

std::future<Status> debug_future = debug_promise.get_future();

void add_need_create_plan_table(int64_t table_id);

std::shared_ptr<MemTrackerLimiter> group_commit_mem_tracker() {

return _group_commit_mem_tracker;

}

private:

void _create_plan_worker();

ExecEnv* _exec_env = nullptr;

std::unique_ptr<doris::ThreadPool> _thread_pool;

// memory consumption of all tables' load block queues, used for memory back pressure.

std::shared_ptr<std::atomic_size_t> _all_block_queues_bytes;

std::mutex _lock;

// TODO remove table when unused

std::unordered_map<int64_t, std::shared_ptr<GroupCommitTable>> _table_map;

std::mutex _need_create_plan_lock;

std::condition_variable _need_create_plan_cv;

std::set<int64_t> _need_create_plan_tables;

bool _stopped = false;

std::thread _create_plan_thread;

std::shared_ptr<MemTrackerLimiter> _group_commit_mem_tracker;

};

} // namespace doris