// 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.

#include "exec/exchange/vdata_stream_sender.h"

#include <fmt/format.h>

#include <fmt/ranges.h> // IWYU pragma: keep

#include <gen_cpp/DataSinks_types.h>

#include <gen_cpp/Metrics_types.h>

#include <gen_cpp/Types_types.h>

#include <gen_cpp/data.pb.h>

#include <gen_cpp/internal_service.pb.h>

#include <glog/logging.h>

#include <stddef.h>

#include <algorithm>

#include <cstdint>

#include <functional>

#include <map>

#include <memory>

#include <random>

#include "common/object_pool.h"

#include "common/status.h"

#include "core/column/column_const.h"

#include "exec/common/sip_hash.h"

#include "exec/exchange/vdata_stream_mgr.h"

#include "exec/exchange/vdata_stream_recvr.h"

#include "exec/operator/exchange_sink_operator.h"

#include "exec/operator/result_file_sink_operator.h"

#include "exec/sink/vrow_distribution.h"

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

#include "exprs/vexpr.h"

#include "runtime/descriptors.h"

#include "runtime/runtime_state.h"

#include "runtime/thread_context.h"

#include "storage/tablet_info.h"

#include "util/proto_util.h"

namespace doris {

Status Channel::init(RuntimeState* state) {

    // only enable_local_exchange() is true and the destination address is localhost, then the channel is local

    _is_local &= state->enable_local_exchange();

    if (_is_local) {

        return Status::OK();

    }

    RETURN_IF_ERROR(_init_brpc_stub(state));

    return Status::OK();

}

Status Channel::_init_brpc_stub(RuntimeState* state) {

    if (_brpc_dest_addr.hostname.empty()) {

        LOG(WARNING) << "there is no brpc destination address's hostname"

                        ", maybe version is not compatible.";

        return Status::InternalError("no brpc destination");

    }

    auto network_address = _brpc_dest_addr;

    if (_brpc_dest_addr.hostname == BackendOptions::get_localhost()) {

        _brpc_stub = state->exec_env()->brpc_internal_client_cache()->get_client(

                "127.0.0.1", _brpc_dest_addr.port);

        network_address.hostname = "127.0.0.1";

    } else {

        _brpc_stub = state->exec_env()->brpc_internal_client_cache()->get_client(_brpc_dest_addr);

    }

    if (!_brpc_stub) {

        std::string msg = fmt::format("Get rpc stub failed, dest_addr={}:{}",

                                      _brpc_dest_addr.hostname, _brpc_dest_addr.port);

        LOG(WARNING) << msg;

        return Status::InternalError(msg);

    }

    if (config::enable_brpc_connection_check) {

        state->get_query_ctx()->add_using_brpc_stub(_brpc_dest_addr, _brpc_stub);

    }

    return Status::OK();

}

Status Channel::open(RuntimeState* state) {

    if (_is_local) {

        RETURN_IF_ERROR(_find_local_recvr(state));

    }

    _be_number = state->be_number();

    _brpc_timeout_ms = get_execution_rpc_timeout_ms(state->execution_timeout());

    _serializer.set_is_local(_is_local);

    // In bucket shuffle join will set fragment_instance_id (-1, -1)

    // to build a camouflaged empty channel. the ip and port is '0.0.0.0:0"

    // so the empty channel not need call function close_internal()

    _need_close = (_fragment_instance_id.hi != -1 && _fragment_instance_id.lo != -1);

    return Status::OK();

}

Status Channel::_find_local_recvr(RuntimeState* state) {

    auto st = _parent->state()->exec_env()->vstream_mgr()->find_recvr(_fragment_instance_id,

                                                                      _dest_node_id, &_local_recvr);

    if (!st.ok()) {

        // If could not find local receiver, then it means the channel is EOF.

        // Maybe downstream task is finished already.

        //if (_receiver_status.ok()) {

        //    _receiver_status = Status::EndOfFile("local data stream receiver is deconstructed");

        //}

        LOG(INFO) << "Query: " << print_id(state->query_id())

                  << " recvr is not found, maybe downstream task is finished. error st is: "

                  << st.to_string();

    }

    return Status::OK();

}

Status Channel::add_rows(Block* block, const uint32_t* data, const uint32_t offset,

                         const uint32_t size, bool eos) {

    if (_fragment_instance_id.lo == -1) {

        return Status::OK();

    }

    bool serialized = false;

    if (_pblock == nullptr) {

        _pblock = std::make_unique<PBlock>();

    }

    RETURN_IF_ERROR(_serializer.next_serialized_block(block, _pblock.get(), 1, &serialized, eos,

                                                      data, offset, size));

    if (serialized) {

        RETURN_IF_ERROR(_send_current_block(eos));

    }

    return Status::OK();

}

std::shared_ptr<Dependency> Channel::get_local_channel_dependency() {

    if (!_local_recvr) {

        return nullptr;

    }

    return _local_recvr->get_local_channel_dependency(_parent->sender_id());

}

int64_t Channel::mem_usage() const {

    return _serializer.mem_usage();

}

Status Channel::send_remote_block(std::unique_ptr<PBlock>&& block, bool eos) {

    COUNTER_UPDATE(_parent->blocks_sent_counter(), 1);

    if (eos) {

        if (_eos_send) {

            return Status::OK();

        } else {

            _eos_send = true;

        }

    }

    if (eos || block->column_metas_size()) {

        RETURN_IF_ERROR(_buffer->add_block(this, {std::move(block), eos}));

    }

    return Status::OK();

}

Status Channel::send_broadcast_block(std::shared_ptr<BroadcastPBlockHolder>& block, bool eos) {

    COUNTER_UPDATE(_parent->blocks_sent_counter(), 1);

    if (eos) {

        if (_eos_send) {

            return Status::OK();

        }

        _eos_send = true;

    }

    if (eos || block->get_block()->column_metas_size()) {

        RETURN_IF_ERROR(_buffer->add_block(this, {block, eos}));

    }

    return Status::OK();

}

Status Channel::_send_current_block(bool eos) {

    if (is_local()) {

        return _send_local_block(eos);

    }

    // here _pblock maybe nullptr , but we must send the eos to the receiver

    return send_remote_block(std::move(_pblock), eos);

}

Status Channel::_send_local_block(bool eos) {

    Block block;

    if (_serializer.get_block() != nullptr) {

        block = _serializer.get_block()->to_block();

        _serializer.get_block()->set_mutable_columns(block.clone_empty_columns());

    }

    if (!block.empty() || eos) { // if eos is true, we MUST to send an empty block

        RETURN_IF_ERROR(send_local_block(&block, eos, true));

    }

    return Status::OK();

}

Status Channel::send_local_block(Block* block, bool eos, bool can_be_moved) {

    SCOPED_TIMER(_parent->local_send_timer());

    if (eos) {

        if (_eos_send) {

            return Status::OK();

        } else {

            _eos_send = true;

        }

    }

    if (is_receiver_eof()) {

        return _receiver_status;

    }

    auto receiver_status = _recvr_status();

    // _local_recvr depdend on ExchangeLocalState* _parent to do some memory counter settings

    // but it only owns a raw pointer, so that the ExchangeLocalState object may be deconstructed.

    // Lock the fragment context to ensure the runtime state and other objects are not deconstructed

    TaskExecutionContextSPtr ctx_lock = nullptr;

    if (receiver_status.ok() && _local_recvr != nullptr) {

        ctx_lock = _local_recvr->task_exec_ctx();

        // Do not return internal error, because when query finished, the downstream node

        // may finish before upstream node. And the object maybe deconstructed. If return error

        // then the upstream node may report error status to FE, the query is failed.

        if (ctx_lock == nullptr) {

            receiver_status = Status::EndOfFile("local data stream receiver is deconstructed");

        }

    }

    if (receiver_status.ok()) {

        COUNTER_UPDATE(_parent->local_bytes_send_counter(), block->bytes());

        COUNTER_UPDATE(_parent->local_sent_rows(), block->rows());

        COUNTER_UPDATE(_parent->blocks_sent_counter(), 1);

        const auto sender_id = _parent->sender_id();

        if (!block->empty()) [[likely]] {

            _local_recvr->add_block(block, sender_id, can_be_moved);

        }

        if (eos) [[unlikely]] {

            _local_recvr->remove_sender(sender_id, _be_number, Status::OK());

            _parent->on_channel_finished(_fragment_instance_id.lo);

        }

        return Status::OK();

    } else {

        _receiver_status = std::move(receiver_status);

        _parent->on_channel_finished(_fragment_instance_id.lo);

        return _receiver_status;

    }

}

std::string Channel::debug_string() const {

    fmt::memory_buffer debug_string_buffer;

    fmt::format_to(debug_string_buffer,

                   "fragment_instance_id: {}, _dest_node_id: {}, _is_local: {}, _receiver_status: "

                   "{}, _closed: {}, _need_close: {}, _be_number: {}, _eos_send: {}",

                   print_id(_fragment_instance_id), _dest_node_id, _is_local,

                   _receiver_status.to_string(), _closed, _need_close, _be_number, _eos_send);

    if (_is_local) {

        fmt::format_to(debug_string_buffer, "_local_recvr: {}", _local_recvr->debug_string());

    }

    return fmt::to_string(debug_string_buffer);

}

Status Channel::close(RuntimeState* state) {

    if (_closed) {

        return Status::OK();

    }

    _closed = true;

    if (!_need_close) {

        return Status::OK();

    }

    if (is_receiver_eof()) {

        _serializer.reset_block();

        return Status::OK();

    } else {

        return _send_current_block(true);

    }

}

BlockSerializer::BlockSerializer(ExchangeSinkLocalState* parent, bool is_local)

        : _parent(parent), _is_local(is_local), _batch_size(parent->state()->batch_size()) {}

Status BlockSerializer::next_serialized_block(Block* block, PBlock* dest, size_t num_receivers,

                                              bool* serialized, bool eos, const uint32_t* data,

                                              const uint32_t offset, const uint32_t size) {

    if (_mutable_block == nullptr) {

        _mutable_block = MutableBlock::create_unique(block->clone_empty());

    }

    {

        SCOPED_TIMER(_parent->merge_block_timer());

        if (data) {

            if (size > 0) {

                RETURN_IF_ERROR(

                        _mutable_block->add_rows(block, data + offset, data + offset + size));

            }

        } else if (!block->empty()) {

            RETURN_IF_ERROR(_mutable_block->merge(*block));

        }

    }

    if (_mutable_block->rows() >= _batch_size || eos ||

        (_mutable_block->rows() > 0 && _mutable_block->allocated_bytes() > _buffer_mem_limit)) {

        if (!_is_local) {

            RETURN_IF_ERROR(_serialize_block(dest, num_receivers));

        }

        *serialized = true;

    } else {

        *serialized = false;

    }

    return Status::OK();

}

Status BlockSerializer::_serialize_block(PBlock* dest, size_t num_receivers) {

    if (_mutable_block && _mutable_block->rows() > 0) {

        auto block = _mutable_block->to_block();

        RETURN_IF_ERROR(serialize_block(&block, dest, num_receivers));

        if (_parent->state()->low_memory_mode()) {

            reset_block();

        } else {

            block.clear_column_data();

            _mutable_block->set_mutable_columns(block.mutate_columns());

        }

    }

    return Status::OK();

}

Status BlockSerializer::serialize_block(const Block* src, PBlock* dest, size_t num_receivers) {

    SCOPED_TIMER(_parent->_serialize_batch_timer);

    dest->Clear();

    size_t uncompressed_bytes = 0, compressed_bytes = 0;

    int64_t compress_time = 0;

    RETURN_IF_ERROR(src->serialize(_parent->_state->be_exec_version(), dest, &uncompressed_bytes,

                                   &compressed_bytes, &compress_time, _parent->compression_type(),

                                   _parent->transfer_large_data_by_brpc()));

    COUNTER_UPDATE(_parent->_bytes_sent_counter, compressed_bytes * num_receivers);

    COUNTER_UPDATE(_parent->_uncompressed_bytes_counter, uncompressed_bytes * num_receivers);

    COUNTER_UPDATE(_parent->_compress_timer, compress_time);

#ifndef BE_TEST

    _parent->state()->get_query_ctx()->resource_ctx()->io_context()->update_shuffle_send_bytes(

            compressed_bytes * num_receivers);

    _parent->state()->get_query_ctx()->resource_ctx()->io_context()->update_shuffle_send_rows(

            src->rows() * num_receivers);

#endif

    return Status::OK();

}

} // namespace doris