// 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 "meta-service/meta_server.h"

#include <brpc/server.h>

#include <butil/endpoint.h>

#include <bvar/window.h>

#include <gen_cpp/cloud.pb.h>

#include <gtest/gtest.h>

#include <atomic>

#include <condition_variable>

#include <cstdint>

#include <cstdlib>

#include <fstream>

#include <memory>

#include <random>

#include <thread>

#include "common/config.h"

#include "common/defer.h"

#include "common/logging.h"

#include "common/stats.h"

#include "cpp/sync_point.h"

#include "meta-service/meta_service.h"

#include "meta-store/keys.h"

#include "meta-store/mem_txn_kv.h"

#include "meta-store/txn_kv.h"

#include "meta-store/txn_kv_error.h"

#include "mock_resource_manager.h"

#include "rate-limiter/rate_limiter.h"

#include "resource-manager/resource_manager.h"

using namespace doris;

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

const std::string conf_file = "doris_cloud.conf";

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

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

return -1;

}

// Clear the secondary package name for testing, to avoid brpc::Server::~Server() SEGV

config::secondary_package_name = "";

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

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

return -1;

}

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

return RUN_ALL_TESTS();

}

namespace doris::cloud {

void notify_refresh_instance(std::shared_ptr<TxnKv> txn_kv, const std::string& instance_id,

KVStats* stats, bool include_self);

} // namespace doris::cloud

TEST(MetaServerTest, FQDNRefreshInstance) {

class MockMetaService : public cloud::MetaServiceImpl {

public:

MockMetaService(std::shared_ptr<TxnKv> txn_kv,

std::shared_ptr<ResourceManager> resource_mgr,

std::shared_ptr<RateLimiter> rate_controller,

std::shared_ptr<SnapshotManager> snapshot_mgr)

: MetaServiceImpl(txn_kv, resource_mgr, rate_controller, snapshot_mgr) {}

~MockMetaService() override = default;

void alter_instance(google::protobuf::RpcController* controller,

const ::doris::cloud::AlterInstanceRequest* request,

::doris::cloud::AlterInstanceResponse* response,

::google::protobuf::Closure* done) override {

(void)controller;

response->mutable_status()->set_code(cloud::MetaServiceCode::OK);

LOG(INFO) << "alter instance " << request->DebugString();

if (request->op() == cloud::AlterInstanceRequest::REFRESH) {

std::unique_lock<std::mutex> lock(mu_);

LOG(INFO) << "refresh instance " << request->instance_id();

refreshed_instances_.insert(request->instance_id());

}

done->Run();

}

bool is_instance_refreshed(std::string instance_id) {

std::unique_lock<std::mutex> lock(mu_);

return refreshed_instances_.count(instance_id) > 0;

}

std::mutex mu_;

std::unordered_set<std::string> refreshed_instances_;

};

std::shared_ptr<cloud::TxnKv> txn_kv = std::make_shared<cloud::MemTxnKv>();

auto resource_mgr = std::make_shared<MockResourceManager>(txn_kv);

auto rate_limiter = std::make_shared<cloud::RateLimiter>();

auto snapshot_mgr = std::make_shared<cloud::SnapshotManager>(txn_kv);

auto mock_service =

std::make_unique<MockMetaService>(txn_kv, resource_mgr, rate_limiter, snapshot_mgr);

MockMetaService* mock_service_ptr = mock_service.get();

MetaServiceProxy meta_service(std::move(mock_service));

brpc::ServerOptions options;

options.num_threads = 1;

brpc::Server server;

ASSERT_EQ(server.AddService(&meta_service, brpc::ServiceOwnership::SERVER_DOESNT_OWN_SERVICE),

0);

ASSERT_EQ(server.Start(0, &options), 0);

auto addr = server.listen_address();

config::hostname = butil::my_hostname();

config::brpc_listen_port = addr.port;

config::meta_server_register_interval_ms = 1;

// Register meta service.

cloud::MetaServerRegister server_register(txn_kv);

ASSERT_EQ(server_register.start(), 0);

while (true) {

std::unique_ptr<cloud::Transaction> txn;

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

auto system_key = cloud::system_meta_service_registry_key();

std::string value;

if (txn->get(system_key, &value) == TxnErrorCode::TXN_OK) {

break;

}

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

}

server_register.stop();

// Refresh instance with FQDN endpoint.

config::hostname = "";

notify_refresh_instance(txn_kv, "fqdn_instance_id", nullptr, false);

bool refreshed = false;

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

if (mock_service_ptr->is_instance_refreshed("fqdn_instance_id")) {

refreshed = true;

break;

}

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

}

EXPECT_TRUE(refreshed);

server.Stop(1);

server.Join();

}

TEST(MetaServerTest, InstanceWatcherStartAndStop) {

class MockResourceManagerWithTracking : public MockResourceManager {

public:

MockResourceManagerWithTracking(std::shared_ptr<TxnKv> txn_kv)

: MockResourceManager(txn_kv) {}

std::pair<MetaServiceCode, std::string> refresh_instance(

const std::string& instance_id) override {

std::unique_lock<std::mutex> lock(mu_);

refreshed_instances_.insert(instance_id);

refresh_count_++;

cv_.notify_all();

return std::make_pair(MetaServiceCode::OK, "");

}

bool wait_for_refresh(const std::string& instance_id, int timeout_ms = 5000) {

std::unique_lock<std::mutex> lock(mu_);

return cv_.wait_for(lock, std::chrono::milliseconds(timeout_ms), [this, &instance_id] {

return refreshed_instances_.count(instance_id) > 0;

});

}

int get_refresh_count() {

std::unique_lock<std::mutex> lock(mu_);

return refresh_count_;

}

std::mutex mu_;

std::condition_variable cv_;

std::unordered_set<std::string> refreshed_instances_;

int refresh_count_ = 0;

};

std::shared_ptr<cloud::TxnKv> txn_kv = std::make_shared<cloud::MemTxnKv>();

auto resource_mgr = std::make_shared<MockResourceManagerWithTracking>(txn_kv);

// Test basic start and stop

{

auto watcher = std::make_unique<cloud::MetaServerInstanceWatcher>(txn_kv, resource_mgr);

ASSERT_EQ(watcher->start(), 0);

// Try to start again should fail

ASSERT_EQ(watcher->start(), -2);

watcher->stop();

}

// Test nullptr txn_kv

{

auto watcher = std::make_unique<cloud::MetaServerInstanceWatcher>(nullptr, resource_mgr);

ASSERT_EQ(watcher->start(), -1);

}

}

TEST(MetaServerTest, InstanceWatcherDetectInstanceUpdate) {

class MockResourceManagerWithTracking : public MockResourceManager {

public:

MockResourceManagerWithTracking(std::shared_ptr<TxnKv> txn_kv)

: MockResourceManager(txn_kv) {}

std::pair<MetaServiceCode, std::string> refresh_instance(

const std::string& instance_id) override {

std::unique_lock<std::mutex> lock(mu_);

refreshed_instances_.insert(instance_id);

refresh_count_++;

cv_.notify_all();

LOG(INFO) << "refresh_instance called for instance_id=" << instance_id

<< " total_count=" << refresh_count_;

return std::make_pair(MetaServiceCode::OK, "");

}

bool wait_for_refresh(const std::string& instance_id, int timeout_ms = 5000) {

std::unique_lock<std::mutex> lock(mu_);

return cv_.wait_for(lock, std::chrono::milliseconds(timeout_ms), [this, &instance_id] {

return refreshed_instances_.count(instance_id) > 0;

});

}

bool wait_for_refresh_count(int count, int timeout_ms = 5000) {

std::unique_lock<std::mutex> lock(mu_);

return cv_.wait_for(lock, std::chrono::milliseconds(timeout_ms),

[this, count] { return refresh_count_ >= count; });

}

int get_refresh_count() {

std::unique_lock<std::mutex> lock(mu_);

return refresh_count_;

}

std::mutex mu_;

std::condition_variable cv_;

std::unordered_set<std::string> refreshed_instances_;

int refresh_count_ = 0;

};

std::shared_ptr<cloud::TxnKv> txn_kv = std::make_shared<cloud::MemTxnKv>();

auto resource_mgr = std::make_shared<MockResourceManagerWithTracking>(txn_kv);

auto watcher = std::make_unique<cloud::MetaServerInstanceWatcher>(txn_kv, resource_mgr);

ASSERT_EQ(watcher->start(), 0);

// Create instance in kv store

std::string instance_id = "test_instance_123";

cloud::InstanceKeyInfo key_info {instance_id};

std::string instance_key = cloud::instance_key(key_info);

cloud::InstanceInfoPB instance_pb;

instance_pb.set_instance_id(instance_id);

instance_pb.set_name("test_instance");

std::string instance_val = instance_pb.SerializeAsString();

{

std::unique_ptr<cloud::Transaction> txn;

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

txn->put(instance_key, instance_val);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

// Trigger instance update

{

std::unique_ptr<cloud::Transaction> txn;

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

txn->atomic_add(cloud::system_meta_service_instance_update_key(), 1);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

// Wait for watcher to detect and refresh

ASSERT_TRUE(resource_mgr->wait_for_refresh(instance_id, 5000))

<< "Failed to detect instance update";

// Update instance again

instance_pb.set_name("test_instance_updated");

instance_val = instance_pb.SerializeAsString();

{

std::unique_ptr<cloud::Transaction> txn;

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

txn->put(instance_key, instance_val);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

// Trigger instance update again

{

std::unique_ptr<cloud::Transaction> txn;

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

txn->atomic_add(cloud::system_meta_service_instance_update_key(), 1);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

// Should detect the second update

ASSERT_TRUE(resource_mgr->wait_for_refresh_count(2, 5000))

<< "Failed to detect second instance update";

watcher->stop();

}

TEST(MetaServerTest, InstanceWatcherMultipleInstances) {

class MockResourceManagerWithTracking : public MockResourceManager {

public:

MockResourceManagerWithTracking(std::shared_ptr<TxnKv> txn_kv)

: MockResourceManager(txn_kv) {}

std::pair<MetaServiceCode, std::string> refresh_instance(

const std::string& instance_id) override {

std::unique_lock<std::mutex> lock(mu_);

refreshed_instances_.insert(instance_id);

refresh_count_[instance_id]++;

cv_.notify_all();

LOG(INFO) << "refresh_instance called for instance_id=" << instance_id;

return std::make_pair(MetaServiceCode::OK, "");

}

bool wait_for_instances(const std::vector<std::string>& instance_ids,

int timeout_ms = 5000) {

std::unique_lock<std::mutex> lock(mu_);

return cv_.wait_for(lock, std::chrono::milliseconds(timeout_ms), [this, &instance_ids] {

for (const auto& id : instance_ids) {

if (refreshed_instances_.count(id) == 0) {

return false;

}

}

return true;

});

}

std::mutex mu_;

std::condition_variable cv_;

std::unordered_set<std::string> refreshed_instances_;

std::unordered_map<std::string, int> refresh_count_;

};

std::shared_ptr<cloud::TxnKv> txn_kv = std::make_shared<cloud::MemTxnKv>();

auto resource_mgr = std::make_shared<MockResourceManagerWithTracking>(txn_kv);

auto watcher = std::make_unique<cloud::MetaServerInstanceWatcher>(txn_kv, resource_mgr);

ASSERT_EQ(watcher->start(), 0);

// Create multiple instances

std::vector<std::string> instance_ids = {"instance_1", "instance_2", "instance_3"};

for (const auto& instance_id : instance_ids) {

cloud::InstanceKeyInfo key_info {instance_id};

std::string instance_key = cloud::instance_key(key_info);

cloud::InstanceInfoPB instance_pb;

instance_pb.set_instance_id(instance_id);

instance_pb.set_name("instance_" + instance_id);

std::unique_ptr<cloud::Transaction> txn;

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

txn->put(instance_key, instance_pb.SerializeAsString());

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

// Trigger instance update

{

std::unique_ptr<cloud::Transaction> txn;

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

txn->atomic_add(cloud::system_meta_service_instance_update_key(), 1);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

// Wait for all instances to be refreshed

ASSERT_TRUE(resource_mgr->wait_for_instances(instance_ids, 5000))

<< "Failed to detect all instance updates";

watcher->stop();

}

TEST(MetaServerTest, StartAndStop) {

std::shared_ptr<cloud::TxnKv> txn_kv = std::make_shared<cloud::MemTxnKv>();

auto server = std::make_unique<MetaServer>(txn_kv);

auto resource_mgr = std::make_shared<MockResourceManager>(txn_kv);

auto rate_limiter = std::make_shared<cloud::RateLimiter>();

brpc::ServerOptions options;

options.num_threads = 1;

brpc::Server brpc_server;

auto sp = SyncPoint::get_instance();

std::array<std::string, 2> sps {"MetaServer::start:1", "MetaServer::start:2"};

// use structured binding for point alias (avoid multi lines of declaration)

auto [meta_server_start_1, meta_server_start_2] = sps;

sp->enable_processing();

DORIS_CLOUD_DEFER {

for (auto& i : sps) {

sp->clear_call_back(i);

} // redundant

sp->disable_processing();

};

auto foo = [](auto&& args) {

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

*ret = 1;

};

// failed to init resource mgr

sp->set_call_back(meta_server_start_1, foo);

ASSERT_EQ(server->start(&brpc_server), 1);

sp->clear_call_back(meta_server_start_1);

// failed to start registry

sp->set_call_back(meta_server_start_2, foo);

ASSERT_EQ(server->start(&brpc_server), -1);

sp->clear_call_back(meta_server_start_2);

ASSERT_EQ(server->start(&brpc_server), 0);

ASSERT_EQ(brpc_server.Start(0, &options), 0);

auto addr = brpc_server.listen_address();

config::hostname = butil::my_hostname();

config::brpc_listen_port = addr.port;

config::meta_server_register_interval_ms = 1;

while (true) {

std::unique_ptr<cloud::Transaction> txn;

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

auto system_key = cloud::system_meta_service_registry_key();

std::string value;

if (txn->get(system_key, &value) == TxnErrorCode::TXN_OK) {

break;

}

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

}

server->stop();

brpc_server.Stop(1);

brpc_server.Join();

}