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

#include <gen_cpp/cloud.pb.h>

#include <gen_cpp/olap_file.pb.h>

#include <gtest/gtest.h>

#include <chrono>

#include <memory>

#include <random>

#include <string>

#include <thread>

#include "common/config.h"

#include "common/logging.h"

#include "common/simple_thread_pool.h"

#include "common/util.h"

#include "cpp/sync_point.h"

#include "meta-service/keys.h"

#include "meta-service/mem_txn_kv.h"

#include "meta-service/meta_service.h"

#include "meta-service/txn_kv_error.h"

#include "mock_accessor.h"

#include "mock_resource_manager.h"

#include "rate-limiter/rate_limiter.h"

#include "recycler/checker.h"

#include "recycler/storage_vault_accessor.h"

#include "recycler/util.h"

#include "recycler/white_black_list.h"

using namespace doris;

static const std::string instance_id = "instance_id_recycle_test";

static int64_t current_time = 0;

static constexpr int64_t db_id = 1000;

static doris::cloud::RecyclerThreadPoolGroup thread_group;

int main(int argc, char** argv) {

auto conf_file = "doris_cloud.conf";

if (!cloud::config::init(conf_file, true)) {

std::cerr << "failed to init config file, conf=" << conf_file << std::endl;

return -1;

}

if (!cloud::init_glog("recycler")) {

std::cerr << "failed to init glog" << std::endl;

return -1;

}

using namespace std::chrono;

current_time = duration_cast<seconds>(system_clock::now().time_since_epoch()).count();

config::recycler_sleep_before_scheduling_seconds = 0; // we dont have to wait in UT

::testing::InitGoogleTest(&argc, argv);

auto s3_producer_pool = std::make_shared<SimpleThreadPool>(config::recycle_pool_parallelism);

s3_producer_pool->start();

auto recycle_tablet_pool = std::make_shared<SimpleThreadPool>(config::recycle_pool_parallelism);

recycle_tablet_pool->start();

auto group_recycle_function_pool =

std::make_shared<SimpleThreadPool>(config::recycle_pool_parallelism);

group_recycle_function_pool->start();

thread_group =

RecyclerThreadPoolGroup(std::move(s3_producer_pool), std::move(recycle_tablet_pool),

std::move(group_recycle_function_pool));

return RUN_ALL_TESTS();

}

namespace doris::cloud {

TEST(RecyclerTest, WhiteBlackList) {

WhiteBlackList filter;

EXPECT_FALSE(filter.filter_out("instance1"));

EXPECT_FALSE(filter.filter_out("instance2"));

filter.reset({}, {"instance1", "instance2"});

EXPECT_TRUE(filter.filter_out("instance1"));

EXPECT_TRUE(filter.filter_out("instance2"));

EXPECT_FALSE(filter.filter_out("instance3"));

filter.reset({"instance1", "instance2"}, {});

EXPECT_FALSE(filter.filter_out("instance1"));

EXPECT_FALSE(filter.filter_out("instance2"));

EXPECT_TRUE(filter.filter_out("instance3"));

filter.reset({"instance1"}, {"instance1"}); // whitelist overrides blacklist

EXPECT_FALSE(filter.filter_out("instance1"));

EXPECT_TRUE(filter.filter_out("instance2"));

}

static std::string next_rowset_id() {

static int64_t cnt = 0;

return std::to_string(++cnt);

}

static doris::RowsetMetaCloudPB create_rowset(const std::string& resource_id, int64_t tablet_id,

int64_t index_id, int num_segments,

const doris::TabletSchemaCloudPB& schema,

int64_t txn_id = 0) {

doris::RowsetMetaCloudPB rowset;

rowset.set_rowset_id(0); // useless but required

rowset.set_rowset_id_v2(next_rowset_id());

rowset.set_txn_id(txn_id);

rowset.set_num_segments(num_segments);

rowset.set_tablet_id(tablet_id);

rowset.set_index_id(index_id);

rowset.set_resource_id(resource_id);

rowset.set_schema_version(schema.schema_version());

rowset.mutable_tablet_schema()->CopyFrom(schema);

return rowset;

}

static doris::RowsetMetaCloudPB create_rowset(const std::string& resource_id, int64_t tablet_id,

int64_t index_id, int num_segments,

const doris::TabletSchemaCloudPB& schema,

RowsetStatePB rowset_state, int64_t txn_id = 0) {

doris::RowsetMetaCloudPB rowset;

rowset.set_rowset_id(0); // useless but required

rowset.set_rowset_id_v2(next_rowset_id());

rowset.set_txn_id(txn_id);

rowset.set_num_segments(num_segments);

rowset.set_tablet_id(tablet_id);

rowset.set_index_id(index_id);

rowset.set_resource_id(resource_id);

rowset.set_schema_version(schema.schema_version());

rowset.mutable_tablet_schema()->CopyFrom(schema);

rowset.set_rowset_state(rowset_state);

return rowset;

}

static int create_recycle_rowset(TxnKv* txn_kv, StorageVaultAccessor* accessor,

const doris::RowsetMetaCloudPB& rowset, RecycleRowsetPB::Type type,

bool write_schema_kv) {

std::string key;

std::string val;

RecycleRowsetKeyInfo key_info {instance_id, rowset.tablet_id(), rowset.rowset_id_v2()};

recycle_rowset_key(key_info, &key);

RecycleRowsetPB rowset_pb;

rowset_pb.set_creation_time(current_time);

if (type != RecycleRowsetPB::UNKNOWN) {

rowset_pb.set_type(type);

rowset_pb.mutable_rowset_meta()->CopyFrom(rowset);

if (write_schema_kv) { // Detach schema

rowset_pb.mutable_rowset_meta()->set_allocated_tablet_schema(nullptr);

}

} else { // old version RecycleRowsetPB

rowset_pb.set_tablet_id(rowset.tablet_id());

rowset_pb.set_resource_id(rowset.resource_id());

}

rowset_pb.SerializeToString(&val);

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(key, val);

std::string schema_key, schema_val;

if (write_schema_kv) {

meta_schema_key({instance_id, rowset.index_id(), rowset.schema_version()}, &schema_key);

rowset.tablet_schema().SerializeToString(&schema_val);

txn->put(schema_key, schema_val);

}

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

for (int i = 0; i < rowset.num_segments(); ++i) {

auto path = segment_path(rowset.tablet_id(), rowset.rowset_id_v2(), i);

accessor->put_file(path, "");

for (auto& index : rowset.tablet_schema().index()) {

auto path = inverted_index_path_v1(rowset.tablet_id(), rowset.rowset_id_v2(), i,

index.index_id(), index.index_suffix_name());

accessor->put_file(path, "");

}

}

return 0;

}

static int create_tmp_rowset(TxnKv* txn_kv, StorageVaultAccessor* accessor,

const doris::RowsetMetaCloudPB& rowset, bool write_schema_kv,

bool is_inverted_idx_v2 = false) {

std::string key, val;

meta_rowset_tmp_key({instance_id, rowset.txn_id(), rowset.tablet_id()}, &key);

if (write_schema_kv) {

auto rowset_copy = rowset;

rowset_copy.clear_tablet_schema();

rowset_copy.SerializeToString(&val);

} else {

rowset.SerializeToString(&val);

}

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(key, val);

std::string schema_key, schema_val;

if (write_schema_kv) {

meta_schema_key({instance_id, rowset.index_id(), rowset.schema_version()}, &schema_key);

rowset.tablet_schema().SerializeToString(&schema_val);

txn->put(schema_key, schema_val);

}

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

for (int i = 0; i < rowset.num_segments(); ++i) {

auto path = segment_path(rowset.tablet_id(), rowset.rowset_id_v2(), i);

accessor->put_file(path, path);

for (auto& index : rowset.tablet_schema().index()) {

std::string path;

if (!is_inverted_idx_v2) {

path = inverted_index_path_v1(rowset.tablet_id(), rowset.rowset_id_v2(), i,

index.index_id(), index.index_suffix_name());

} else {

path = inverted_index_path_v2(rowset.tablet_id(), rowset.rowset_id_v2(), i);

}

accessor->put_file(path, path);

}

}

return 0;

}

static int create_committed_rowset(TxnKv* txn_kv, StorageVaultAccessor* accessor,

const std::string& resource_id, int64_t tablet_id,

int64_t version, int num_segments = 1,

int num_inverted_indexes = 1) {

std::string key;

std::string val;

auto rowset_id = next_rowset_id();

MetaRowsetKeyInfo key_info {instance_id, tablet_id, version};

meta_rowset_key(key_info, &key);

doris::RowsetMetaCloudPB rowset_pb;

rowset_pb.set_rowset_id(0); // useless but required

rowset_pb.set_rowset_id_v2(rowset_id);

rowset_pb.set_num_segments(num_segments);

rowset_pb.set_tablet_id(tablet_id);

rowset_pb.set_resource_id(resource_id);

rowset_pb.set_creation_time(current_time);

if (num_inverted_indexes > 0) {

auto schema = rowset_pb.mutable_tablet_schema();

for (int i = 0; i < num_inverted_indexes; ++i) {

schema->add_index()->set_index_id(i);

}

}

rowset_pb.SerializeToString(&val);

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(key, val);

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

for (int i = 0; i < num_segments; ++i) {

auto path = segment_path(tablet_id, rowset_id, i);

accessor->put_file(path, "");

for (int j = 0; j < num_inverted_indexes; ++j) {

auto path = inverted_index_path_v1(tablet_id, rowset_id, i, j, "");

accessor->put_file(path, "");

}

}

return 0;

}

static int create_committed_rowset_with_rowset_id(TxnKv* txn_kv, StorageVaultAccessor* accessor,

const std::string& resource_id, int64_t tablet_id,

int64_t start_version, int64_t end_version,

std::string rowset_id, bool segments_overlap,

int num_segments) {

std::string key;

std::string val;

MetaRowsetKeyInfo key_info {instance_id, tablet_id, end_version};

meta_rowset_key(key_info, &key);

doris::RowsetMetaCloudPB rowset_pb;

rowset_pb.set_rowset_id(0); // useless but required

rowset_pb.set_rowset_id_v2(rowset_id);

rowset_pb.set_num_segments(num_segments);

rowset_pb.set_tablet_id(tablet_id);

rowset_pb.set_resource_id(resource_id);

rowset_pb.set_creation_time(current_time);

rowset_pb.set_start_version(start_version);

rowset_pb.set_end_version(end_version);

rowset_pb.set_segments_overlap_pb(segments_overlap ? OVERLAPPING : NONOVERLAPPING);

rowset_pb.SerializeToString(&val);

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(key, val);

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

for (int i = 0; i < num_segments; ++i) {

auto path = segment_path(tablet_id, rowset_id, i);

accessor->put_file(path, "");

}

return 0;

}

static void create_delete_bitmaps(Transaction* txn, int64_t tablet_id, std::string rowset_id,

int64_t start_version, int64_t end_version) {

for (int64_t ver {start_version}; ver <= end_version; ver++) {

auto key = meta_delete_bitmap_key({instance_id, tablet_id, rowset_id, ver, 0});

std::string val {"test_data"};

txn->put(key, val);

}

}

static int create_tablet(TxnKv* txn_kv, int64_t table_id, int64_t index_id, int64_t partition_id,

int64_t tablet_id, bool is_mow = false) {

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

doris::TabletMetaCloudPB tablet_meta;

tablet_meta.set_tablet_id(tablet_id);

tablet_meta.set_enable_unique_key_merge_on_write(is_mow);

auto val = tablet_meta.SerializeAsString();

auto key = meta_tablet_key({instance_id, table_id, index_id, partition_id, tablet_id});

txn->put(key, val);

key = stats_tablet_key({instance_id, table_id, index_id, partition_id, tablet_id});

txn->put(key, val); // val is not necessary

key = job_tablet_key({instance_id, table_id, index_id, partition_id, tablet_id});

txn->put(key, val); // val is not necessary

// mock tablet index

TabletIndexPB tablet_idx_pb;

tablet_idx_pb.set_db_id(db_id);

tablet_idx_pb.set_table_id(table_id);

tablet_idx_pb.set_index_id(index_id);

tablet_idx_pb.set_partition_id(partition_id);

tablet_idx_pb.set_tablet_id(tablet_id);

auto idx_val = tablet_idx_pb.SerializeAsString();

key = meta_tablet_idx_key({instance_id, tablet_id});

txn->put(key, idx_val);

LOG(INFO) << "tablet_idx_pb:" << tablet_idx_pb.DebugString() << " key=" << hex(key);

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

return 0;

}

static int create_recycle_partiton(TxnKv* txn_kv, int64_t table_id, int64_t partition_id,

const std::vector<int64_t>& index_ids) {

std::string key;

std::string val;

RecyclePartKeyInfo key_info {instance_id, partition_id};

recycle_partition_key(key_info, &key);

RecyclePartitionPB partition_pb;

partition_pb.set_db_id(db_id);

partition_pb.set_table_id(table_id);

for (auto index_id : index_ids) {

partition_pb.add_index_id(index_id);

}

partition_pb.set_creation_time(current_time);

partition_pb.set_state(RecyclePartitionPB::DROPPED);

partition_pb.SerializeToString(&val);

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(key, val);

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

return 0;

}

static int create_partition_version_kv(TxnKv* txn_kv, int64_t table_id, int64_t partition_id) {

auto key = partition_version_key({instance_id, db_id, table_id, partition_id});

VersionPB version;

version.set_version(1);

auto val = version.SerializeAsString();

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(key, val);

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

return 0;

}

static int create_table_version_kv(TxnKv* txn_kv, int64_t table_id) {

auto key = table_version_key({instance_id, db_id, table_id});

std::string val(sizeof(int64_t), 0);

*reinterpret_cast<int64_t*>(val.data()) = (int64_t)1;

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(key, val);

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

return 0;

}

static int create_txn_label_kv(TxnKv* txn_kv, std::string label, int64_t db_id) {

std::string txn_label_key_;

std::string txn_label_val;

auto keyinfo = TxnLabelKeyInfo({instance_id, db_id, label});

txn_label_key(keyinfo, &txn_label_key_);

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(txn_label_key_, label);

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

return 0;

}

static int create_recycle_index(TxnKv* txn_kv, int64_t table_id, int64_t index_id) {

std::string key;

std::string val;

RecycleIndexKeyInfo key_info {instance_id, index_id};

recycle_index_key(key_info, &key);

RecycleIndexPB index_pb;

index_pb.set_table_id(table_id);

index_pb.set_creation_time(current_time);

index_pb.set_state(RecycleIndexPB::DROPPED);

index_pb.SerializeToString(&val);

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(key, val);

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

return 0;

}

static int get_txn_info(std::shared_ptr<TxnKv> txn_kv, std::string instance_id, int64_t db_id,

int64_t txn_id, TxnInfoPB& txn_info_pb) {

std::string txn_inf_key;

std::string txn_inf_val;

TxnInfoKeyInfo txn_inf_key_info {instance_id, db_id, txn_id};

LOG(INFO) << instance_id << "|" << db_id << "|" << txn_id;

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn_info_key(txn_inf_key_info, &txn_inf_key);

LOG(INFO) << "txn_inf_key:" << hex(txn_inf_key);

TxnErrorCode err = txn->get(txn_inf_key, &txn_inf_val);

if (err != TxnErrorCode::TXN_OK) {

LOG(WARNING) << "txn->get failed, err=" << err;

return -2;

}

if (!txn_info_pb.ParseFromString(txn_inf_val)) {

LOG(WARNING) << "ParseFromString failed";

return -3;

}

LOG(INFO) << "txn_info_pb" << txn_info_pb.DebugString();

err = txn->commit();

if (err != TxnErrorCode::TXN_OK) {

LOG(WARNING) << "txn->commit failed, err=" << err;

return -4;

}

return 0;

}

static int check_recycle_txn_keys(std::shared_ptr<TxnKv> txn_kv, std::string instance_id,

int64_t db_id, int64_t txn_id, const std::string& label) {

std::string txn_inf_key;

std::string txn_inf_val;

TxnInfoKeyInfo txn_inf_key_info {instance_id, db_id, txn_id};

LOG(INFO) << instance_id << "|" << db_id << "|" << txn_id;

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn_info_key(txn_inf_key_info, &txn_inf_key);

TxnErrorCode err = txn->get(txn_inf_key, &txn_inf_val);

if (err != TxnErrorCode::TXN_KEY_NOT_FOUND) {

return -2;

}

std::string label_key, label_val;

txn_label_key({instance_id, db_id, label}, &label_key);

err = txn->get(label_key, &label_val);

if (err != TxnErrorCode::TXN_KEY_NOT_FOUND) {

return -3;

}

std::string index_key, index_val;

index_key = txn_index_key({instance_id, txn_id});

err = txn->get(index_key, &index_val);

if (err != TxnErrorCode::TXN_KEY_NOT_FOUND) {

return -4;

}

std::string running_key;

std::string running_value;

TxnRunningKeyInfo running_key_info {instance_id, db_id, txn_id};

txn_running_key(running_key_info, &running_key);

err = txn->get(running_key, &running_value);

if (err != TxnErrorCode::TXN_KEY_NOT_FOUND) {

return -5;

}

std::string rec_txn_key;

std::string rec_txn_val;

RecycleTxnKeyInfo recycle_txn_key_info {instance_id, db_id, txn_id};

recycle_txn_key(recycle_txn_key_info, &rec_txn_key);

err = txn->get(rec_txn_key, &rec_txn_val);

if (err != TxnErrorCode::TXN_KEY_NOT_FOUND) {

return -6;

}

return 0;

}

static int create_instance(const std::string& internal_stage_id,

const std::string& external_stage_id, InstanceInfoPB& instance_info) {

// create internal stage

{

std::string s3_prefix = "internal_prefix";

std::string stage_prefix = fmt::format("{}/stage/root/{}/", s3_prefix, internal_stage_id);

ObjectStoreInfoPB object_info;

object_info.set_id(internal_stage_id); // test use accessor_map_ in recycle

StagePB internal_stage;

internal_stage.set_type(StagePB::INTERNAL);

internal_stage.set_stage_id(internal_stage_id);

ObjectStoreInfoPB internal_object_info;

internal_object_info.set_prefix(stage_prefix);

internal_object_info.set_id("0");

internal_stage.mutable_obj_info()->CopyFrom(internal_object_info);

instance_info.add_obj_info()->CopyFrom(object_info);

instance_info.add_stages()->CopyFrom(internal_stage);

}

// create external stage

{

ObjectStoreInfoPB object_info;

object_info.set_id(external_stage_id);

StagePB external_stage;

external_stage.set_type(StagePB::EXTERNAL);

external_stage.set_stage_id(external_stage_id);

external_stage.mutable_obj_info()->CopyFrom(object_info);

instance_info.add_obj_info()->CopyFrom(object_info);

instance_info.add_stages()->CopyFrom(external_stage);

}

instance_info.set_instance_id(instance_id);

return 0;

}

static int create_copy_job(TxnKv* txn_kv, const std::string& stage_id, int64_t table_id,

StagePB::StageType stage_type, CopyJobPB::JobStatus job_status,

std::vector<ObjectFilePB> object_files, int64_t timeout_time,

int64_t start_time = 0, int64_t finish_time = 0) {

std::string key;

std::string val;

CopyJobKeyInfo key_info {instance_id, stage_id, table_id, "copy_id", 0};

copy_job_key(key_info, &key);

CopyJobPB copy_job;

copy_job.set_stage_type(stage_type);

copy_job.set_job_status(job_status);

copy_job.set_timeout_time_ms(timeout_time);

if (start_time != 0) {

copy_job.set_start_time_ms(start_time);

}

if (finish_time != 0) {

copy_job.set_finish_time_ms(finish_time);

}

for (const auto& file : object_files) {

copy_job.add_object_files()->CopyFrom(file);

}

copy_job.SerializeToString(&val);

std::vector<std::string> file_keys;

std::string file_val;

CopyFilePB copy_file;

copy_file.set_copy_id("copy_id");

copy_file.set_group_id(0);

file_val = copy_file.SerializeAsString();

// create job files

for (const auto& file : object_files) {

CopyFileKeyInfo file_info {instance_id, stage_id, table_id, file.relative_path(),

file.etag()};

std::string file_key;

copy_file_key(file_info, &file_key);

file_keys.push_back(file_key);

}

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

txn->put(key, val);

for (const auto& file_key : file_keys) {

txn->put(file_key, file_val);

}

if (txn->commit() != TxnErrorCode::TXN_OK) {

return -1;

}

return 0;

}

static int copy_job_exists(TxnKv* txn_kv, const std::string& stage_id, int64_t table_id,

bool* exist) {

std::string key;

std::string val;

CopyJobKeyInfo key_info {instance_id, stage_id, table_id, "copy_id", 0};

copy_job_key(key_info, &key);

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

TxnErrorCode err = txn->get(key, &val);

if (err != TxnErrorCode::TXN_OK && err != TxnErrorCode::TXN_KEY_NOT_FOUND) {

return -1;

}

*exist = err == TxnErrorCode::TXN_OK;

return 0;

}

static int create_object_files(StorageVaultAccessor* accessor,

std::vector<ObjectFilePB>* object_files) {

for (auto& file : *object_files) {

auto key = file.relative_path();

if (accessor->put_file(key, "") != 0) {

return -1;

}

file.set_etag("");

}

return 0;

}

static int get_copy_file_num(TxnKv* txn_kv, const std::string& stage_id, int64_t table_id,

int* file_num) {

*file_num = 0;

std::string key0;

std::string key1;

CopyFileKeyInfo key_info0 {instance_id, stage_id, table_id, "", ""};

CopyFileKeyInfo key_info1 {instance_id, stage_id, table_id + 1, "", ""};

copy_file_key(key_info0, &key0);

copy_file_key(key_info1, &key1);

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return -1;

}

std::unique_ptr<RangeGetIterator> it;

do {

if (txn->get(key0, key1, &it) != TxnErrorCode::TXN_OK) {

return -1;

}

while (it->has_next()) {

it->next();

++(*file_num);

}

key0.push_back('\x00');

} while (it->more());

return 0;

}

TEST(RecyclerTest, recycle_empty) {

auto txn_kv = std::make_shared<MemTxnKv>();

ASSERT_EQ(txn_kv->init(), 0);

InstanceInfoPB instance;

instance.set_instance_id(instance_id);

auto obj_info = instance.add_obj_info();

obj_info->set_id("recycle_empty");

obj_info->set_ak(config::test_s3_ak);

obj_info->set_sk(config::test_s3_sk);

obj_info->set_endpoint(config::test_s3_endpoint);

obj_info->set_region(config::test_s3_region);

obj_info->set_bucket(config::test_s3_bucket);

obj_info->set_prefix("recycle_empty");

InstanceRecycler recycler(txn_kv, instance, thread_group,

std::make_shared<TxnLazyCommitter>(txn_kv));

ASSERT_EQ(recycler.init(), 0);

ASSERT_EQ(recycler.recycle_rowsets(), 0);

}

TEST(RecyclerTest, recycle_rowsets) {

config::retention_seconds = 0;

auto txn_kv = std::make_shared<MemTxnKv>();

ASSERT_EQ(txn_kv->init(), 0);

InstanceInfoPB instance;

instance.set_instance_id(instance_id);

auto obj_info = instance.add_obj_info();

obj_info->set_id("recycle_rowsets");

obj_info->set_ak(config::test_s3_ak);

obj_info->set_sk(config::test_s3_sk);

obj_info->set_endpoint(config::test_s3_endpoint);

obj_info->set_region(config::test_s3_region);

obj_info->set_bucket(config::test_s3_bucket);

obj_info->set_prefix("recycle_rowsets");

config::instance_recycler_worker_pool_size = 1;

int insert_no_inverted_index = 0;

int insert_inverted_index = 0;

auto sp = SyncPoint::get_instance();

std::unique_ptr<int, std::function<void(int*)>> defer(

(int*)0x01, [](int*) { SyncPoint::get_instance()->clear_all_call_backs(); });

sp->set_call_back("InvertedIndexIdCache::insert1", [&](auto&&) { ++insert_no_inverted_index; });

sp->set_call_back("InvertedIndexIdCache::insert2", [&](auto&&) { ++insert_inverted_index; });

sp->enable_processing();

InstanceRecycler recycler(txn_kv, instance, thread_group,

std::make_shared<TxnLazyCommitter>(txn_kv));

ASSERT_EQ(recycler.init(), 0);

std::vector<doris::TabletSchemaCloudPB> schemas;

for (int i = 0; i < 5; ++i) {

auto& schema = schemas.emplace_back();

schema.set_schema_version(i);

schema.set_inverted_index_storage_format(InvertedIndexStorageFormatPB::V1);

for (int j = 0; j < i; ++j) {

auto index = schema.add_index();

index->set_index_id(j);

index->set_index_type(IndexType::INVERTED);

}

}

auto accessor = recycler.accessor_map_.begin()->second;

constexpr int index_id = 10001, tablet_id = 10002;

for (int i = 0; i < 1000; ++i) {

auto rowset = create_rowset("recycle_rowsets", tablet_id, index_id, 5, schemas[i % 5]);

create_recycle_rowset(

txn_kv.get(), accessor.get(), rowset,

static_cast<RecycleRowsetPB::Type>(i % (RecycleRowsetPB::Type_MAX + 1)), i & 1);

}

for (int i = 0; i < 1000; ++i) {

auto rowset = create_rowset("recycle_rowsets", tablet_id, index_id, 5, schemas[i % 5]);

create_recycle_rowset(txn_kv.get(), accessor.get(), rowset, RecycleRowsetPB::COMPACT, true);

}

ASSERT_EQ(recycler.recycle_rowsets(), 0);

// check rowset does not exist on obj store

std::unique_ptr<ListIterator> list_iter;

ASSERT_EQ(0, accessor->list_directory(tablet_path_prefix(tablet_id), &list_iter));

EXPECT_FALSE(list_iter->has_next());

// check all recycle rowset kv have been deleted

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

std::unique_ptr<RangeGetIterator> it;

auto begin_key = recycle_key_prefix(instance_id);

auto end_key = recycle_key_prefix(instance_id + '\xff');

ASSERT_EQ(txn->get(begin_key, end_key, &it), TxnErrorCode::TXN_OK);

EXPECT_EQ(it->size(), 0);

// Check InvertedIndexIdCache

EXPECT_EQ(insert_inverted_index, 4);

EXPECT_EQ(insert_no_inverted_index, 1);

}

TEST(RecyclerTest, bench_recycle_rowsets) {

config::retention_seconds = 0;

auto txn_kv = std::make_shared<MemTxnKv>();

ASSERT_EQ(txn_kv->init(), 0);

InstanceInfoPB instance;

instance.set_instance_id(instance_id);

auto obj_info = instance.add_obj_info();

obj_info->set_id("recycle_rowsets");

obj_info->set_ak(config::test_s3_ak);

obj_info->set_sk(config::test_s3_sk);

obj_info->set_endpoint(config::test_s3_endpoint);

obj_info->set_region(config::test_s3_region);

obj_info->set_bucket(config::test_s3_bucket);

obj_info->set_prefix("recycle_rowsets");

config::instance_recycler_worker_pool_size = 10;

config::recycle_task_threshold_seconds = 0;

InstanceRecycler recycler(txn_kv, instance, thread_group,

std::make_shared<TxnLazyCommitter>(txn_kv));

ASSERT_EQ(recycler.init(), 0);

auto sp = SyncPoint::get_instance();

std::unique_ptr<int, std::function<void(int*)>> defer(

(int*)0x01, [](int*) { SyncPoint::get_instance()->clear_all_call_backs(); });

sp->set_call_back("memkv::Transaction::get", [](auto&& args) {

auto* limit = try_any_cast<int*>(args[0]);

*limit = 100;

std::this_thread::sleep_for(std::chrono::milliseconds(5));

});

sp->set_call_back("MockAccessor::delete_files", [&](auto&& args) {

std::this_thread::sleep_for(std::chrono::seconds(1));

bool found = recycler.check_recycle_tasks();

ASSERT_EQ(found, true);

});

sp->set_call_back("MockAccessor::delete_prefix",

[&](auto&&) { std::this_thread::sleep_for(std::chrono::milliseconds(20)); });

sp->enable_processing();

std::vector<doris::TabletSchemaCloudPB> schemas;

for (int i = 0; i < 5; ++i) {

auto& schema = schemas.emplace_back();

schema.set_schema_version(i);

schema.set_inverted_index_storage_format(InvertedIndexStorageFormatPB::V1);

for (int j = 0; j < i; ++j) {

auto index = schema.add_index();

index->set_index_id(j);

index->set_index_type(IndexType::INVERTED);

}

}

auto accessor = recycler.accessor_map_.begin()->second;

constexpr int index_id = 10001, tablet_id = 10002;

for (int i = 0; i < 2000; ++i) {

auto rowset = create_rowset("recycle_rowsets", tablet_id, index_id, 5, schemas[i % 5]);

create_recycle_rowset(txn_kv.get(), accessor.get(), rowset,

i % 10 < 2 ? RecycleRowsetPB::PREPARE : RecycleRowsetPB::COMPACT,

i & 1);

}

ASSERT_EQ(recycler.recycle_rowsets(), 0);

ASSERT_EQ(recycler.check_recycle_tasks(), false);

// check rowset does not exist on obj store

std::unique_ptr<ListIterator> list_iter;

ASSERT_EQ(0, accessor->list_directory(tablet_path_prefix(tablet_id), &list_iter));

ASSERT_FALSE(list_iter->has_next());

// check all recycle rowset kv have been deleted

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

std::unique_ptr<RangeGetIterator> it;

auto begin_key = recycle_key_prefix(instance_id);

auto end_key = recycle_key_prefix(instance_id + '\xff');

ASSERT_EQ(txn->get(begin_key, end_key, &it), TxnErrorCode::TXN_OK);

ASSERT_EQ(it->size(), 0);

}

TEST(RecyclerTest, recycle_tmp_rowsets) {

config::retention_seconds = 0;

auto txn_kv = std::make_shared<MemTxnKv>();

ASSERT_EQ(txn_kv->init(), 0);

InstanceInfoPB instance;

instance.set_instance_id(instance_id);

auto obj_info = instance.add_obj_info();

obj_info->set_id("recycle_tmp_rowsets");

obj_info->set_ak(config::test_s3_ak);

obj_info->set_sk(config::test_s3_sk);

obj_info->set_endpoint(config::test_s3_endpoint);

obj_info->set_region(config::test_s3_region);

obj_info->set_bucket(config::test_s3_bucket);

obj_info->set_prefix("recycle_tmp_rowsets");

int insert_no_inverted_index = 0;

int insert_inverted_index = 0;

auto sp = SyncPoint::get_instance();

std::unique_ptr<int, std::function<void(int*)>> defer(

(int*)0x01, [](int*) { SyncPoint::get_instance()->clear_all_call_backs(); });

sp->set_call_back("InvertedIndexIdCache::insert1", [&](auto&&) { ++insert_no_inverted_index; });

sp->set_call_back("InvertedIndexIdCache::insert2", [&](auto&&) { ++insert_inverted_index; });

sp->enable_processing();

InstanceRecycler recycler(txn_kv, instance, thread_group,

std::make_shared<TxnLazyCommitter>(txn_kv));

ASSERT_EQ(recycler.init(), 0);

std::vector<doris::TabletSchemaCloudPB> schemas;

for (int i = 0; i < 5; ++i) {

auto& schema = schemas.emplace_back();

schema.set_inverted_index_storage_format(InvertedIndexStorageFormatPB::V1);

schema.set_schema_version(i);

for (int j = 0; j < i; ++j) {

auto index = schema.add_index();

index->set_index_id(j);

index->set_index_type(IndexType::INVERTED);

}

}

auto accessor = recycler.accessor_map_.begin()->second;

int64_t txn_id_base = 114115;

int64_t tablet_id_base = 10015;

int64_t index_id_base = 1000;

for (int i = 0; i < 100; ++i) {

int64_t txn_id = txn_id_base + i;

for (int j = 0; j < 20; ++j) {

auto rowset = create_rowset("recycle_tmp_rowsets", tablet_id_base + j,

index_id_base + j % 4, 5, schemas[i % 5], txn_id);

create_tmp_rowset(txn_kv.get(), accessor.get(), rowset, i & 1);

}

}

ASSERT_EQ(recycler.recycle_tmp_rowsets(), 0);

// check rowset does not exist on obj store

std::unique_ptr<ListIterator> list_iter;

ASSERT_EQ(0, accessor->list_directory("data/", &list_iter));

ASSERT_FALSE(list_iter->has_next());

// check all tmp rowset kv have been deleted

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

std::unique_ptr<RangeGetIterator> it;

auto begin_key = meta_rowset_tmp_key({instance_id, 0, 0});

auto end_key = meta_rowset_tmp_key({instance_id, INT64_MAX, 0});

ASSERT_EQ(txn->get(begin_key, end_key, &it), TxnErrorCode::TXN_OK);

ASSERT_EQ(it->size(), 0);

// Check InvertedIndexIdCache

EXPECT_EQ(insert_inverted_index, 16);

EXPECT_EQ(insert_no_inverted_index, 4);

}

TEST(RecyclerTest, recycle_tmp_rowsets_partial_update) {

config::retention_seconds = 0;

auto txn_kv = std::make_shared<MemTxnKv>();

ASSERT_EQ(txn_kv->init(), 0);

InstanceInfoPB instance;

instance.set_instance_id(instance_id);

auto obj_info = instance.add_obj_info();

obj_info->set_id("recycle_tmp_rowsets_partial_update");

obj_info->set_ak(config::test_s3_ak);

obj_info->set_sk(config::test_s3_sk);

obj_info->set_endpoint(config::test_s3_endpoint);

obj_info->set_region(config::test_s3_region);

obj_info->set_bucket(config::test_s3_bucket);

obj_info->set_prefix("recycle_tmp_rowsets_partial_update");

InstanceRecycler recycler(txn_kv, instance, thread_group,

std::make_shared<TxnLazyCommitter>(txn_kv));

ASSERT_EQ(recycler.init(), 0);

doris::TabletSchemaCloudPB schema;

auto accessor = recycler.accessor_map_.begin()->second;

int64_t tablet_id = 10015;

int64_t index_id = 1000;

int64_t txn_id_base = 293039;

for (int j = 0; j < 20; ++j) {

int64_t txn_id = txn_id_base + j;

int segment_num = 5;

if (j < 15) {

auto rowset = create_rowset("recycle_tmp_rowsets_partial_update", tablet_id, index_id,

segment_num, schema, RowsetStatePB::VISIBLE, txn_id);

create_tmp_rowset(txn_kv.get(), accessor.get(), rowset, false);

} else {

auto rowset =

create_rowset("recycle_tmp_rowsets_partial_update", tablet_id, tablet_id,

segment_num, schema, RowsetStatePB::BEGIN_PARTIAL_UPDATE, txn_id);

create_tmp_rowset(txn_kv.get(), accessor.get(), rowset, false);

// partial update may write new segment to an existing tmp rowsets

// we simulate that partial update load fails after it writes a segment

// and before it updates the segments num in tmp rowset meta

int extra_segment_id = segment_num;

auto path = segment_path(rowset.tablet_id(), rowset.rowset_id_v2(), extra_segment_id);

accessor->put_file(path, path);

}

}

ASSERT_EQ(recycler.recycle_tmp_rowsets(), 0);