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

#include <fmt/format.h>

#include <gen_cpp/olap_file.pb.h>

#include <limits.h>

#include <algorithm>

#include <atomic>

#include <condition_variable>

#include <filesystem>

#include <map>

#include <memory>

#include <mutex>

#include <sstream>

#include <string>

#include <utility>

#include "beta_rowset_writer.h"

#include "common/compiler_util.h" // IWYU pragma: keep

#include "common/logging.h"

#include "gutil/stringprintf.h"

#include "gutil/strings/substitute.h"

#include "io/fs/file_system.h"

#include "io/fs/file_writer.h"

#include "io/io_common.h"

#include "olap/data_dir.h"

#include "olap/iterators.h"

#include "olap/merger.h"

#include "olap/olap_common.h"

#include "olap/olap_define.h"

#include "olap/rowset/beta_rowset.h"

#include "olap/rowset/rowset_meta.h"

#include "olap/rowset/rowset_writer_context.h"

#include "olap/rowset/segment_v2/inverted_index_cache.h"

#include "olap/rowset/segment_v2/inverted_index_desc.h"

#include "olap/rowset/segment_v2/segment.h"

#include "olap/rowset/segment_v2/segment_writer.h"

#include "olap/schema.h"

#include "olap/storage_engine.h"

#include "olap/tablet_reader.h"

#include "olap/tablet_schema.h"

#include "runtime/memory/global_memory_arbitrator.h"

#include "runtime/thread_context.h"

#include "util/debug_points.h"

#include "util/mem_info.h"

#include "util/time.h"

#include "vec/olap/vertical_block_reader.h"

#include "vec/olap/vertical_merge_iterator.h"

namespace doris {

using namespace ErrorCode;

SegcompactionWorker::SegcompactionWorker(BetaRowsetWriter* writer) : _writer(writer) {}

void SegcompactionWorker::init_mem_tracker(const RowsetWriterContext& rowset_writer_context) {

    _seg_compact_mem_tracker = MemTrackerLimiter::create_shared(

            MemTrackerLimiter::Type::COMPACTION,

            fmt::format("segcompaction-txnID_{}-loadID_{}-tabletID_{}-indexID_{}-"

                        "partitionID_{}-version_{}",

                        std::to_string(rowset_writer_context.txn_id),

                        print_id(rowset_writer_context.load_id),

                        std::to_string(rowset_writer_context.tablet_id),

                        std::to_string(rowset_writer_context.index_id),

                        std::to_string(rowset_writer_context.partition_id),

                        rowset_writer_context.version.to_string()));

}

Status SegcompactionWorker::_get_segcompaction_reader(

        SegCompactionCandidatesSharedPtr segments, TabletSharedPtr tablet,

        std::shared_ptr<Schema> schema, OlapReaderStatistics* stat,

        vectorized::RowSourcesBuffer& row_sources_buf, bool is_key,

        std::vector<uint32_t>& return_columns, std::vector<uint32_t>& key_group_cluster_key_idxes,

        std::unique_ptr<vectorized::VerticalBlockReader>* reader) {

    const auto& ctx = _writer->_context;

    bool record_rowids = need_convert_delete_bitmap() && is_key;

    StorageReadOptions read_options;

    read_options.stats = stat;

    read_options.use_page_cache = false;

    read_options.tablet_schema = ctx.tablet_schema;

    read_options.record_rowids = record_rowids;

    if (!tablet->tablet_schema()->cluster_key_uids().empty()) {

        DeleteBitmapPtr delete_bitmap = std::make_shared<DeleteBitmap>(tablet->tablet_id());

        RETURN_IF_ERROR(tablet->calc_delete_bitmap_between_segments(ctx.rowset_id, *segments,

                                                                    delete_bitmap));

        for (auto& seg_ptr : *segments) {

            auto d = delete_bitmap->get_agg(

                    {ctx.rowset_id, seg_ptr->id(), DeleteBitmap::TEMP_VERSION_COMMON});

            if (d->isEmpty()) {

                continue; // Empty delete bitmap for the segment

            }

            read_options.delete_bitmap.emplace(seg_ptr->id(), std::move(d));

        }

    }

    std::vector<std::unique_ptr<RowwiseIterator>> seg_iterators;

    std::map<uint32_t, uint32_t> segment_rows;

    for (auto& seg_ptr : *segments) {

        std::unique_ptr<RowwiseIterator> iter;

        auto s = seg_ptr->new_iterator(schema, read_options, &iter);

        if (!s.ok()) {

            return Status::Error<INIT_FAILED>("failed to create iterator[{}]: {}", seg_ptr->id(),

                                              s.to_string());

        }

        seg_iterators.push_back(std::move(iter));

        segment_rows.emplace(seg_ptr->id(), seg_ptr->num_rows());

    }

    if (record_rowids && _rowid_conversion != nullptr) {

        _rowid_conversion->reset_segment_map(segment_rows);

    }

    *reader = std::make_unique<vectorized::VerticalBlockReader>(&row_sources_buf);

    TabletReader::ReaderParams reader_params;

    reader_params.is_segcompaction = true;

    reader_params.segment_iters_ptr = &seg_iterators;

    // no reader_params.version shouldn't break segcompaction

    reader_params.tablet_schema = ctx.tablet_schema;

    reader_params.tablet = tablet;

    reader_params.return_columns = return_columns;

    reader_params.is_key_column_group = is_key;

    reader_params.use_page_cache = false;

    reader_params.record_rowids = record_rowids;

    reader_params.key_group_cluster_key_idxes = key_group_cluster_key_idxes;

    return (*reader)->init(reader_params, nullptr);

}

std::unique_ptr<segment_v2::SegmentWriter> SegcompactionWorker::_create_segcompaction_writer(

        uint32_t begin, uint32_t end) {

    Status status;

    std::unique_ptr<segment_v2::SegmentWriter> writer = nullptr;

    status = _create_segment_writer_for_segcompaction(&writer, begin, end);

    if (!status.ok() || writer == nullptr) {

        LOG(ERROR) << "failed to create segment writer for begin:" << begin << " end:" << end

                   << " status:" << status;

        return nullptr;

    } else {

        return writer;

    }

}

Status SegcompactionWorker::_delete_original_segments(uint32_t begin, uint32_t end) {

    DCHECK(_writer->rowset_meta()->is_local());

    const auto& fs = io::global_local_filesystem();

    auto ctx = _writer->_context;

    auto schema = ctx.tablet_schema;

    for (uint32_t i = begin; i <= end; ++i) {

        auto seg_path = local_segment_path(ctx.tablet_path, ctx.rowset_id.to_string(), i);

        // Even if an error is encountered, these files that have not been cleaned up

        // will be cleaned up by the GC background. So here we only print the error

        // message when we encounter an error.

        RETURN_NOT_OK_STATUS_WITH_WARN(fs->delete_file(seg_path),

                                       strings::Substitute("Failed to delete file=$0", seg_path));

        if (schema->has_inverted_index() &&

            schema->get_inverted_index_storage_format() >= InvertedIndexStorageFormatPB::V2) {

            auto idx_path = InvertedIndexDescriptor::get_index_file_path_v2(

                    InvertedIndexDescriptor::get_index_file_path_prefix(seg_path));

            VLOG_DEBUG << "segcompaction index. delete file " << idx_path;

            RETURN_NOT_OK_STATUS_WITH_WARN(

                    fs->delete_file(idx_path),

                    strings::Substitute("Failed to delete file=$0", idx_path));

        }

        // Delete inverted index files

        for (auto&& column : schema->columns()) {

            if (const auto* index_info = schema->inverted_index(*column); index_info != nullptr) {

                auto index_id = index_info->index_id();

                if (schema->get_inverted_index_storage_format() ==

                    InvertedIndexStorageFormatPB::V1) {

                    auto idx_path = InvertedIndexDescriptor::get_index_file_path_v1(

                            InvertedIndexDescriptor::get_index_file_path_prefix(seg_path), index_id,

                            index_info->get_index_suffix());

                    VLOG_DEBUG << "segcompaction index. delete file " << idx_path;

                    RETURN_NOT_OK_STATUS_WITH_WARN(

                            fs->delete_file(idx_path),

                            strings::Substitute("Failed to delete file=$0", idx_path));

                }

                // Erase the origin index file cache

                auto idx_file_cache_key = InvertedIndexDescriptor::get_index_file_cache_key(

                        InvertedIndexDescriptor::get_index_file_path_prefix(seg_path), index_id,

                        index_info->get_index_suffix());

                RETURN_IF_ERROR(InvertedIndexSearcherCache::instance()->erase(idx_file_cache_key));

            }

        }

    }

    return Status::OK();

}

Status SegcompactionWorker::_check_correctness(OlapReaderStatistics& reader_stat,

                                               Merger::Statistics& merger_stat, uint32_t begin,

                                               uint32_t end, bool is_mow_with_cluster_keys) {

    uint64_t raw_rows_read = reader_stat.raw_rows_read; /* total rows read before merge */

    uint64_t rows_del_by_bitmap = reader_stat.rows_del_by_bitmap;

    uint64_t sum_src_row = 0; /* sum of rows in each involved source segments */

    uint64_t filtered_rows = merger_stat.filtered_rows; /* rows filtered by del conditions */

    uint64_t output_rows = merger_stat.output_rows;     /* rows after merge */

    uint64_t merged_rows = merger_stat.merged_rows;     /* dup key merged by unique/agg */

    {

        std::lock_guard<std::mutex> lock(_writer->_segid_statistics_map_mutex);

        for (int i = begin; i <= end; ++i) {

            sum_src_row += _writer->_segid_statistics_map[i].row_num;

        }

    }

    DBUG_EXECUTE_IF("SegcompactionWorker._check_correctness_wrong_sum_src_row", { sum_src_row++; });

    uint64_t raw_rows = raw_rows_read;

    if (is_mow_with_cluster_keys) {

        raw_rows += rows_del_by_bitmap;

    }

    if (raw_rows != sum_src_row) {

        return Status::Error<CHECK_LINES_ERROR>(

                "segcompaction read row num does not match source. expect read row:{}, actual read "

                "row:{}(raw_rows_read: {}, rows_del_by_bitmap: {})",

                sum_src_row, raw_rows, raw_rows_read, rows_del_by_bitmap);

    }

    DBUG_EXECUTE_IF("SegcompactionWorker._check_correctness_wrong_merged_rows", { merged_rows++; });

    if ((output_rows + merged_rows) != raw_rows_read) {

        return Status::Error<CHECK_LINES_ERROR>(

                "segcompaction total row num does not match after merge. expect total row:{},  "

                "actual total row:{}, (output_rows:{},merged_rows:{})",

                raw_rows_read, output_rows + merged_rows, output_rows, merged_rows);

    }

    DBUG_EXECUTE_IF("SegcompactionWorker._check_correctness_wrong_filtered_rows",

                    { filtered_rows++; });

    if (filtered_rows != 0) {

        return Status::Error<CHECK_LINES_ERROR>(

                "segcompaction should not have filtered rows but actual filtered rows:{}",

                filtered_rows);

    }

    return Status::OK();

}

Status SegcompactionWorker::_create_segment_writer_for_segcompaction(

        std::unique_ptr<segment_v2::SegmentWriter>* writer, uint32_t begin, uint32_t end) {

    return _writer->create_segment_writer_for_segcompaction(writer, begin, end);

}

Status SegcompactionWorker::_do_compact_segments(SegCompactionCandidatesSharedPtr segments) {

    DCHECK(_seg_compact_mem_tracker != nullptr);

    SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(_seg_compact_mem_tracker);

    /* throttle segcompaction task if memory depleted */

    if (GlobalMemoryArbitrator::is_exceed_soft_mem_limit(GB_EXCHANGE_BYTE)) {

        return Status::Error<FETCH_MEMORY_EXCEEDED>("skip segcompaction due to memory shortage");

    }

    uint32_t begin = (*(segments->begin()))->id();

    uint32_t end = (*(segments->end() - 1))->id();

    uint64_t begin_time = GetCurrentTimeMicros();

    uint64_t index_size = 0;

    uint64_t total_index_size = 0;

    auto ctx = _writer->_context;

    auto writer = _create_segcompaction_writer(begin, end);

    if (UNLIKELY(writer == nullptr)) {

        return Status::Error<SEGCOMPACTION_INIT_WRITER>("failed to get segcompaction writer");

    }

    DCHECK(ctx.tablet);

    auto tablet = std::static_pointer_cast<Tablet>(ctx.tablet);

    if (need_convert_delete_bitmap() && _rowid_conversion == nullptr) {

        _rowid_conversion = std::make_unique<SimpleRowIdConversion>(_writer->rowset_id());

    }

    std::vector<std::vector<uint32_t>> column_groups;

    std::vector<uint32_t> key_group_cluster_key_idxes;

    Merger::vertical_split_columns(*ctx.tablet_schema, &column_groups,

                                   &key_group_cluster_key_idxes);

    vectorized::RowSourcesBuffer row_sources_buf(tablet->tablet_id(), tablet->tablet_path(),

                                                 ReaderType::READER_SEGMENT_COMPACTION);

    KeyBoundsPB key_bounds;

    Merger::Statistics key_merger_stats;

    OlapReaderStatistics key_reader_stats;

    /* compact group one by one */

    for (auto i = 0; i < column_groups.size(); ++i) {

        VLOG_NOTICE << "row source size: " << row_sources_buf.total_size();

        bool is_key = (i == 0);

        std::vector<uint32_t> column_ids = column_groups[i];

        writer->clear();

        RETURN_IF_ERROR(writer->init(column_ids, is_key));

        auto schema = std::make_shared<Schema>(ctx.tablet_schema->columns(), column_ids);

        OlapReaderStatistics reader_stats;

        std::unique_ptr<vectorized::VerticalBlockReader> reader;

        auto s =

                _get_segcompaction_reader(segments, tablet, schema, &reader_stats, row_sources_buf,

                                          is_key, column_ids, key_group_cluster_key_idxes, &reader);

        if (UNLIKELY(reader == nullptr || !s.ok())) {

            return Status::Error<SEGCOMPACTION_INIT_READER>(

                    "failed to get segcompaction reader. err: {}", s.to_string());

        }

        Merger::Statistics merger_stats;

        RETURN_IF_ERROR(Merger::vertical_compact_one_group(

                tablet->tablet_id(), ReaderType::READER_SEGMENT_COMPACTION, *ctx.tablet_schema,

                is_key, column_ids, &row_sources_buf, *reader, *writer, &merger_stats, &index_size,

                key_bounds, _rowid_conversion.get()));

        total_index_size += index_size;

        if (is_key) {

            RETURN_IF_ERROR(row_sources_buf.flush());

            key_merger_stats = merger_stats;

            key_reader_stats = reader_stats;

        }

        RETURN_IF_ERROR(row_sources_buf.seek_to_begin());

    }

    /* check row num after merge/aggregation */

    bool is_mow_with_cluster_keys = !tablet->tablet_schema()->cluster_key_uids().empty();

    RETURN_NOT_OK_STATUS_WITH_WARN(_check_correctness(key_reader_stats, key_merger_stats, begin,

                                                      end, is_mow_with_cluster_keys),

                                   "check correctness failed");

    {

        std::lock_guard<std::mutex> lock(_writer->_segid_statistics_map_mutex);

        _writer->_clear_statistics_for_deleting_segments_unsafe(begin, end);

    }

    RETURN_IF_ERROR(

            _writer->flush_segment_writer_for_segcompaction(&writer, total_index_size, key_bounds));

    if (_file_writer != nullptr && _file_writer->state() != io::FileWriter::State::CLOSED) {

        RETURN_IF_ERROR(_file_writer->close());

    }

    RETURN_IF_ERROR(_delete_original_segments(begin, end));

    if (_rowid_conversion != nullptr) {

        convert_segment_delete_bitmap(ctx.mow_context->delete_bitmap, begin, end,

                                      _writer->_num_segcompacted);

    }

    RETURN_IF_ERROR(_writer->_rename_compacted_segments(begin, end));

    if (_inverted_index_file_writer != nullptr) {

        _inverted_index_file_writer.reset();

    }

    if (VLOG_DEBUG_IS_ON) {

        _writer->vlog_buffer.clear();

        for (const auto& entry : std::filesystem::directory_iterator(ctx.tablet_path)) {

            fmt::format_to(_writer->vlog_buffer, "[{}]", string(entry.path()));

        }

        VLOG_DEBUG << "tablet_id:" << ctx.tablet_id << " rowset_id:" << ctx.rowset_id

                   << "_segcompacted_point:" << _writer->_segcompacted_point

                   << " _num_segment:" << _writer->_num_segment

                   << " _num_segcompacted:" << _writer->_num_segcompacted

                   << " list directory:" << fmt::to_string(_writer->vlog_buffer);

    }

    _writer->_segcompacted_point += (end - begin + 1);

    uint64_t elapsed = GetCurrentTimeMicros() - begin_time;

    LOG(INFO) << "segcompaction completed. tablet_id:" << ctx.tablet_id

              << " rowset_id:" << ctx.rowset_id << " elapsed time:" << elapsed

              << "us. update segcompacted_point:" << _writer->_segcompacted_point

              << " segment num:" << segments->size() << " begin:" << begin << " end:" << end;

    return Status::OK();

}

void SegcompactionWorker::compact_segments(SegCompactionCandidatesSharedPtr segments) {

    Status status = Status::OK();

    if (_is_compacting_state_mutable.exchange(false)) {

        status = _do_compact_segments(segments);

    } else {

        // note: be aware that _writer maybe released when the task is cancelled

        LOG(INFO) << "segcompaction worker is cancelled, skipping segcompaction task";

        return;

    }

    if (!status.ok()) {

        int16_t errcode = status.code();

        switch (errcode) {

        case FETCH_MEMORY_EXCEEDED:

        case SEGCOMPACTION_INIT_READER:

        case SEGCOMPACTION_INIT_WRITER:

            LOG(WARNING) << "segcompaction failed, try next time:" << status;

            break;

        default:

            auto ctx = _writer->_context;

            LOG(WARNING) << "segcompaction fatal, terminating the write job."

                         << " tablet_id:" << ctx.tablet_id << " rowset_id:" << ctx.rowset_id

                         << " status:" << status;

            // status will be checked by the next trigger of segcompaction or the final wait

            _writer->_segcompaction_status.store(ErrorCode::INTERNAL_ERROR);

        }

    }

    DCHECK_EQ(_writer->_is_doing_segcompaction, true);

    {

        std::lock_guard lk(_writer->_is_doing_segcompaction_lock);

        _writer->_is_doing_segcompaction = false;

        _writer->_segcompacting_cond.notify_all();

    }

    _is_compacting_state_mutable = true;

}

bool SegcompactionWorker::need_convert_delete_bitmap() {

    if (_writer == nullptr) {

        return false;

    }

    auto tablet = _writer->context().tablet;

    return tablet != nullptr && tablet->keys_type() == KeysType::UNIQUE_KEYS &&

           tablet->enable_unique_key_merge_on_write() &&

           tablet->tablet_schema()->has_sequence_col();

}

void SegcompactionWorker::convert_segment_delete_bitmap(DeleteBitmapPtr src_delete_bitmap,

                                                        uint32_t src_seg_id, uint32_t dest_seg_id) {

    // lazy init

    if (nullptr == _converted_delete_bitmap) {

        _converted_delete_bitmap = std::make_shared<DeleteBitmap>(_writer->context().tablet_id);

    }

    auto rowset_id = _writer->context().rowset_id;

    const auto* seg_map =

            src_delete_bitmap->get({rowset_id, src_seg_id, DeleteBitmap::TEMP_VERSION_COMMON});

    if (seg_map != nullptr) {

        _converted_delete_bitmap->set({rowset_id, dest_seg_id, DeleteBitmap::TEMP_VERSION_COMMON},

                                      *seg_map);

    }

}

void SegcompactionWorker::convert_segment_delete_bitmap(DeleteBitmapPtr src_delete_bitmap,

                                                        uint32_t src_begin, uint32_t src_end,

                                                        uint32_t dst_seg_id) {

    // lazy init

    if (nullptr == _converted_delete_bitmap) {

        _converted_delete_bitmap = std::make_shared<DeleteBitmap>(_writer->context().tablet_id);

    }

    auto rowset_id = _writer->context().rowset_id;

    RowLocation src(rowset_id, 0, 0);

    for (uint32_t seg_id = src_begin; seg_id <= src_end; seg_id++) {

        const auto* seg_map =

                src_delete_bitmap->get({rowset_id, seg_id, DeleteBitmap::TEMP_VERSION_COMMON});

        if (!seg_map) {

            continue;

        }

        src.segment_id = seg_id;

        for (unsigned int row_id : *seg_map) {

            src.row_id = row_id;

            auto dst_row_id = _rowid_conversion->get(src);

            if (dst_row_id < 0) {

                continue;

            }

            _converted_delete_bitmap->add(

                    {rowset_id, dst_seg_id, DeleteBitmap::TEMP_VERSION_COMMON}, dst_row_id);

        }

    }

}

bool SegcompactionWorker::cancel() {

    // return true if the task is canncellable (actual compaction is not started)

    // return false when the task is not cancellable (it is in the middle of segcompaction)

    return _is_compacting_state_mutable.exchange(false);

}

} // namespace doris