3FS/tests/lib/UnitTestFabric.h
2025-02-27 21:53:53 +08:00

304 lines
12 KiB
C++

#pragma once
#include <map>
#include <vector>
#include "client/mgmtd/MgmtdClientForClient.h"
#include "client/mgmtd/MgmtdClientForServer.h"
#include "client/storage/StorageClient.h"
#include "fbs/mgmtd/MgmtdTypes.h"
#include "fbs/mgmtd/RoutingInfo.h"
#include "storage/service/StorageServer.h"
#include "tests/GtestHelpers.h"
#include "tests/client/ClientWithConfig.h"
#include "tests/client/ServerWithConfig.h"
#include "tests/mgmtd/MgmtdTestHelper.h"
namespace hf3fs::test {
class FakeMgmtdClient : public hf3fs::client::IMgmtdClientForClient, public hf3fs::client::IMgmtdClientForServer {
public:
FakeMgmtdClient(std::shared_ptr<hf3fs::flat::RoutingInfo> routingInfo) { setRoutingInfo(routingInfo); }
std::shared_ptr<hf3fs::client::RoutingInfo> getRoutingInfo() override { return routingInfo_.load(); }
void setRoutingInfo(std::shared_ptr<hf3fs::flat::RoutingInfo> routingInfo) {
routingInfo_.store(std::make_shared<hf3fs::client::RoutingInfo>(routingInfo, SteadyClock::now()));
for (auto &listener : listeners_) {
listener(routingInfo_.load());
}
}
CoTryTask<void> refreshRoutingInfo(bool) override { co_return Void{}; }
CoTryTask<mgmtd::HeartbeatRsp> heartbeat() override { co_return makeError(StatusCode::kNotImplemented); }
Result<Void> triggerHeartbeat() override { return makeError(StatusCode::kNotImplemented); }
CoTryTask<std::optional<flat::ConfigInfo>> getConfig(flat::NodeType, flat::ConfigVersion) override {
co_return makeError(StatusCode::kNotImplemented);
}
CoTryTask<RHStringHashMap<flat::ConfigVersion>> getConfigVersions() override {
co_return makeError(StatusCode::kNotImplemented);
}
CoTryTask<std::vector<flat::TagPair>> getUniversalTags(const String &universalId) override {
co_return makeError(StatusCode::kNotImplemented);
}
void setAppInfoForHeartbeat(flat::AppInfo) override {}
bool addRoutingInfoListener(String, RoutingInfoListener listener) override {
listeners_.push_back(listener);
return true;
}
bool removeRoutingInfoListener(std::string_view) override { return true; }
void setConfigListener(ConfigListener) override {}
void updateHeartbeatPayload(HeartbeatPayload) override {}
CoTryTask<void> extendClientSession() override { co_return makeError(StatusCode::kNotImplemented); }
void setClientSessionPayload(ClientSessionPayload) override {}
CoTryTask<mgmtd::ListClientSessionsRsp> listClientSessions() override {
co_return makeError(StatusCode::kNotImplemented);
}
CoTryTask<mgmtd::GetClientSessionRsp> getClientSession(const String &) override {
co_return makeError(StatusCode::kNotImplemented);
}
// The diamond inheritance forbids directly upcasting from FakeMgmtdClient to ICommonMgmtdClient
hf3fs::client::ICommonMgmtdClient *asCommon() {
hf3fs::client::IMgmtdClientForServer *s = this;
return s;
}
private:
folly::atomic_shared_ptr<hf3fs::client::RoutingInfo> routingInfo_;
std::vector<RoutingInfoListener> listeners_;
};
struct SystemSetupConfig : public ConfigBase<SystemSetupConfig> {
CONFIG_ITEM(chunk_size, 128_KB);
CONFIG_ITEM(num_chains, 1u);
CONFIG_ITEM(num_replicas, 1u);
CONFIG_ITEM(num_storage_nodes, 1u);
CONFIG_ITEM(data_paths, (std::vector<hf3fs::Path>{}));
CONFIG_ITEM(client_config, hf3fs::Path());
CONFIG_ITEM(server_config, hf3fs::Path());
CONFIG_ITEM(storage_endpoints, (std::map<std::string, net::Address>{}));
CONFIG_ITEM(service_level, 0u);
CONFIG_ITEM(listen_port, 0u);
CONFIG_ITEM(client_impl_type, storage::client::StorageClient::ImplementationType::RPC);
CONFIG_ITEM(meta_store_type, kv::KVStore::Type::LevelDB);
CONFIG_ITEM(use_fake_mgmtd_client, true);
CONFIG_ITEM(start_storage_server, true);
CONFIG_ITEM(use_temp_path, true);
public:
std::vector<Size> getChunkSizeList() const {
auto chunkSize = chunk_size();
std::vector<Size> chunkSizeList;
for (auto size : std::vector<Size>{chunkSize / 4, chunkSize / 2, chunkSize, 2 * chunkSize, 4 * chunkSize}) {
if (size >= 512 && (size & size - 1) == 0 /*pow of 2*/) {
chunkSizeList.push_back(size);
}
}
return chunkSizeList;
}
public:
SystemSetupConfig() = default;
SystemSetupConfig(uint32_t chunkSize,
uint32_t numChains,
uint32_t numReplicas,
uint32_t numStorageNodes,
std::vector<hf3fs::Path> dataPaths,
hf3fs::Path clientConfig = hf3fs::Path(),
hf3fs::Path serverConfig = hf3fs::Path(),
std::map<storage::NodeId, net::Address> storageEndpoints = {},
uint32_t serviceLevel = 0,
uint32_t listenPort = 0,
storage::client::StorageClient::ImplementationType clientImplType =
storage::client::StorageClient::ImplementationType::RPC,
kv::KVStore::Type metaStoreType = kv::KVStore::Type::LevelDB,
bool useFakeMgmtdClient = true,
bool startStorageServer = true,
bool useTempPath = true) {
set_chunk_size(chunkSize);
set_num_chains(numChains);
set_num_replicas(numReplicas);
set_num_storage_nodes(numStorageNodes);
set_data_paths(dataPaths);
set_client_config(clientConfig);
set_server_config(serverConfig);
std::map<std::string, net::Address> map{};
for (auto &[nodeId, addr] : storageEndpoints) {
map[std::to_string(nodeId)] = addr;
}
set_storage_endpoints(map);
set_service_level(serviceLevel);
set_listen_port(listenPort);
set_client_impl_type(clientImplType);
set_meta_store_type(metaStoreType);
set_use_fake_mgmtd_client(useFakeMgmtdClient);
set_start_storage_server(startStorageServer);
set_use_temp_path(useTempPath);
}
static std::string prettyPrintConfig(const testing::TestParamInfo<SystemSetupConfig> &info) {
return fmt::format("{}replicas_{}nodes_client{}_metastore{}_fakemgmtd{}",
info.param.num_replicas(),
info.param.num_storage_nodes(),
static_cast<int>(info.param.client_impl_type()),
static_cast<int>(info.param.meta_store_type()),
info.param.use_fake_mgmtd_client());
}
};
using NodeTargetMap = std::map<storage::NodeId, std::vector<storage::TargetId>>;
class UnitTestFabric {
public:
static NodeTargetMap buildNodeTargetMap(uint32_t numNodes, uint32_t numTargetsPerNode);
/*
Suppose numChains = 4, numReplicas = 3 and nodeTargets is a 2-d array:
----------------------------------------------
node1 | target1001 | target1002 | target1003
node2 | target2001 | target2002 | target2003
node3 | target3001 | target3002 | target3003
node4 | target4001 | target4002 | target4003
----------------------------------------------
Then this function generates the following routing table:
-----------------------------------------------
chain1 | target1001 | target2001 | target3001
chain2 | target4001 | target1002 | target2002
chain3 | target3002 | target4002 | target1003
chain4 | target2003 | target3003 | target4003
-----------------------------------------------
*/
static std::map<storage::ChainId, storage::client::SlimChainInfo> buildRepliaChainMap(uint32_t numChains,
uint32_t numReplicas,
NodeTargetMap nodeTargets);
static std::shared_ptr<hf3fs::flat::RoutingInfo> createRoutingInfo(
const std::map<storage::ChainId, storage::client::SlimChainInfo> &replicaChains,
const std::map<storage::NodeId, net::Address> &nodeEndpoints);
/* Set up and tear down */
public:
UnitTestFabric(const SystemSetupConfig &setupConfig)
: setupConfig_(setupConfig),
requestExe_(std::max(1U, std::thread::hardware_concurrency() / 2),
std::make_shared<folly::NamedThreadFactory>("UnitTestFabric")) {}
bool setUpStorageSystem();
void tearDownStorageSystem();
std::unique_ptr<storage::StorageServer> createStorageServer(size_t nodeIndex);
std::shared_ptr<hf3fs::flat::RoutingInfo> getRoutingInfo();
bool updateRoutingInfo(std::function<void(hf3fs::flat::RoutingInfo &)> callback);
auto copyRoutingInfo() { return std::make_shared<hf3fs::flat::RoutingInfo>(*rawRoutingInfo_); }
void setTargetOffline(hf3fs::flat::RoutingInfo &routingInfo, uint32_t targetIndex, uint32_t chainIndex = 0);
bool stopAndRemoveStorageServer(uint32_t serverIndex);
bool stopAndRemoveStorageServer(storage::NodeId nodeId);
/* Helper functions */
storage::IOResult writeToChunk(storage::ChainId chainId,
storage::ChunkId chunkId,
std::span<uint8_t> chunkData,
uint32_t offset = 0,
uint32_t length = 0,
const storage::client::WriteOptions &options = storage::client::WriteOptions());
bool writeToChunks(storage::ChainId chainId,
storage::ChunkId chunkBegin,
storage::ChunkId chunkEnd,
std::span<uint8_t> chunkData,
uint32_t offset = 0,
uint32_t length = 0,
const storage::client::WriteOptions &options = storage::client::WriteOptions(),
std::vector<storage::IOResult> *results = nullptr);
storage::IOResult readFromChunk(storage::ChainId chainId,
storage::ChunkId chunkId,
std::span<uint8_t> chunkData,
uint32_t offset = 0,
uint32_t length = 0,
const storage::client::ReadOptions &options = storage::client::ReadOptions());
bool readFromChunks(storage::ChainId chainId,
storage::ChunkId chunkBegin,
storage::ChunkId chunkEnd,
std::vector<std::vector<uint8_t>> &chunkData,
uint32_t offset = 0,
uint32_t length = 0,
const storage::client::ReadOptions &options = storage::client::ReadOptions(),
std::vector<storage::IOResult> *results = nullptr);
protected:
// the id of the only table
static flat::ChainTableId kTableId() {
static const auto v = flat::ChainTableId(1);
return v;
}
// the initial routing table version
static const flat::RoutingInfoVersion kRoutingVersion() {
static const auto v = flat::RoutingInfoVersion(10);
return v;
}
// the initial chain version
static const storage::ChainVer kInitChainVer() {
static const auto v = storage::ChainVer(1);
return v;
}
constexpr static auto kClusterId = "test";
char const *const kTempDataFolderPrefix = "temp_storage_data";
ClientId clientId_ = ClientId::random();
storage::ChainId firstChainId_;
std::vector<storage::ChainId> chainIds_;
NodeTargetMap nodeTargets_;
std::vector<storage::NodeId> storageNodeIds_;
std::map<storage::ChainId, storage::client::SlimChainInfo> replicaChains_;
std::shared_ptr<hf3fs::flat::RoutingInfo> rawRoutingInfo_;
std::vector<std::vector<folly::test::TemporaryDirectory>> tmpDataPaths_;
std::vector<storage::StorageTargets::Config> targetsConfigs_;
std::vector<storage::StorageServer::Config> serverConfigs_;
std::vector<std::unique_ptr<storage::StorageServer>> storageServers_;
MgmtdServerWithConfig mgmtdServer_{kClusterId, flat::NodeId(10000)};
SystemSetupConfig setupConfig_;
storage::client::StorageClient::Config clientConfig_;
net::Client::Config netClientConfig_;
net::Client client_{netClientConfig_};
hf3fs::client::MgmtdClientForClient::Config mgmtdClientConfig_;
std::unique_ptr<hf3fs::client::ICommonMgmtdClient> mgmtdForClient_;
std::shared_ptr<storage::client::StorageClient> storageClient_;
folly::CPUThreadPoolExecutor requestExe_;
};
} // namespace hf3fs::test