be/src/load/channel/tablets_channel.cpp

Source
// 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 "load/channel/tablets_channel.h"

#include <bvar/bvar.h>
#include <fmt/format.h>
#include <gen_cpp/internal_service.pb.h>
#include <gen_cpp/types.pb.h>

#include <ctime>

#include "common/compiler_util.h" // IWYU pragma: keep
#include "common/status.h"
// IWYU pragma: no_include <bits/chrono.h>
#include <chrono> // IWYU pragma: keep
#include <initializer_list>
#include <optional>
#include <set>
#include <thread>
#include <utility>

#ifdef DEBUG
#include <unordered_set>
#endif

#include "common/logging.h"
#include "common/metrics/doris_metrics.h"
#include "common/metrics/metrics.h"
#include "core/block/block.h"
#include "load/channel/load_channel.h"
#include "load/delta_writer/delta_writer.h"
#include "storage/storage_engine.h"
#include "storage/tablet/tablet_manager.h"
#include "storage/tablet_info.h"
#include "storage/txn/txn_manager.h"
#include "util/defer_op.h"

namespace doris {
class SlotDescriptor;

bvar::Adder<int64_t> g_tablets_channel_send_data_allocated_size(
        "tablets_channel_send_data_allocated_size");

DEFINE_GAUGE_METRIC_PROTOTYPE_2ARG(tablet_writer_count, MetricUnit::NOUNIT);

std::atomic<uint64_t> BaseTabletsChannel::_s_tablet_writer_count;

BaseTabletsChannel::BaseTabletsChannel(const TabletsChannelKey& key, const UniqueId& load_id,
                                       bool is_high_priority, RuntimeProfile* profile)
        : _key(key),
          _state(kInitialized),
          _load_id(load_id),
          _closed_senders(64),
          _is_high_priority(is_high_priority) {
    static std::once_flag once_flag;
    if (profile != nullptr) {
        _init_profile(profile);
    }
    std::call_once(once_flag, [] {
        REGISTER_HOOK_METRIC(tablet_writer_count, [&]() { return _s_tablet_writer_count.load(); });
    });
}

TabletsChannel::TabletsChannel(StorageEngine& engine, const TabletsChannelKey& key,
                               const UniqueId& load_id, bool is_high_priority,
                               RuntimeProfile* profile)
        : BaseTabletsChannel(key, load_id, is_high_priority, profile), _engine(engine) {}

BaseTabletsChannel::~BaseTabletsChannel() {
    _s_tablet_writer_count -= _tablet_writers.size();
}

TabletsChannel::~TabletsChannel() = default;

Status BaseTabletsChannel::_get_current_seq(int64_t& cur_seq,
                                            const PTabletWriterAddBlockRequest& request) {
    std::lock_guard<std::mutex> l(_lock);
    if (_state != kOpened) {
        return _state == kFinished ? _close_status
                                   : Status::InternalError("TabletsChannel {} state: {}",
                                                           _key.to_string(), _state);
    }
    cur_seq = _next_seqs[request.sender_id()];
    // check packet
    if (request.packet_seq() > cur_seq) {
        LOG(WARNING) << "lost data packet, expect_seq=" << cur_seq
                     << ", recept_seq=" << request.packet_seq();
        return Status::InternalError("lost data packet");
    }
    return Status::OK();
}

void BaseTabletsChannel::_init_profile(RuntimeProfile* profile) {
    DCHECK(profile != nullptr);
    _profile =
            profile->create_child(fmt::format("TabletsChannel {}", _key.to_string()), true, true);
    _add_batch_number_counter = ADD_COUNTER(_profile, "NumberBatchAdded", TUnit::UNIT);

    auto* memory_usage = _profile->create_child("PeakMemoryUsage", true, true);
    _add_batch_timer = ADD_TIMER(_profile, "AddBatchTime");
    _write_block_timer = ADD_TIMER(_profile, "WriteBlockTime");
    _incremental_open_timer = ADD_TIMER(_profile, "IncrementalOpenTabletTime");
    _memory_usage_counter = memory_usage->AddHighWaterMarkCounter("Total", TUnit::BYTES);
    _write_memory_usage_counter = memory_usage->AddHighWaterMarkCounter("Write", TUnit::BYTES);
    _flush_memory_usage_counter = memory_usage->AddHighWaterMarkCounter("Flush", TUnit::BYTES);
    _max_tablet_memory_usage_counter =
            memory_usage->AddHighWaterMarkCounter("MaxTablet", TUnit::BYTES);
    _max_tablet_write_memory_usage_counter =
            memory_usage->AddHighWaterMarkCounter("MaxTabletWrite", TUnit::BYTES);
    _max_tablet_flush_memory_usage_counter =
            memory_usage->AddHighWaterMarkCounter("MaxTabletFlush", TUnit::BYTES);
}

Status BaseTabletsChannel::open(const PTabletWriterOpenRequest& request) {
    std::lock_guard<std::mutex> l(_lock);
    // if _state is kOpened, it's a normal case, already open by other sender
    // if _state is kFinished, already cancelled by other sender
    if (_state == kOpened) {
        RETURN_IF_ERROR(_init_adaptive_random_bucket_state(request));
        return Status::OK();
    }
    if (_state == kFinished) {
        return Status::OK();
    }
    _txn_id = request.txn_id();
    _index_id = request.index_id();
    _schema = std::make_shared<OlapTableSchemaParam>();
    RETURN_IF_ERROR(_schema->init(request.schema()));
    _tuple_desc = _schema->tuple_desc();
    int max_sender = request.num_senders();
    /*
     * a tablets channel in reciever is related to a bulk of VNodeChannel of sender. each instance one or none.
     * there are two possibilities:
     *  1. there's partitions originally broadcasted by FE. so all sender(instance) know it at start. and open() will be 
     *     called directly, not by incremental_open(). and after _state changes to kOpened. _open_by_incremental will never 
     *     be true. in this case, _num_remaining_senders will keep same with senders number. when all sender sent close rpc,
     *     the tablets channel will close. and if for auto partition table, these channel's closing will hang on reciever and
     *     return together to avoid close-then-incremental-open problem.
     *  2. this tablets channel is opened by incremental_open of sender's sink node. so only this sender will know this partition
     *     (this TabletsChannel) at that time. and we are not sure how many sender will know in the end. it depends on data
     *     distribution. in this situation open() is called by incremental_open() at first time. so _open_by_incremental is true.
     *     then _num_remaining_senders will not be set here. but inc every time when incremental_open() called. so it's dynamic
     *     and also need same number of senders' close to close. but will not hang.
     */
    if (_open_by_incremental) {
        DCHECK(_num_remaining_senders == 0) << _num_remaining_senders;
    } else {
        _num_remaining_senders = max_sender;
    }
    LOG(INFO) << fmt::format(
            "open tablets channel {}, tablets num: {} timeout(s): {}, init senders {} with "
            "incremental {}",
            _key.to_string(), request.tablets().size(), request.load_channel_timeout_s(),
            _num_remaining_senders, _open_by_incremental ? "on" : "off");
    // just use max_sender no matter incremental or not cuz we dont know how many senders will open.
    _next_seqs.resize(max_sender, 0);
    _closed_senders.Reset(max_sender);

    RETURN_IF_ERROR(_open_all_writers(request));
    RETURN_IF_ERROR(_init_adaptive_random_bucket_state(request));

    _state = kOpened;
    return Status::OK();
}

Status BaseTabletsChannel::incremental_open(const PTabletWriterOpenRequest& params) {
    SCOPED_TIMER(_incremental_open_timer);

    // current node first opened by incremental open
    if (_state == kInitialized) {
        _open_by_incremental = true;
        RETURN_IF_ERROR(open(params));
    }

    std::lock_guard<std::mutex> l(_lock);

    // one sender may incremental_open many times. but only close one time. so dont count duplicately.
    if (_open_by_incremental) {
        if (params.has_sender_id() && !_recieved_senders.contains(params.sender_id())) {
            _recieved_senders.insert(params.sender_id());
            _num_remaining_senders++;
        } else if (!params.has_sender_id()) { // for compatible
            _num_remaining_senders++;
        }
        VLOG_DEBUG << fmt::format("txn {}: TabletsChannel {} inc senders to {}", _txn_id, _index_id,
                                  _num_remaining_senders);
    }

    std::vector<SlotDescriptor*>* index_slots = nullptr;
    int32_t schema_hash = 0;

    for (const auto& index : _schema->indexes()) {
        if (index->index_id == _index_id) {
            index_slots = &index->slots;
            schema_hash = index->schema_hash;
            break;
        }
    }
    if (index_slots == nullptr) {
        return Status::InternalError("unknown index id, key={}", _key.to_string());
    }
    // update tablets
    size_t incremental_tablet_num = 0;
    std::stringstream ss;
    ss << "LocalTabletsChannel txn_id: " << _txn_id << " load_id: " << print_id(params.id())
       << " incremental open delta writer: ";

    // every change will hold _lock. this find in under _lock too. so no need _tablet_writers_lock again.
    for (const auto& tablet : params.tablets()) {
        if (_tablet_writers.find(tablet.tablet_id()) != _tablet_writers.end()) {
            continue;
        }
        incremental_tablet_num++;

        WriteRequest wrequest;
        wrequest.index_id = params.index_id();
        wrequest.tablet_id = tablet.tablet_id();
        wrequest.schema_hash = schema_hash;
        wrequest.txn_id = _txn_id;
        wrequest.partition_id = tablet.partition_id();
        wrequest.load_id = params.id();
        wrequest.tuple_desc = _tuple_desc;
        wrequest.slots = index_slots;
        wrequest.is_high_priority = _is_high_priority;
        wrequest.table_schema_param = _schema;
        wrequest.txn_expiration = params.txn_expiration(); // Required by CLOUD.
        wrequest.write_file_cache = params.write_file_cache();
        wrequest.storage_vault_id = params.storage_vault_id();
        wrequest.enable_table_memtable_backpressure = params.is_adaptive_random_bucket();
        if (tablet.has_binlog_tablet_id()) {
            wrequest.binlog_tablet_id = tablet.binlog_tablet_id();
        }

        auto delta_writer = create_delta_writer(wrequest);
        {
            // here we modify _tablet_writers. so need lock.
            std::lock_guard<std::mutex> lt(_tablet_writers_lock);
            _tablet_writers.emplace(tablet.tablet_id(), std::move(delta_writer));
        }

        ss << "[" << tablet.tablet_id() << "]";
    }

    _s_tablet_writer_count += incremental_tablet_num;
    LOG(INFO) << ss.str();
    RETURN_IF_ERROR(_init_adaptive_random_bucket_state(params));

    _state = kOpened;
    return Status::OK();
}

Status BaseTabletsChannel::_init_adaptive_random_bucket_state(
        const PTabletWriterOpenRequest& request) {
    if (!request.is_adaptive_random_bucket() || request.random_bucket_partitions_size() == 0) {
        return Status::OK();
    }
    if (_adaptive_random_bucket_state == nullptr) {
        _adaptive_random_bucket_state = std::make_shared<AdaptiveRandomBucketState>(_load_id);
    }
    for (const auto& partition : request.random_bucket_partitions()) {
        if (partition.ordered_tablet_ids_size() == 0) {
            return Status::InternalError(
                    "ordered_tablet_ids is empty for adaptive random bucket, load_id={}, "
                    "sender_id={}, partition_id={}",
                    print_id(_load_id), request.sender_id(), partition.partition_id());
        }
        std::vector<int32_t> ordered_positions;
        ordered_positions.reserve(partition.ordered_tablet_ids_size());
        for (int i = 0; i < partition.ordered_tablet_ids_size(); ++i) {
            ordered_positions.push_back(cast_set<int32_t>(i));
        }
        std::vector<int64_t> ordered_tablet_ids;
        ordered_tablet_ids.reserve(partition.ordered_tablet_ids_size());
        for (auto tablet_id : partition.ordered_tablet_ids()) {
            ordered_tablet_ids.push_back(tablet_id);
        }
        RETURN_IF_ERROR(_adaptive_random_bucket_state->init_partition(
                partition.partition_id(), ordered_tablet_ids, ordered_positions, 0));
    }
    return Status::OK();
}

std::unique_ptr<BaseDeltaWriter> TabletsChannel::create_delta_writer(const WriteRequest& request) {
    DCHECK(request.write_req_type == WriteRequestType::DATA);
    DCHECK(request.table_schema_param != nullptr);

    // whether to write binlog is decided by binlog_tablet_id: it is set only when this backend
    // also owns the binlog tablet that is paired with the base tablet.
    if (request.binlog_tablet_id <= 0) {
        return std::make_unique<DeltaWriter>(_engine, request, _profile, _load_id);
    }

    int64_t row_binlog_index_id = 0;
    for (const auto* index_schema : request.table_schema_param->indexes()) {
        if (index_schema->index_id == request.index_id) {
            row_binlog_index_id = index_schema->row_binlog_id;
            break;
        }
    }
    DCHECK(row_binlog_index_id > 0);

    const auto* row_binlog_index_schema =
            request.table_schema_param->row_binlog_index_schema(row_binlog_index_id);
    DCHECK(row_binlog_index_schema != nullptr);

    // group_build_req is only for the group wrapper itself. It provides the group semantics and
    // metadata used by BaseDeltaWriter/GroupRowsetBuilder to expose tablet_id, txn_id,
    // partition_id, load_id and profile information, while concrete rowset builders use the
    // sub requests below.
    WriteRequest group_build_req = request;
    group_build_req.write_req_type = WriteRequestType::GROUP;

    WriteRequest sub_data_req = request;
    sub_data_req.write_req_type = WriteRequestType::DATA;

    WriteRequest sub_row_binlog_req = request;
    sub_row_binlog_req.write_req_type = WriteRequestType::ROW_BINLOG;
    sub_row_binlog_req.tablet_id = request.binlog_tablet_id;
    sub_row_binlog_req.index_id = row_binlog_index_schema->index_id;
    sub_row_binlog_req.schema_hash = row_binlog_index_schema->schema_hash;

    return std::make_unique<DeltaWriter>(_engine, group_build_req, sub_data_req, sub_row_binlog_req,
                                         _profile, _load_id);
}

Status TabletsChannel::close(LoadChannel* parent, const PTabletWriterAddBlockRequest& req,
                             PTabletWriterAddBlockResult* res, bool* finished) {
    int sender_id = req.sender_id();
    int64_t backend_id = req.backend_id();
    const auto& partition_ids = req.partition_ids();
    auto* tablet_errors = res->mutable_tablet_errors();
    std::lock_guard<std::mutex> l(_lock);
    if (_state == kFinished) {
        return _close_status;
    }
    if (_closed_senders.Get(sender_id)) {
        // Double close from one sender, just return OK
        *finished = (_num_remaining_senders == 0);
        return _close_status;
    }

    for (auto pid : partition_ids) {
        _partition_ids.emplace(pid);
    }
    _closed_senders.Set(sender_id, true);
    _num_remaining_senders--;
    *finished = (_num_remaining_senders == 0);

    LOG(INFO) << fmt::format(
            "txn {}: close tablets channel of index {} , sender id: {}, backend {}, remain "
            "senders: {}",
            _txn_id, _index_id, sender_id, backend_id, _num_remaining_senders);

    if (!*finished) {
        return Status::OK();
    }

    _state = kFinished;
    // All senders are closed
    // 1. close all delta writers
    std::set<DeltaWriter*> need_wait_writers;
    // under _lock. no need _tablet_writers_lock again.
    for (auto&& [tablet_id, writer] : _tablet_writers) {
        if (_partition_ids.contains(writer->partition_id())) {
            auto st = writer->close();
            if (!st.ok()) {
                auto err_msg = fmt::format(
                        "close tablet writer failed, tablet_id={}, "
                        "transaction_id={}, err={}",
                        tablet_id, _txn_id, st.to_string());
                LOG(WARNING) << err_msg;
                PTabletError* tablet_error = tablet_errors->Add();
                tablet_error->set_tablet_id(tablet_id);
                tablet_error->set_msg(st.to_string());
                // just skip this tablet(writer) and continue to close others
                continue;
            }
            // tablet writer in `_broken_tablets` should not call `build_rowset` and
            // `commit_txn` method, after that, the publish-version task will success,
            // which can cause the replica inconsistency.
            if (_is_broken_tablet(writer->tablet_id())) {
                LOG(WARNING) << "SHOULD NOT HAPPEN, tablet writer is broken but not cancelled"
                             << ", tablet_id=" << tablet_id << ", transaction_id=" << _txn_id;
                continue;
            }
            need_wait_writers.insert(static_cast<DeltaWriter*>(writer.get()));
        } else {
            auto st = writer->cancel();
            if (!st.ok()) {
                LOG(WARNING) << "cancel tablet writer failed, tablet_id=" << tablet_id
                             << ", transaction_id=" << _txn_id;
                // just skip this tablet(writer) and continue to close others
                continue;
            }
            VLOG_PROGRESS << "cancel tablet writer successfully, tablet_id=" << tablet_id
                          << ", transaction_id=" << _txn_id;
        }
    }

    // 2. wait all writer finished flush.
    for (auto* writer : need_wait_writers) {
        RETURN_IF_ERROR((writer->wait_flush()));
    }

    // 3. build rowset
    for (auto it = need_wait_writers.begin(); it != need_wait_writers.end();) {
        Status st = (*it)->build_rowset();
        if (!st.ok()) {
            _add_error_tablet(tablet_errors, (*it)->tablet_id(), st);
            it = need_wait_writers.erase(it);
            continue;
        }
        // 3.1 calculate delete bitmap for Unique Key MoW tables
        st = (*it)->submit_calc_delete_bitmap_task();
        if (!st.ok()) {
            _add_error_tablet(tablet_errors, (*it)->tablet_id(), st);
            it = need_wait_writers.erase(it);
            continue;
        }
        it++;
    }

    // 4. wait for delete bitmap calculation complete if necessary
    for (auto it = need_wait_writers.begin(); it != need_wait_writers.end();) {
        Status st = (*it)->wait_calc_delete_bitmap();
        if (!st.ok()) {
            _add_error_tablet(tablet_errors, (*it)->tablet_id(), st);
            it = need_wait_writers.erase(it);
            continue;
        }
        it++;
    }

    // 5. commit all writers

    for (auto* writer : need_wait_writers) {
        // close may return failed, but no need to handle it here.
        // tablet_vec will only contains success tablet, and then let FE judge it.
        _commit_txn(writer, res);
    }

    return Status::OK();
}

void TabletsChannel::_commit_txn(DeltaWriter* writer, PTabletWriterAddBlockResult* res) {
    Status st = writer->commit_txn();
    if (st.ok()) [[likely]] {
        auto* tablet_vec = res->mutable_tablet_vec();
        PTabletInfo* tablet_info = tablet_vec->Add();
        tablet_info->set_tablet_id(writer->tablet_id());
        // unused required field.
        tablet_info->set_schema_hash(0);
        tablet_info->set_received_rows(writer->total_received_rows());
        tablet_info->set_num_rows_filtered(writer->num_rows_filtered());
        // report the row binlog tablet as a normal tablet so FE advances its version.
        if (writer->binlog_tablet_id() > 0) {
            PTabletInfo* binlog_tablet_info = tablet_vec->Add();
            binlog_tablet_info->set_tablet_id(writer->binlog_tablet_id());
            binlog_tablet_info->set_schema_hash(0);
            binlog_tablet_info->set_received_rows(writer->total_received_rows());
            binlog_tablet_info->set_num_rows_filtered(writer->num_rows_filtered());
        }
        _total_received_rows += writer->total_received_rows();
        _num_rows_filtered += writer->num_rows_filtered();
    } else {
        _add_error_tablet(res->mutable_tablet_errors(), writer->tablet_id(), st);
    }
}

void BaseTabletsChannel::_add_error_tablet(
        google::protobuf::RepeatedPtrField<PTabletError>* tablet_errors, int64_t tablet_id,
        Status error) const {
    PTabletError* tablet_error = tablet_errors->Add();
    tablet_error->set_tablet_id(tablet_id);
    tablet_error->set_msg(error.to_string());
    VLOG_PROGRESS << "close wait failed tablet " << tablet_id << " transaction_id " << _txn_id
                  << "err msg " << error;
}

void BaseTabletsChannel::refresh_profile() {
    int64_t write_mem_usage = 0;
    int64_t flush_mem_usage = 0;
    int64_t max_tablet_mem_usage = 0;
    int64_t max_tablet_write_mem_usage = 0;
    int64_t max_tablet_flush_mem_usage = 0;
    {
        std::lock_guard<std::mutex> l(_tablet_writers_lock);
        for (auto&& [tablet_id, writer] : _tablet_writers) {
            int64_t write_mem = writer->mem_consumption(MemType::WRITE_FINISHED);
            write_mem_usage += write_mem;
            int64_t flush_mem = writer->mem_consumption(MemType::FLUSH);
            flush_mem_usage += flush_mem;
            if (write_mem > max_tablet_write_mem_usage) {
                max_tablet_write_mem_usage = write_mem;
            }
            if (flush_mem > max_tablet_flush_mem_usage) {
                max_tablet_flush_mem_usage = flush_mem;
            }
            if (write_mem + flush_mem > max_tablet_mem_usage) {
                max_tablet_mem_usage = write_mem + flush_mem;
            }
        }
    }
    COUNTER_SET(_memory_usage_counter, write_mem_usage + flush_mem_usage);
    COUNTER_SET(_write_memory_usage_counter, write_mem_usage);
    COUNTER_SET(_flush_memory_usage_counter, flush_mem_usage);
    COUNTER_SET(_max_tablet_memory_usage_counter, max_tablet_mem_usage);
    COUNTER_SET(_max_tablet_write_memory_usage_counter, max_tablet_write_mem_usage);
    COUNTER_SET(_max_tablet_flush_memory_usage_counter, max_tablet_flush_mem_usage);
}

Status BaseTabletsChannel::_open_all_writers(const PTabletWriterOpenRequest& request) {
    std::vector<SlotDescriptor*>* index_slots = nullptr;
    int32_t schema_hash = 0;
    for (const auto& index : _schema->indexes()) {
        if (index->index_id == _index_id) {
            index_slots = &index->slots;
            schema_hash = index->schema_hash;
            break;
        }
    }
    if (index_slots == nullptr) {
        return Status::InternalError("unknown index id, key={}", _key.to_string());
    }

#ifdef DEBUG
    // check: tablet ids should be unique
    {
        std::unordered_set<int64_t> tablet_ids;
        for (const auto& tablet : request.tablets()) {
            CHECK(tablet_ids.count(tablet.tablet_id()) == 0)
                    << "found duplicate tablet id: " << tablet.tablet_id();
            tablet_ids.insert(tablet.tablet_id());
        }
    }
#endif

    int tablet_cnt = 0;
    // under _lock. no need _tablet_writers_lock again.
    for (const auto& tablet : request.tablets()) {
        if (_tablet_writers.find(tablet.tablet_id()) != _tablet_writers.end()) {
            continue;
        }
        tablet_cnt++;
        WriteRequest wrequest {
                .tablet_id = tablet.tablet_id(),
                .schema_hash = schema_hash,
                .txn_id = _txn_id,
                .txn_expiration = request.txn_expiration(), // Required by CLOUD.
                .index_id = request.index_id(),
                .partition_id = tablet.partition_id(),
                .load_id = request.id(),
                .tuple_desc = _tuple_desc,
                .slots = index_slots,
                .table_schema_param = _schema,
                .is_high_priority = _is_high_priority,
                .write_file_cache = request.write_file_cache(),
                .storage_vault_id = request.storage_vault_id(),
                .enable_table_memtable_backpressure = request.is_adaptive_random_bucket(),
        };
        if (tablet.has_binlog_tablet_id()) {
            wrequest.binlog_tablet_id = tablet.binlog_tablet_id();
        }

        auto delta_writer = create_delta_writer(wrequest);
        {
            std::lock_guard<std::mutex> l(_tablet_writers_lock);
            _tablet_writers.emplace(tablet.tablet_id(), std::move(delta_writer));
        }
    }
    _s_tablet_writer_count += _tablet_writers.size();
    DCHECK_EQ(_tablet_writers.size(), tablet_cnt);
    return Status::OK();
}

Status BaseTabletsChannel::cancel() {
    std::lock_guard<std::mutex> l(_lock);
    if (_state == kFinished) {
        return _close_status;
    }
    for (auto& it : _tablet_writers) {
        static_cast<void>(it.second->cancel());
    }
    _state = kFinished;

    return Status::OK();
}

std::string TabletsChannelKey::to_string() const {
    std::stringstream ss;
    ss << *this;
    return ss.str();
}

std::ostream& operator<<(std::ostream& os, const TabletsChannelKey& key) {
    os << "(load_id=" << key.id << ", index_id=" << key.index_id << ")";
    return os;
}

Status BaseTabletsChannel::_write_block_data(
        const PTabletWriterAddBlockRequest& request, int64_t cur_seq,
        std::unordered_map<int64_t, TabletAddRowsPayload>& tablet_to_rows,
        PTabletWriterAddBlockResult* response) {
    Block send_data;
    [[maybe_unused]] size_t uncompressed_size = 0;
    [[maybe_unused]] int64_t uncompressed_time = 0;
    RETURN_IF_ERROR(send_data.deserialize(request.block(), &uncompressed_size, &uncompressed_time));
    if (send_data.rows() != request.tablet_ids_size()) {
        return Status::InternalError(
                "invalid add block request row count, load_id={}, index_id={}, packet_seq={}, "
                "block_rows={}, tablet_ids_size={}",
                print_id(_load_id), _index_id, request.packet_seq(), send_data.rows(),
                request.tablet_ids_size());
    }
    bool has_row_binlog_lsn = request.row_binlog_lsns_size() > 0;
    if (has_row_binlog_lsn) {
        if (send_data.rows() != request.row_binlog_lsns_size()) {
            return Status::InternalError(
                    "invalid add block request row-binlog lsn count, load_id={}, index_id={}, "
                    "packet_seq={}, block_rows={}, row_binlog_lsns_size={}",
                    print_id(_load_id), _index_id, request.packet_seq(), send_data.rows(),
                    request.row_binlog_lsns_size());
        }
    }

    g_tablets_channel_send_data_allocated_size << send_data.allocated_bytes();
    Defer defer {
            [&]() { g_tablets_channel_send_data_allocated_size << -send_data.allocated_bytes(); }};

    auto write_tablet_data = [&](int64_t tablet_id,
                                 std::function<Status(BaseDeltaWriter * writer)> write_func) {
        google::protobuf::RepeatedPtrField<PTabletError>* tablet_errors =
                response->mutable_tablet_errors();

        // add_batch may concurrency with inc_open but not under _lock.
        // so need to protect it with _tablet_writers_lock.
        BaseDeltaWriter* tablet_writer = nullptr;
        {
            std::lock_guard<std::mutex> l(_tablet_writers_lock);
            auto tablet_writer_it = _tablet_writers.find(tablet_id);
            if (tablet_writer_it == _tablet_writers.end()) {
                return Status::InternalError("unknown tablet to append data, tablet={}", tablet_id);
            }
            tablet_writer = tablet_writer_it->second.get();
        }

        Status st = write_func(tablet_writer);
        if (!st.ok()) {
            auto err_msg =
                    fmt::format("tablet writer write failed, tablet_id={}, txn_id={}, err={}",
                                tablet_id, _txn_id, st.to_string());
            LOG(WARNING) << err_msg;
            PTabletError* error = tablet_errors->Add();
            error->set_tablet_id(tablet_id);
            error->set_msg(err_msg);
            static_cast<void>(tablet_writer->cancel_with_status(st));
            _add_broken_tablet(tablet_id);
            // continue write to other tablet.
            // the error will return back to sender.
        }
        return Status::OK();
    };

    SCOPED_TIMER(_write_block_timer);
    auto* tablet_load_infos = response->mutable_tablet_load_rowset_num_infos();
    for (const auto& tablet_to_rows_it : tablet_to_rows) {
        bool memtable_flushed = false;
        RETURN_IF_ERROR(write_tablet_data(tablet_to_rows_it.first, [&](BaseDeltaWriter* writer) {
            return writer->write(&send_data, tablet_to_rows_it.second, &memtable_flushed);
        }));

        BaseDeltaWriter* tablet_writer = nullptr;
        {
            std::lock_guard<std::mutex> l(_tablet_writers_lock);
            auto tablet_writer_it = _tablet_writers.find(tablet_to_rows_it.first);
            if (tablet_writer_it != _tablet_writers.end()) {
                tablet_writer = tablet_writer_it->second.get();
            }
        }
        if (tablet_writer != nullptr) {
            tablet_writer->set_tablet_load_rowset_num_info(tablet_load_infos);
        }
    }

    {
        std::lock_guard<std::mutex> l(_lock);
        _next_seqs[request.sender_id()] = cur_seq + 1;
    }
    return Status::OK();
}

std::shared_ptr<std::mutex> BaseTabletsChannel::_get_partition_route_lock(int64_t partition_id) {
    std::lock_guard<std::mutex> l(_partition_route_locks_lock);
    auto& lock = _partition_route_locks[partition_id];
    if (lock == nullptr) {
        lock = std::make_shared<std::mutex>();
    }
    return lock;
}

Status BaseTabletsChannel::_write_block_data_for_adaptive_random_bucket(
        const PTabletWriterAddBlockRequest& request, int64_t cur_seq,
        std::unordered_map<int64_t, DorisVector<uint32_t>>& partition_to_rowidxs,
        PTabletWriterAddBlockResult* response) {
    Block send_data;
    [[maybe_unused]] size_t uncompressed_size = 0;
    [[maybe_unused]] int64_t uncompressed_time = 0;
    RETURN_IF_ERROR(send_data.deserialize(request.block(), &uncompressed_size, &uncompressed_time));
    if (send_data.rows() != request.partition_ids_size()) {
        return Status::InternalError(
                "invalid adaptive random bucket add block request row count, load_id={}, "
                "index_id={}, packet_seq={}, block_rows={}, partition_ids_size={}",
                print_id(_load_id), _index_id, request.packet_seq(), send_data.rows(),
                request.partition_ids_size());
    }

    {
        std::lock_guard<std::mutex> l(_lock);
        for (const auto& [partition_id, _] : partition_to_rowidxs) {
            _partition_ids.emplace(partition_id);
        }
    }

    g_tablets_channel_send_data_allocated_size << send_data.allocated_bytes();
    Defer defer {
            [&]() { g_tablets_channel_send_data_allocated_size << -send_data.allocated_bytes(); }};

    auto* tablet_errors = response->mutable_tablet_errors();
    auto* tablet_load_infos = response->mutable_tablet_load_rowset_num_infos();

    auto write_partition_data = [&](int64_t partition_id,
                                    const DorisVector<uint32_t>& row_idxs) -> Status {
        auto partition_lock = _get_partition_route_lock(partition_id);
        std::lock_guard<std::mutex> partition_guard(*partition_lock);
        int64_t tablet_id = -1;
        if (_adaptive_random_bucket_state == nullptr) {
            return Status::InternalError(
                    "adaptive random bucket state is not initialized, load_id={}, "
                    "index_id={}, packet_seq={}, partition_id={}",
                    print_id(_load_id), _index_id, request.packet_seq(), partition_id);
        }
        tablet_id = _adaptive_random_bucket_state->current_tablet(partition_id);
        if (tablet_id < 0) {
            return Status::InternalError(
                    "invalid current tablet for adaptive random bucket, load_id={}, "
                    "index_id={}, sender_id={}, packet_seq={}, partition_id={}",
                    print_id(_load_id), _index_id, request.sender_id(), request.packet_seq(),
                    partition_id);
        }
        VLOG_DEBUG << "FIND_TABLET_RANDOM_BUCKET: route+write begin"
                   << ", load_id=" << _load_id << ", index_id=" << _index_id
                   << ", sender_id=" << request.sender_id()
                   << ", packet_seq=" << request.packet_seq() << ", partition_id=" << partition_id
                   << ", tablet_id=" << tablet_id << ", row_count=" << row_idxs.size();

        {
            std::shared_lock<std::shared_mutex> broken_rlock(_broken_tablets_lock);
            if (_is_broken_tablet(tablet_id)) {
                return Status::InternalError(
                        "current tablet is broken for adaptive random bucket, load_id={}, "
                        "index_id={}, sender_id={}, packet_seq={}, partition_id={}, tablet_id={}",
                        print_id(_load_id), _index_id, request.sender_id(), request.packet_seq(),
                        partition_id, tablet_id);
            }
        }

        BaseDeltaWriter* tablet_writer = nullptr;
        {
            std::lock_guard<std::mutex> l(_tablet_writers_lock);
            auto tablet_writer_it = _tablet_writers.find(tablet_id);
            if (tablet_writer_it == _tablet_writers.end()) {
                return Status::InternalError("unknown tablet to append data, tablet={}", tablet_id);
            }
            tablet_writer = tablet_writer_it->second.get();
        }
        RETURN_IF_ERROR(_prepare_adaptive_random_bucket_writer(tablet_writer));

        TabletAddRowsPayload rows {.row_idxs = row_idxs};
        if (request.row_binlog_lsns_size() > 0) {
            rows.row_binlog_lsns.reserve(row_idxs.size());
            for (auto row_idx : row_idxs) {
                rows.row_binlog_lsns.emplace_back(request.row_binlog_lsns(row_idx));
            }
        }
        bool memtable_flushed = false;
        Status st = tablet_writer->write(&send_data, rows, &memtable_flushed);
        if (!st.ok()) {
            auto err_msg =
                    fmt::format("tablet writer write failed, tablet_id={}, txn_id={}, err={}",
                                tablet_id, _txn_id, st.to_string());
            LOG(WARNING) << err_msg;
            PTabletError* error = tablet_errors->Add();
            error->set_tablet_id(tablet_id);
            error->set_msg(err_msg);
            static_cast<void>(tablet_writer->cancel_with_status(st));
            _add_broken_tablet(tablet_id);
            return Status::OK();
        }

        VLOG_DEBUG << "FIND_TABLET_RANDOM_BUCKET: route+write done"
                   << ", load_id=" << _load_id << ", index_id=" << _index_id
                   << ", sender_id=" << request.sender_id()
                   << ", packet_seq=" << request.packet_seq() << ", partition_id=" << partition_id
                   << ", tablet_id=" << tablet_id << ", row_count=" << row_idxs.size()
                   << ", memtable_flushed=" << memtable_flushed;
        if (memtable_flushed) {
            _adaptive_random_bucket_state->rotate_by_tablet(partition_id, tablet_id);
        }
        tablet_writer->set_tablet_load_rowset_num_info(tablet_load_infos);
        return Status::OK();
    };

    SCOPED_TIMER(_write_block_timer);
    for (const auto& [partition_id, row_idxs] : partition_to_rowidxs) {
        RETURN_IF_ERROR(write_partition_data(partition_id, row_idxs));
    }

    {
        std::lock_guard<std::mutex> l(_lock);
        _next_seqs[request.sender_id()] = cur_seq + 1;
    }
    return Status::OK();
}

Status BaseTabletsChannel::_prepare_adaptive_random_bucket_writer(BaseDeltaWriter*) {
    return Status::OK();
}

Status BaseTabletsChannel::_build_partition_to_rowidxs_for_adaptive_random_bucket(
        const PTabletWriterAddBlockRequest& request,
        std::unordered_map<int64_t, DorisVector<uint32_t>>* partition_to_rowidxs) {
    if (_adaptive_random_bucket_state == nullptr) {
        return Status::InternalError(
                "adaptive random bucket state is not initialized, load_id={}, index_id={}, "
                "packet_seq={}",
                print_id(_load_id), _index_id, request.packet_seq());
    }
    if (request.partition_ids_size() == 0) {
        return Status::InternalError(
                "empty partition ids for adaptive random bucket add block, load_id={}, "
                "index_id={}, packet_seq={}",
                print_id(_load_id), _index_id, request.packet_seq());
    }
    for (uint32_t i = 0; i < request.partition_ids_size(); ++i) {
        int64_t partition_id = request.partition_ids(i);
        auto it = partition_to_rowidxs->find(partition_id);
        if (it == partition_to_rowidxs->end()) {
            partition_to_rowidxs->emplace(partition_id, std::initializer_list<uint32_t> {i});
        } else {
            it->second.emplace_back(i);
        }
    }
    return Status::OK();
}

Status TabletsChannel::add_batch(const PTabletWriterAddBlockRequest& request,
                                 PTabletWriterAddBlockResult* response) {
    SCOPED_TIMER(_add_batch_timer);
    int64_t cur_seq = 0;
    if (_add_batch_number_counter) {
        _add_batch_number_counter->update(1);
    }

    auto status = _get_current_seq(cur_seq, request);
    if (UNLIKELY(!status.ok())) {
        return status;
    }

    if (request.packet_seq() < cur_seq) {
        LOG(INFO) << "packet has already recept before, expect_seq=" << cur_seq
                  << ", recept_seq=" << request.packet_seq();
        return Status::OK();
    }

    // Adaptive random bucket add-block RPCs carry partition ids. The receiver maps rows to
    // the current tablet and advances its selected-tablet state after memtable flush.
    if (request.is_adaptive_random_bucket()) {
        std::unordered_map<int64_t /* partition_id */, DorisVector<uint32_t> /* row index */>
                partition_to_rowidxs;
        RETURN_IF_ERROR(_build_partition_to_rowidxs_for_adaptive_random_bucket(
                request, &partition_to_rowidxs));
        return _write_block_data_for_adaptive_random_bucket(request, cur_seq, partition_to_rowidxs,
                                                            response);
    }

    std::unordered_map<int64_t /* tablet_id */, TabletAddRowsPayload> tablet_to_rows;
    _build_tablet_to_rows(request, &tablet_to_rows);

    return _write_block_data(request, cur_seq, tablet_to_rows, response);
}

void BaseTabletsChannel::_add_broken_tablet(int64_t tablet_id) {
    std::unique_lock<std::shared_mutex> wlock(_broken_tablets_lock);
    _broken_tablets.insert(tablet_id);
}

bool BaseTabletsChannel::_is_broken_tablet(int64_t tablet_id) const {
    return _broken_tablets.find(tablet_id) != _broken_tablets.end();
}

void BaseTabletsChannel::_build_tablet_to_rows(
        const PTabletWriterAddBlockRequest& request,
        std::unordered_map<int64_t, TabletAddRowsPayload>* tablet_to_rows) {
    // just add a coarse-grained read lock here rather than each time when visiting _broken_tablets
    // tests show that a relatively coarse-grained read lock here performs better under multicore scenario
    // see: https://github.com/apache/doris/pull/28552
    std::shared_lock<std::shared_mutex> rlock(_broken_tablets_lock);
    bool has_row_binlog_lsn = request.row_binlog_lsns_size() > 0;
    if (request.is_single_tablet_block()) {
        // The cloud mode need the tablet ids to prepare rowsets.
        int64_t tablet_id = request.tablet_ids(0);
        auto& rows = (*tablet_to_rows)[tablet_id];
        rows.row_idxs.emplace_back(0);
        if (has_row_binlog_lsn) {
            rows.row_binlog_lsns.emplace_back(request.row_binlog_lsns(0));
        }
        return;
    }
    for (uint32_t i = 0; i < request.tablet_ids_size(); ++i) {
        int64_t tablet_id = request.tablet_ids(i);
        if (_is_broken_tablet(tablet_id)) {
            // skip broken tablets
            VLOG_PROGRESS << "skip broken tablet tablet=" << tablet_id;
            continue;
        }
        auto& rows = (*tablet_to_rows)[tablet_id];
        rows.row_idxs.emplace_back(i);
        if (has_row_binlog_lsn) {
            rows.row_binlog_lsns.emplace_back(request.row_binlog_lsns(i));
        }
    }
}

} // namespace doris

Coverage Report

Created: 2026-08-14 13:39