// 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/sink/load_stream_stub.h"

#include <sstream>

#include "common/cast_set.h"

#include "runtime/query_context.h"

#include "storage/rowset/rowset_writer.h"

#include "util/brpc_client_cache.h"

#include "util/debug_points.h"

#include "util/network_util.h"

#include "util/thrift_util.h"

#include "util/uid_util.h"

namespace doris {

int LoadStreamReplyHandler::on_received_messages(brpc::StreamId id, butil::IOBuf* const messages[],

                                                 size_t size) {

    auto stub = _stub.lock();

    if (!stub) {

        LOG(WARNING) << "stub is not exist when on_received_messages, " << *this

                     << ", stream_id=" << id;

        return 0;

    }

    for (size_t i = 0; i < size; i++) {

        butil::IOBufAsZeroCopyInputStream wrapper(*messages[i]);

        PLoadStreamResponse response;

        response.ParseFromZeroCopyStream(&wrapper);

        if (response.eos()) {

            stub->_is_eos.store(true);

        }

        Status st = Status::create<false>(response.status());

        std::stringstream ss;

        ss << "on_received_messages, " << *this << ", stream_id=" << id;

        if (response.success_tablet_ids_size() > 0) {

            ss << ", success tablet ids:";

            for (auto tablet_id : response.success_tablet_ids()) {

                ss << " " << tablet_id;

            }

            std::lock_guard<bthread::Mutex> lock(stub->_success_tablets_mutex);

            for (auto tablet_id : response.success_tablet_ids()) {

                stub->_success_tablets.push_back(tablet_id);

            }

        }

        if (response.failed_tablets_size() > 0) {

            ss << ", failed tablet ids:";

            for (auto pb : response.failed_tablets()) {

                ss << " " << pb.id() << ":" << Status::create(pb.status());

            }

            std::lock_guard<bthread::Mutex> lock(stub->_failed_tablets_mutex);

            for (auto pb : response.failed_tablets()) {

                stub->_failed_tablets.emplace(pb.id(), Status::create(pb.status()));

            }

        }

        if (response.tablet_schemas_size() > 0) {

            ss << ", tablet schema num: " << response.tablet_schemas_size();

            std::lock_guard<bthread::Mutex> lock(stub->_schema_mutex);

            for (const auto& schema : response.tablet_schemas()) {

                auto tablet_schema = std::make_unique<TabletSchema>();

                tablet_schema->init_from_pb(schema.tablet_schema());

                stub->_tablet_schema_for_index->emplace(schema.index_id(),

                                                        std::move(tablet_schema));

                stub->_enable_unique_mow_for_index->emplace(

                        schema.index_id(), schema.enable_unique_key_merge_on_write());

            }

            stub->_schema_cv.notify_all();

        }

        ss << ", status: " << st;

        LOG(INFO) << ss.str();

        if (response.tablet_load_rowset_num_infos_size() > 0) {

            stub->_refresh_back_pressure_version_wait_time(response.tablet_load_rowset_num_infos());

        }

        if (response.has_load_stream_profile()) {

            TRuntimeProfileTree tprofile;

            const uint8_t* buf =

                    reinterpret_cast<const uint8_t*>(response.load_stream_profile().data());

            uint32_t len = cast_set<uint32_t>(response.load_stream_profile().size());

            auto status = deserialize_thrift_msg(buf, &len, false, &tprofile);

            if (status.ok()) {

                // TODO

                //_sink->_state->load_channel_profile()->update(tprofile);

            } else {

                LOG(WARNING) << "load stream TRuntimeProfileTree deserialize failed, errmsg="

                             << status;

            }

        }

    }

    return 0;

}

void LoadStreamReplyHandler::on_closed(brpc::StreamId id) {

    Defer defer {[this]() { delete this; }};

    LOG(INFO) << "on_closed, " << *this << ", stream_id=" << id;

    auto stub = _stub.lock();

    if (!stub) {

        LOG(WARNING) << "stub is not exist when on_closed, " << *this;

        return;

    }

    stub->_is_closed.store(true);

}

inline std::ostream& operator<<(std::ostream& ostr, const LoadStreamReplyHandler& handler) {

    ostr << "LoadStreamReplyHandler load_id=" << UniqueId(handler._load_id)

         << ", dst_id=" << handler._dst_id;

    return ostr;

}

LoadStreamStub::LoadStreamStub(PUniqueId load_id, int64_t src_id,

                               std::shared_ptr<IndexToTabletSchema> schema_map,

                               std::shared_ptr<IndexToEnableMoW> mow_map, bool incremental)

        : _load_id(load_id),

          _src_id(src_id),

          _tablet_schema_for_index(schema_map),

          _enable_unique_mow_for_index(mow_map),

          _is_incremental(incremental) {};

LoadStreamStub::~LoadStreamStub() {

    if (_is_open.load() && !_is_closed.load()) {

        auto ret = brpc::StreamClose(_stream_id);

        LOG(INFO) << *this << " is deconstructed, close " << (ret == 0 ? "success" : "failed");

    }

}

// open_load_stream

Status LoadStreamStub::open(BrpcClientCache<PBackendService_Stub>* client_cache,

                            const NodeInfo& node_info, int64_t txn_id,

                            const OlapTableSchemaParam& schema,

                            const std::vector<PTabletID>& tablets_for_schema, int total_streams,

                            int64_t idle_timeout_ms, bool enable_profile) {

    std::unique_lock<bthread::Mutex> lock(_open_mutex);

    if (_is_init.load()) {

        return _status;

    }

    _is_init.store(true);

    _dst_id = node_info.id;

    brpc::StreamOptions opt;

    opt.max_buf_size = cast_set<int>(config::load_stream_max_buf_size);

    opt.idle_timeout_ms = idle_timeout_ms;

    opt.messages_in_batch = config::load_stream_messages_in_batch;

    opt.handler = new LoadStreamReplyHandler(_load_id, _dst_id, shared_from_this());

    brpc::Controller cntl;

    if (int ret = brpc::StreamCreate(&_stream_id, cntl, &opt)) {

        delete opt.handler;

        _status = Status::Error<true>(ret, "Failed to create stream");

        return _status;

    }

    cntl.set_timeout_ms(config::open_load_stream_timeout_ms);

    POpenLoadStreamRequest request;

    *request.mutable_load_id() = _load_id;

    request.set_src_id(_src_id);

    request.set_txn_id(txn_id);

    request.set_enable_profile(enable_profile);

    if (_is_incremental) {

        request.set_total_streams(0);

    } else if (total_streams > 0) {

        request.set_total_streams(total_streams);

    } else {

        _status = Status::InternalError("total_streams should be greator than 0");

        return _status;

    }

    request.set_idle_timeout_ms(idle_timeout_ms);

    schema.to_protobuf(request.mutable_schema());

    for (auto& tablet : tablets_for_schema) {

        *request.add_tablets() = tablet;

    }

    POpenLoadStreamResponse response;

    // set connection_group "streaming" to distinguish with non-streaming connections

    const auto& stub = client_cache->get_client(node_info.host, node_info.brpc_port);

    if (stub == nullptr) {

        return Status::InternalError("failed to init brpc client to {}:{}", node_info.host,

                                     node_info.brpc_port);

    }

    stub->open_load_stream(&cntl, &request, &response, nullptr);

    for (const auto& resp : response.tablet_schemas()) {

        auto tablet_schema = std::make_unique<TabletSchema>();

        tablet_schema->init_from_pb(resp.tablet_schema());

        _tablet_schema_for_index->emplace(resp.index_id(), std::move(tablet_schema));

        _enable_unique_mow_for_index->emplace(resp.index_id(),

                                              resp.enable_unique_key_merge_on_write());

    }

    if (response.tablet_load_rowset_num_infos_size() > 0) {

        _refresh_back_pressure_version_wait_time(response.tablet_load_rowset_num_infos());

    }

    if (cntl.Failed()) {

        brpc::StreamClose(_stream_id);

        _status = Status::InternalError("Failed to connect to backend {}: {}", _dst_id,

                                        cntl.ErrorText());

        return _status;

    }

    LOG(INFO) << "open load stream to host=" << node_info.host << ", port=" << node_info.brpc_port

              << ", " << *this;

    _is_open.store(true);

    _status = Status::OK();

    return _status;

}

// APPEND_DATA

Status LoadStreamStub::append_data(int64_t partition_id, int64_t index_id, int64_t tablet_id,

                                   int32_t segment_id, uint64_t offset, std::span<const Slice> data,

                                   bool segment_eos, FileType file_type) {

    if (!_is_open.load()) {

        add_failed_tablet(tablet_id, _status);

        return _status;

    }

    DBUG_EXECUTE_IF("LoadStreamStub.skip_send_segment", { return Status::OK(); });

    PStreamHeader header;

    header.set_src_id(_src_id);

    *header.mutable_load_id() = _load_id;

    header.set_partition_id(partition_id);

    header.set_index_id(index_id);

    header.set_tablet_id(tablet_id);

    header.set_segment_id(segment_id);

    header.set_segment_eos(segment_eos);

    header.set_offset(offset);

    header.set_opcode(doris::PStreamHeader::APPEND_DATA);

    header.set_file_type(file_type);

    return _encode_and_send(header, data);

}

// ADD_SEGMENT

Status LoadStreamStub::add_segment(int64_t partition_id, int64_t index_id, int64_t tablet_id,

                                   int32_t segment_id, const SegmentStatistics& segment_stat) {

    if (!_is_open.load()) {

        add_failed_tablet(tablet_id, _status);

        return _status;

    }

    DBUG_EXECUTE_IF("LoadStreamStub.skip_send_segment", { return Status::OK(); });

    PStreamHeader header;

    header.set_src_id(_src_id);

    *header.mutable_load_id() = _load_id;

    header.set_partition_id(partition_id);

    header.set_index_id(index_id);

    header.set_tablet_id(tablet_id);

    header.set_segment_id(segment_id);

    header.set_opcode(doris::PStreamHeader::ADD_SEGMENT);

    segment_stat.to_pb(header.mutable_segment_statistics());

    return _encode_and_send(header);

}

// CLOSE_LOAD

Status LoadStreamStub::close_load(const std::vector<PTabletID>& tablets_to_commit,

                                  int num_incremental_streams) {

    if (!_is_open.load()) {

        return _status;

    }

    PStreamHeader header;

    *header.mutable_load_id() = _load_id;

    header.set_src_id(_src_id);

    header.set_opcode(doris::PStreamHeader::CLOSE_LOAD);

    for (const auto& tablet : tablets_to_commit) {

        *header.add_tablets() = tablet;

    }

    header.set_num_incremental_streams(num_incremental_streams);

    _status = _encode_and_send(header);

    if (!_status.ok()) {

        LOG(WARNING) << "stream " << _stream_id << " close failed: " << _status;

        return _status;

    }

    _is_closing.store(true);

    return Status::OK();

}

// GET_SCHEMA

Status LoadStreamStub::get_schema(const std::vector<PTabletID>& tablets) {

    if (!_is_open.load()) {

        return _status;

    }

    PStreamHeader header;

    *header.mutable_load_id() = _load_id;

    header.set_src_id(_src_id);

    header.set_opcode(doris::PStreamHeader::GET_SCHEMA);

    std::ostringstream oss;

    oss << "fetching tablet schema from stream " << _stream_id

        << ", load id: " << print_id(_load_id) << ", tablet id:";

    for (const auto& tablet : tablets) {

        *header.add_tablets() = tablet;

        oss << " " << tablet.tablet_id();

    }

    if (tablets.size() == 0) {

        oss << " none";

    }

    LOG(INFO) << oss.str();

    return _encode_and_send(header);

}

Status LoadStreamStub::wait_for_schema(int64_t partition_id, int64_t index_id, int64_t tablet_id,

                                       int64_t timeout_ms) {

    if (!_is_open.load()) {

        return _status;

    }

    if (_tablet_schema_for_index->contains(index_id)) {

        return Status::OK();

    }

    PTabletID tablet;

    tablet.set_partition_id(partition_id);

    tablet.set_index_id(index_id);

    tablet.set_tablet_id(tablet_id);

    RETURN_IF_ERROR(get_schema({tablet}));

    MonotonicStopWatch watch;

    watch.start();

    while (!_tablet_schema_for_index->contains(index_id) &&

           watch.elapsed_time() / 1000 / 1000 < timeout_ms) {

        RETURN_IF_ERROR(check_cancel());

        static_cast<void>(wait_for_new_schema(100));

    }

    if (!_tablet_schema_for_index->contains(index_id)) {

        return Status::TimedOut("timeout to get tablet schema for index {}", index_id);

    }

    return Status::OK();

}

Status LoadStreamStub::close_finish_check(RuntimeState* state, bool* is_closed) {

    DBUG_EXECUTE_IF("LoadStreamStub::close_wait.long_wait", DBUG_BLOCK);

    DBUG_EXECUTE_IF("LoadStreamStub::close_finish_check.close_failed",

                    { return Status::InternalError("close failed"); });

    *is_closed = true;

    if (!_is_open.load()) {

        // we don't need to close wait on non-open streams

        return Status::OK();

    }

    // If stream is cancelled (e.g., due to connection failure), treat it as closed

    // to avoid waiting indefinitely for a stream that will never respond.

    if (_is_cancelled.load()) {

        return check_cancel();

    }

    if (state->get_query_ctx()->is_cancelled()) {

        return state->get_query_ctx()->exec_status();

    }

    if (!_is_closing.load()) {

        *is_closed = false;

        return _status;

    }

    if (_is_closed.load()) {

        RETURN_IF_ERROR(check_cancel());

        if (!_is_eos.load()) {

            return Status::InternalError("Stream closed without EOS, {}", to_string());

        }

        return Status::OK();

    }

    *is_closed = false;

    return Status::OK();

}

void LoadStreamStub::cancel(Status reason) {

    LOG(WARNING) << *this << " is cancelled because of " << reason;

    if (_is_open.load()) {

        brpc::StreamClose(_stream_id);

    }

    {

        std::lock_guard<bthread::Mutex> lock(_cancel_mutex);

        _cancel_st = reason;

        _is_cancelled.store(true);

    }

    _is_closed.store(true);

}

Status LoadStreamStub::_encode_and_send(PStreamHeader& header, std::span<const Slice> data) {

    butil::IOBuf buf;

    size_t header_len = header.ByteSizeLong();

    buf.append(reinterpret_cast<uint8_t*>(&header_len), sizeof(header_len));

    buf.append(header.SerializeAsString());

    size_t data_len = std::transform_reduce(data.begin(), data.end(), 0, std::plus(),

                                            [](const Slice& s) { return s.get_size(); });

    buf.append(reinterpret_cast<uint8_t*>(&data_len), sizeof(data_len));

    for (const auto& slice : data) {

        buf.append(slice.get_data(), slice.get_size());

    }

    bool eos = header.opcode() == doris::PStreamHeader::CLOSE_LOAD;

    bool get_schema = header.opcode() == doris::PStreamHeader::GET_SCHEMA;

    add_bytes_written(buf.size());

    return _send_with_buffer(buf, eos || get_schema);

}

Status LoadStreamStub::_send_with_buffer(butil::IOBuf& buf, bool sync) {

    butil::IOBuf output;

    std::unique_lock<decltype(_buffer_mutex)> buffer_lock(_buffer_mutex);

    _buffer.append(buf);

    if (!sync && _buffer.size() < config::brpc_streaming_client_batch_bytes) {

        return Status::OK();

    }

    output.swap(_buffer);

    // acquire send lock while holding buffer lock, to ensure the message order

    std::lock_guard<decltype(_send_mutex)> send_lock(_send_mutex);

    buffer_lock.unlock();

    VLOG_DEBUG << "send buf size : " << output.size() << ", sync: " << sync;

    auto st = _send_with_retry(output);

    if (!st.ok()) {

        _handle_failure(output, st);

    }

    return st;

}

void LoadStreamStub::_handle_failure(butil::IOBuf& buf, Status st) {

    while (buf.size() > 0) {

        // step 1: parse header

        size_t hdr_len = 0;

        buf.cutn((void*)&hdr_len, sizeof(size_t));

        butil::IOBuf hdr_buf;

        PStreamHeader hdr;

        buf.cutn(&hdr_buf, hdr_len);

        butil::IOBufAsZeroCopyInputStream wrapper(hdr_buf);

        hdr.ParseFromZeroCopyStream(&wrapper);

        // step 2: cut data

        size_t data_len = 0;

        buf.cutn((void*)&data_len, sizeof(size_t));

        butil::IOBuf data_buf;

        buf.cutn(&data_buf, data_len);

        // step 3: handle failure

        switch (hdr.opcode()) {

        case PStreamHeader::ADD_SEGMENT:

        case PStreamHeader::APPEND_DATA: {

            DBUG_EXECUTE_IF("LoadStreamStub._handle_failure.append_data_failed", {

                add_failed_tablet(hdr.tablet_id(), st);

                return;

            });

            DBUG_EXECUTE_IF("LoadStreamStub._handle_failure.add_segment_failed", {

                add_failed_tablet(hdr.tablet_id(), st);

                return;

            });

            add_failed_tablet(hdr.tablet_id(), st);

        } break;

        case PStreamHeader::CLOSE_LOAD: {

            DBUG_EXECUTE_IF("LoadStreamStub._handle_failure.close_load_failed", {

                brpc::StreamClose(_stream_id);

                return;

            });

            brpc::StreamClose(_stream_id);

        } break;

        case PStreamHeader::GET_SCHEMA: {

            DBUG_EXECUTE_IF("LoadStreamStub._handle_failure.get_schema_failed", {

                // Just log and let wait_for_schema timeout

                std::ostringstream oss;

                for (const auto& tablet : hdr.tablets()) {

                    oss << " " << tablet.tablet_id();

                }

                LOG(WARNING) << "failed to send GET_SCHEMA request, tablet_id:" << oss.str() << ", "

                             << *this;

                return;

            });

            // Just log and let wait_for_schema timeout

            std::ostringstream oss;

            for (const auto& tablet : hdr.tablets()) {

                oss << " " << tablet.tablet_id();

            }

            LOG(WARNING) << "failed to send GET_SCHEMA request, tablet_id:" << oss.str() << ", "

                         << *this;

        } break;

        default:

            LOG(WARNING) << "unexpected stream message " << hdr.opcode() << ", " << *this;

            DCHECK(false);

        }

    }

}

Status LoadStreamStub::_send_with_retry(butil::IOBuf& buf) {

    for (;;) {

        RETURN_IF_ERROR(check_cancel());

        int ret;

        {

            DBUG_EXECUTE_IF("LoadStreamStub._send_with_retry.delay_before_send", {

                int64_t delay_ms = dp->param<int64_t>("delay_ms", 1000);

                bthread_usleep(delay_ms * 1000);

            });

            brpc::StreamWriteOptions options;

            options.write_in_background = config::enable_brpc_stream_write_background;

            ret = brpc::StreamWrite(_stream_id, buf, &options);

        }

        DBUG_EXECUTE_IF("LoadStreamStub._send_with_retry.stream_write_failed", { ret = EPIPE; });

        switch (ret) {

        case 0:

            return Status::OK();

        case EAGAIN: {

            const timespec time = butil::seconds_from_now(config::load_stream_eagain_wait_seconds);

            int wait_ret = brpc::StreamWait(_stream_id, &time);

            if (wait_ret != 0) {

                return Status::InternalError("StreamWait failed, err={}, {}", wait_ret,

                                             to_string());

            }

            break;

        }

        default:

            return Status::InternalError("StreamWrite failed, err={}, {}", ret, to_string());

        }

    }

}

void LoadStreamStub::_refresh_back_pressure_version_wait_time(

        const ::google::protobuf::RepeatedPtrField<::doris::PTabletLoadRowsetInfo>&

                tablet_load_infos) {

    int64_t max_rowset_num_gap = 0;

    // if any one tablet is under high load pressure, we would make the whole procedure

    // sleep to prevent the corresponding BE return -235

    std::for_each(

            tablet_load_infos.begin(), tablet_load_infos.end(),

            [&max_rowset_num_gap](auto& load_info) {

                int64_t cur_rowset_num = load_info.current_rowset_nums();

                int64_t high_load_point = load_info.max_config_rowset_nums() *

                                          (config::load_back_pressure_version_threshold / 100);

                DCHECK(cur_rowset_num > high_load_point);

                max_rowset_num_gap = std::max(max_rowset_num_gap, cur_rowset_num - high_load_point);

            });

    // to slow down the high load pressure

    // we would use the rowset num gap to calculate one sleep time

    // for example:

    // if the max tablet version is 2000, there are 3 BE

    // A: ====================  1800

    // B: ===================   1700

    // C: ==================    1600

    //    ==================    1600

    //                      ^

    //                      the high load point

    // then then max gap is 1800 - (max tablet version * config::load_back_pressure_version_threshold / 100) = 200,

    // we would make the whole send procesure sleep

    // 1200ms for compaction to be done toe reduce the high pressure

    auto max_time = config::max_load_back_pressure_version_wait_time_ms;

    if (UNLIKELY(max_rowset_num_gap > 0)) {

        _load_back_pressure_version_wait_time_ms.store(

                std::min(max_rowset_num_gap + 1000, max_time));

        LOG(INFO) << "try to back pressure version, wait time(ms): "

                  << _load_back_pressure_version_wait_time_ms << ", load id: " << print_id(_load_id)

                  << ", max_rowset_num_gap: " << max_rowset_num_gap;

    }

}

std::string LoadStreamStub::to_string() {

    std::ostringstream ss;

    ss << *this;

    return ss.str();

}

inline std::ostream& operator<<(std::ostream& ostr, const LoadStreamStub& stub) {

    ostr << "LoadStreamStub load_id=" << print_id(stub._load_id) << ", src_id=" << stub._src_id

         << ", dst_id=" << stub._dst_id << ", stream_id=" << stub._stream_id;

    return ostr;

}

Status LoadStreamStubs::open(BrpcClientCache<PBackendService_Stub>* client_cache,

                             const NodeInfo& node_info, int64_t txn_id,

                             const OlapTableSchemaParam& schema,

                             const std::vector<PTabletID>& tablets_for_schema, int total_streams,

                             int64_t idle_timeout_ms, bool enable_profile) {

    bool get_schema = true;

    auto status = Status::OK();

    bool first_stream = true;

    for (auto& stream : _streams) {

        Status st;

        if (get_schema) {

            st = stream->open(client_cache, node_info, txn_id, schema, tablets_for_schema,

                              total_streams, idle_timeout_ms, enable_profile);

        } else {

            st = stream->open(client_cache, node_info, txn_id, schema, {}, total_streams,

                              idle_timeout_ms, enable_profile);

        }

        // Simulate one stream open failure within LoadStreamStubs.

        // This causes the successfully opened streams to be cancelled,

        // reproducing the bug where cancelled streams cause close_wait timeout.

        DBUG_EXECUTE_IF("LoadStreamStubs.open.fail_one_stream", {

            if (st.ok() && !first_stream) {

                st = Status::InternalError("Injected stream open failure");

            }

        });

        if (st.ok()) {

            get_schema = false;

        } else {

            LOG(WARNING) << "open stream failed: " << st << "; stream: " << *stream;

            status = st;

            // no break here to try get schema from the rest streams

        }

        first_stream = false;

    }

    // only mark open when all streams open success

    _open_success.store(status.ok());

    // cancel all streams if open failed

    if (!status.ok()) {

        cancel(status);

    }

    return status;

}

Status LoadStreamStubs::close_load(const std::vector<PTabletID>& tablets_to_commit,

                                   int num_incremental_streams) {

    if (!_open_success.load()) {

        return Status::InternalError("streams not open");

    }

    bool first = true;

    auto status = Status::OK();

    for (auto& stream : _streams) {

        Status st;

        if (first) {

            st = stream->close_load(tablets_to_commit, num_incremental_streams);

            first = false;

        } else {

            st = stream->close_load({}, num_incremental_streams);

        }

        if (!st.ok()) {

            LOG(WARNING) << "close_load failed: " << st << "; stream: " << *stream;

        }

    }

    return status;

}

} // namespace doris