Initial commit

tests/mgmtd/CMakeLists.txt

@@ -0,0 +1,3 @@
+add_library(mgmtd-test-helper STATIC MgmtdTestHelper.cpp)
+target_link_libraries(mgmtd-test-helper mgmtd)
+target_add_test(test_mgmtd mgmtd-test-helper fdb)

tests/mgmtd/MgmtdTestHelper.cpp

@@ -0,0 +1,163 @@
+#include "MgmtdTestHelper.h"
+
+namespace hf3fs::mgmtd::testing {
+CoTryTask<void> MgmtdTestHelper::setRoutingInfo(std::function<void(flat::RoutingInfo &)> callback) {
+  co_await extendLease();
+
+  auto handler = [&]() -> CoTryTask<void> {
+    auto writerLock = co_await op_.state_.coScopedLock<"test">();
+    std::optional<flat::RoutingInfo> tmpRoutingInfo;
+    {
+      auto dataPtr = co_await op_.state_.data_.coSharedLock();
+
+      tmpRoutingInfo = dataPtr->getRoutingInfo(flat::RoutingInfoVersion(0), op_.state_.config_);
+      XLOGF_IF(FATAL, !tmpRoutingInfo, "Empty routingInfo!");
+      CO_RETURN_ON_ERROR(checkDataIntegrity(*tmpRoutingInfo));
+
+      auto rv = tmpRoutingInfo->routingInfoVersion;
+      callback(*tmpRoutingInfo);
+
+      CO_RETURN_ON_ERROR(checkDataIntegrity(*tmpRoutingInfo));
+      tmpRoutingInfo->routingInfoVersion =
+          std::max(tmpRoutingInfo->routingInfoVersion, flat::RoutingInfoVersion(rv + 1));
+    }
+
+    auto storeHandler = [&](kv::IReadWriteTransaction &txn) -> CoTryTask<void> {
+      CO_RETURN_ON_ERROR(co_await op_.state_.store_.storeRoutingInfo(txn, *tmpRoutingInfo));
+      co_return Void{};
+    };
+    CO_RETURN_ON_ERROR(co_await kv::WithTransaction(op_.state_.createRetryStrategy())
+                           .run(op_.state_.env_->kvEngine()->createReadWriteTransaction(), std::move(storeHandler)));
+
+    {
+      auto dataPtr = co_await op_.state_.data_.coLock();
+      auto &ri = dataPtr->routingInfo;
+
+      ri.nodeMap.clear();
+      ri.chainTables.clear();
+      ri.targets.clear();
+      ri.newBornChains.clear();
+
+      ++ri.routingInfoVersion.toUnderType();
+      for ([[maybe_unused]] const auto &[nodeId, node] : tmpRoutingInfo->nodes) {
+        ri.nodeMap.try_emplace(nodeId, node);
+      }
+      ri.chainTables = tmpRoutingInfo->chainTables;
+      ri.chains = tmpRoutingInfo->chains;
+      for ([[maybe_unused]] const auto &[tid, ti] : tmpRoutingInfo->targets) {
+        ri.targets[tid].base() = ti;
+      }
+    }
+    co_return Void{};
+  };
+
+  co_return co_await doAsPrimary(op_.state_, std::move(handler));
+}
+
+Result<Void> MgmtdTestHelper::checkDataIntegrity(const flat::RoutingInfo &ri) {
+  using PS = flat::PublicTargetState;
+  using LS = flat::LocalTargetState;
+  // check all nodes valid
+  {
+    for ([[maybe_unused]] const auto &[nodeId, node] : ri.nodes) {
+      if (nodeId == 0) {
+        return makeError(StatusCode::kDataCorruption, "NodeId is 0");
+      }
+    }
+  }
+  // check all chaintables
+  {
+    for (const auto &[chainTableId, chainTables] : ri.chainTables) {
+      if (chainTableId == 0) {
+        return makeError(StatusCode::kDataCorruption, "ChainTableId is 0");
+      }
+      if (chainTables.empty()) {
+        return makeError(StatusCode::kDataCorruption, fmt::format("Empty versions of {}", chainTableId));
+      }
+      for (const auto &[v, table] : chainTables) {
+        if (v == 0) {
+          return makeError(StatusCode::kDataCorruption, fmt::format("ChainTableVersion of {} is 0", chainTableId));
+        }
+        if (table.chains.empty()) {
+          return makeError(StatusCode::kDataCorruption, fmt::format("Empty chains of {} with {}", chainTableId, v));
+        }
+        robin_hood::unordered_set<flat::ChainId> s;
+        size_t targets = 0;
+        for (auto cid : table.chains) {
+          if (cid == 0) {
+            return makeError(StatusCode::kDataCorruption,
+                             fmt::format("ChainId(0) found in {} with {}", chainTableId, v));
+          }
+          if (!ri.chains.contains(cid)) {
+            return makeError(StatusCode::kDataCorruption,
+                             fmt::format("{} in {} with {} not existed", cid, chainTableId, v));
+          }
+          if (!s.insert(cid).second) {
+            return makeError(StatusCode::kDataCorruption,
+                             fmt::format("{} duplicated in {} with {}", cid, chainTableId, v));
+          }
+          const auto &chain = ri.chains.at(cid);
+          if (targets == 0) {
+            targets = chain.targets.size();
+          } else if (targets != chain.targets.size()) {
+            return makeError(StatusCode::kDataCorruption,
+                             fmt::format("Multiple replica count found in {} with {}", chainTableId, v));
+          }
+        }
+      }
+    }
+  }
+  // check all chains
+  for ([[maybe_unused]] const auto &[_, chain] : ri.chains) {
+    if (chain.chainId == 0) {
+      return makeError(StatusCode::kDataCorruption, "ChainId is 0");
+    }
+    if (chain.targets.empty()) {
+      return makeError(StatusCode::kDataCorruption, fmt::format("Empty targets of {}", chain.chainId));
+    }
+    auto lastPs = PS::INVALID;
+    for (const auto &cti : chain.targets) {
+      if (cti.targetId == 0) {
+        return makeError(StatusCode::kDataCorruption, fmt::format("TargetId of {} is 0", chain.chainId));
+      }
+      if (cti.publicState == flat::PublicTargetState::INVALID) {
+        return makeError(StatusCode::kDataCorruption,
+                         fmt::format("PublicTargetState of {} of {} is INVALID", cti.targetId, chain.chainId));
+      }
+      if (cti.publicState < lastPs) {
+        return makeError(
+            StatusCode::kDataCorruption,
+            fmt::format("PublicTargetState order inversion: {} after {}", toString(cti.publicState), toString(lastPs)));
+      }
+      lastPs = cti.publicState;
+    }
+  }
+
+  // check all targets valid
+  {
+    for (const auto &[tid, ti] : ri.targets) {
+      if (tid == 0) {
+        return makeError(StatusCode::kDataCorruption, "TargetId is 0");
+      }
+      if (ti.publicState == PS::INVALID) {
+        return makeError(StatusCode::kDataCorruption, fmt::format("PublicTargetState of {} is INVALID", tid));
+      }
+      if (ti.localState == LS::INVALID) {
+        return makeError(StatusCode::kDataCorruption, fmt::format("LocalTargetState of {} is INVALID", tid));
+      }
+      if (ti.chainId == 0) {
+        return makeError(StatusCode::kDataCorruption, fmt::format("ChainId of {} is 0", tid));
+      }
+
+      if (!ri.getChain(ti.chainId)) {
+        return makeError(StatusCode::kDataCorruption, fmt::format("ChainId {} of {} not found", ti.chainId, tid));
+      }
+
+      if (ti.nodeId && !ri.getNode(*ti.nodeId)) {
+        return makeError(StatusCode::kDataCorruption, fmt::format("NodeId {} of {} not found", *ti.nodeId, tid));
+      }
+    }
+  }
+  return Void{};
+}
+}  // namespace hf3fs::mgmtd::testing

tests/mgmtd/MgmtdTestHelper.h

@@ -0,0 +1,37 @@
+#pragma once
+
+#include "common/utils/StringUtils.h"
+#include "mgmtd/MgmtdServer.h"
+#include "mgmtd/background/MgmtdLeaseExtender.h"
+#include "mgmtd/service/MgmtdOperator.h"
+#include "mgmtd/service/helpers.h"
+
+namespace hf3fs::mgmtd::testing {
+
+class MgmtdTestHelper {
+ public:
+  MgmtdTestHelper(MgmtdOperator &op)
+      : op_(op) {}
+  MgmtdTestHelper(MgmtdServer &server)
+      : MgmtdTestHelper(*server.mgmtdOperator_) {}
+
+  MgmtdStore &getStore() { return op_.state_.store_; }
+  kv::IKVEngine &getEngine() { return *op_.state_.env_->kvEngine(); }
+
+  CoTask<NodeInfoWrapper *> getNodeInfo(flat::NodeId id) {
+    auto dataPtr = co_await op_.state_.data_.coLock();
+    auto &nm = dataPtr->routingInfo.nodeMap;
+    assert(nm.contains(id));
+    co_return &nm.at(id);
+  }
+
+  CoTask<void> extendLease() { co_await op_.backgroundRunner_.leaseExtender_->extend(); }
+
+  CoTryTask<void> setRoutingInfo(std::function<void(flat::RoutingInfo &)> callback);
+
+ private:
+  Result<Void> checkDataIntegrity(const flat::RoutingInfo &ri);
+
+  MgmtdOperator &op_;
+};
+}  // namespace hf3fs::mgmtd::testing

tests/mgmtd/TestMgmtdOperator.cc

@@ -0,0 +1,629 @@
+#include <folly/experimental/coro/BlockingWait.h>
+
+#include "MgmtdTestHelper.h"
+#include "common/kv/KeyPrefix.h"
+#include "common/kv/mem/MemKVEngine.h"
+#include "mgmtd/ops/Include.h"
+#include "mgmtd/service/MgmtdOperator.h"
+#include "tests/GtestHelpers.h"
+
+namespace hf3fs::mgmtd::testing {
+class MgmtdOperatorTest : public ::testing::Test {
+ protected:
+  MgmtdOperatorTest() {
+    flat::AppInfo info;
+    info.nodeId = flat::NodeId(10000);
+    info.clusterId = "clusterId";
+    info.hostname = "hostname0";
+    info.pid = 100;
+
+    flat::ServiceGroupInfo sgi0;
+    sgi0.services.insert("MgmtdService");
+    sgi0.endpoints.push_back(net::Address::fromString("127.0.0.1:9999"));
+    info.serviceGroups.push_back(std::move(sgi0));
+
+    defaultConfig_.set_lease_length(1_ms);
+    defaultConfig_.set_extend_lease_interval(0_s);
+    defaultConfig_.set_suspicious_lease_interval(100_us);
+    defaultConfig_.set_heartbeat_timestamp_valid_window(1_ms);
+    defaultConfig_.set_allow_heartbeat_from_unregistered(false);
+
+    env_ = std::make_shared<core::ServerEnv>();
+    env_->setAppInfo(info);
+    env_->setKvEngine(std::make_shared<kv::MemKVEngine>());
+    env_->setUtcTimeGenerator(generator_);
+    env_->setBackgroundExecutor(&executor_);
+  }
+
+  flat::HeartbeatInfo from(flat::NodeId id) {
+    flat::HeartbeatInfo info;
+    info.app.nodeId = id;
+    return info;
+  }
+
+  flat::HeartbeatInfo from(flat::NodeId id, flat::HeartbeatInfo info) {
+    info.app.nodeId = id;
+    return info;
+  }
+
+  CoTryTask<void> setPrimary(MgmtdOperator &mgmtd, flat::NodeId primaryId, UtcTime now) {
+    MgmtdStore store;
+    flat::PersistentNodeInfo info;
+    info.nodeId = primaryId;
+
+    MgmtdTestHelper helper(mgmtd);
+    now_ = now;
+    auto txn = helper.getEngine().createReadWriteTransaction();
+    auto result = co_await store.extendLease(*txn, info, std::chrono::microseconds(1000), now);
+    CO_RETURN_ON_ERROR(result);
+    co_return co_await txn->commit();
+  }
+
+  CoTryTask<mgmtd::GetPrimaryMgmtdRsp> getPrimaryMgmtd(MgmtdOperator &mgmtd, const String &clusterId) {
+    co_return co_await mgmtd.getPrimaryMgmtd(mgmtd::GetPrimaryMgmtdReq::create(clusterId), {});
+  }
+
+  CoTryTask<mgmtd::GetRoutingInfoRsp> getRoutingInfo(MgmtdOperator &mgmtd,
+                                                     const String &clusterId,
+                                                     flat::RoutingInfoVersion version) {
+    co_return co_await mgmtd.getRoutingInfo(mgmtd::GetRoutingInfoReq::create(clusterId, version), {});
+  }
+
+  CoTryTask<mgmtd::HeartbeatRsp> heartbeat(MgmtdOperator &mgmtd,
+                                           const String &clusterId,
+                                           const flat::HeartbeatInfo &hb,
+                                           UtcTime timestamp) {
+    co_return co_await mgmtd.heartbeat(mgmtd::HeartbeatReq::create(clusterId, hb, timestamp), {});
+  }
+
+  CoTryTask<mgmtd::EnableNodeRsp> enableNode(MgmtdOperator &mgmtd, const String &clusterId, flat::NodeId nodeId) {
+    co_return co_await mgmtd.enableNode(mgmtd::EnableNodeReq::create(clusterId, nodeId), {});
+  }
+
+  CoTryTask<mgmtd::DisableNodeRsp> disableNode(MgmtdOperator &mgmtd, const String &clusterId, flat::NodeId nodeId) {
+    co_return co_await mgmtd.disableNode(mgmtd::DisableNodeReq::create(clusterId, nodeId), {});
+  }
+
+  CoTryTask<mgmtd::RegisterNodeRsp> registerNode(MgmtdOperator &mgmtd,
+                                                 const String &clusterId,
+                                                 flat::NodeId nodeId,
+                                                 flat::NodeType type) {
+    co_return co_await mgmtd.registerNode(mgmtd::RegisterNodeReq::create(clusterId, nodeId, type), {});
+  }
+
+  CoTryTask<mgmtd::SetConfigRsp> setConfig(MgmtdOperator &mgmtd,
+                                           const String &clusterId,
+                                           flat::NodeType nodeType,
+                                           const String &content) {
+    co_return co_await mgmtd.setConfig(mgmtd::SetConfigReq::create(clusterId, nodeType, content), {});
+  }
+
+  CoTryTask<mgmtd::GetConfigRsp> getConfig(MgmtdOperator &mgmtd,
+                                           const String &clusterId,
+                                           flat::NodeType nodeType,
+                                           flat::ConfigVersion version) {
+    co_return co_await mgmtd.getConfig(mgmtd::GetConfigReq::create(clusterId, nodeType, version), {});
+  }
+
+  CoTryTask<mgmtd::SetChainTableRsp> setChainTable(MgmtdOperator &mgmtd,
+                                                   const String &clusterId,
+                                                   flat::ChainTableId tableId,
+                                                   const std::vector<flat::ChainId> &chains) {
+    co_return co_await mgmtd.setChainTable(mgmtd::SetChainTableReq::create(clusterId, tableId, chains), {});
+  }
+
+  CoTryTask<mgmtd::ExtendClientSessionRsp> extendClientSession(MgmtdOperator &mgmtd, auto &&...args) {
+    co_return co_await mgmtd.extendClientSession(
+        mgmtd::ExtendClientSessionReq::create(std::forward<decltype(args)>(args)...),
+        {});
+  }
+
+  CoTryTask<mgmtd::ListClientSessionsRsp> listClientSessions(MgmtdOperator &mgmtd, const String &clusterId) {
+    co_return co_await mgmtd.listClientSessions(mgmtd::ListClientSessionsReq::create(clusterId), {});
+  }
+
+  CoTryTask<mgmtd::SetNodeTagsRsp> setNodeTags(MgmtdOperator &mgmtd,
+                                               const String &clusterId,
+                                               flat::NodeId nodeId,
+                                               std::vector<flat::TagPair> tags,
+                                               flat::SetTagMode mode) {
+    co_return co_await mgmtd.setNodeTags(mgmtd::SetNodeTagsReq::create(clusterId, nodeId, std::move(tags), mode), {});
+  }
+
+  UtcTime now_ = UtcTime::fromMicroseconds(10000);
+  std::function<UtcTime()> generator_ = [this] { return now_; };
+  MgmtdConfig defaultConfig_;
+  std::shared_ptr<core::ServerEnv> env_;
+  CPUExecutorGroup executor_{4, "TEST"};
+};
+
+TEST_F(MgmtdOperatorTest, testGetPrimaryMgmtd) {
+  MgmtdOperator mgmtd(env_, defaultConfig_);
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    // wrong clusterid
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kClusterIdMismatch, getPrimaryMgmtd(mgmtd, "wrong"));
+
+    // no lease holder
+    {
+      auto ret = co_await getPrimaryMgmtd(mgmtd, "clusterId");
+      CO_ASSERT_OK(ret);
+      CO_ASSERT_TRUE(!ret->primary.has_value());
+    }
+
+    // node1 got lease
+    CO_AWAIT_ASSERT_OK(setPrimary(mgmtd, flat::NodeId(1), UtcTime::fromMicroseconds(10000)));
+
+    // node0 doesn't know node1's lease, load from store
+    {
+      auto ret = co_await getPrimaryMgmtd(mgmtd, "clusterId");
+      CO_ASSERT_OK(ret);
+      CO_ASSERT_TRUE(ret->primary.has_value());
+      CO_ASSERT_EQ(ret->primary->nodeId, flat::NodeId(1));
+    }
+
+    // node0 has known node1's lease, directly return
+    {
+      MgmtdTestHelper helper(mgmtd);
+      co_await helper.extendLease();
+      // hack: erase lease from engine to ensure the result come from cached lease
+      auto txn = helper.getEngine().createReadWriteTransaction();
+      auto key = fmt::format("{}MgmtdLease", kv::toStr(kv::KeyPrefix::Single));
+      CO_AWAIT_ASSERT_OK(txn->clear(key));
+      CO_AWAIT_ASSERT_OK(txn->commit());
+
+      auto ret = co_await getPrimaryMgmtd(mgmtd, "clusterId");
+      CO_ASSERT_OK(ret);
+      CO_ASSERT_TRUE(ret->primary.has_value());
+      CO_ASSERT_EQ(ret->primary->nodeId, flat::NodeId(1));
+    }
+
+    // node1's lease isn't trustable
+    {
+      // node2 got lease
+      CO_AWAIT_ASSERT_OK(setPrimary(mgmtd, flat::NodeId(2), UtcTime::fromMicroseconds(11001)));
+      now_ = UtcTime::fromMicroseconds(10950);
+
+      auto ret = co_await getPrimaryMgmtd(mgmtd, "clusterId");
+      CO_ASSERT_OK(ret);
+      CO_ASSERT_TRUE(ret->primary.has_value());
+      CO_ASSERT_EQ(ret->primary->nodeId, flat::NodeId(2));
+    }
+  }());
+}
+
+TEST_F(MgmtdOperatorTest, testRegisterNode) {
+  MgmtdOperator mgmtd(env_, defaultConfig_);
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    // wrong clusterid
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kClusterIdMismatch,
+                          registerNode(mgmtd, "wrong", flat::NodeId(1), flat::NodeType::MGMTD));
+
+    // not primary: no primary right now
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary,
+                          registerNode(mgmtd, "clusterId", flat::NodeId(1), flat::NodeType::MGMTD));
+
+    // not primary: not self
+    CO_AWAIT_ASSERT_OK(setPrimary(mgmtd, flat::NodeId(1), UtcTime::fromMicroseconds(10000)));
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary,
+                          registerNode(mgmtd, "clusterId", flat::NodeId(2), flat::NodeType::MGMTD));
+
+    // not primary: cached lease isn't trustable
+    {
+      now_ = UtcTime::fromMicroseconds(10950);
+      CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary,
+                            registerNode(mgmtd, "clusterId", flat::NodeId(2), flat::NodeType::MGMTD));
+    }
+
+    // succeed: self holds un-suspicious lease
+    {
+      now_ = UtcTime::fromMicroseconds(12000);
+      co_await MgmtdTestHelper(mgmtd).extendLease();
+
+      CO_AWAIT_ASSERT_OK(registerNode(mgmtd, "clusterId", flat::NodeId(1), flat::NodeType::MGMTD));
+
+      auto ret = co_await getRoutingInfo(mgmtd, "clusterId", flat::RoutingInfoVersion(0));
+      CO_ASSERT_OK(ret);
+      CO_ASSERT_TRUE(ret->info.has_value());
+      CO_ASSERT_TRUE(ret->info->nodes.contains(flat::NodeId(1)));
+      const auto &node = ret->info->nodes[flat::NodeId(1)];
+      CO_ASSERT_EQ(node.app.nodeId, flat::NodeId(1));
+      CO_ASSERT_EQ(node.type, flat::NodeType::MGMTD);
+      CO_ASSERT_EQ(node.status, flat::NodeStatus::HEARTBEAT_CONNECTING);
+    }
+
+    // duplicated
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kRegisterFail,
+                          registerNode(mgmtd, "clusterId", flat::NodeId(1), flat::NodeType::MGMTD));
+
+    // not primary: self lease is suspicious
+    {
+      now_ = UtcTime::fromMicroseconds(12950);
+      CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary,
+                            registerNode(mgmtd, "clusterId", flat::NodeId(2), flat::NodeType::MGMTD));
+    }
+  }());
+}
+
+TEST_F(MgmtdOperatorTest, testGetRoutingInfo) {
+  MgmtdOperator mgmtd(env_, defaultConfig_);
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    // wrong clusterid
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kClusterIdMismatch, getRoutingInfo(mgmtd, "wrong", flat::RoutingInfoVersion(0)));
+
+    // not primary: no primary right now
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary, getRoutingInfo(mgmtd, "clusterId", flat::RoutingInfoVersion(1)));
+
+    // init
+    co_await MgmtdTestHelper(mgmtd).extendLease();
+    {
+      auto ret = co_await getRoutingInfo(mgmtd, "clusterId", flat::RoutingInfoVersion(0));
+      CO_ASSERT_OK(ret);
+      CO_ASSERT_TRUE(ret->info.has_value());
+      CO_ASSERT_EQ(ret->info->routingInfoVersion, flat::RoutingInfoVersion(1));
+      CO_ASSERT_EQ(ret->info->nodes.size(), 1);
+      CO_ASSERT_TRUE(ret->info->nodes.contains(flat::NodeId(10000)));  // self
+      const auto &node = ret->info->nodes[flat::NodeId(10000)];
+      CO_ASSERT_EQ(node.app.nodeId, flat::NodeId(10000));
+      CO_ASSERT_EQ(node.type, flat::NodeType::MGMTD);
+      CO_ASSERT_EQ(node.status, flat::NodeStatus::PRIMARY_MGMTD);
+    }
+
+    // return nullopt when version is up-to-date
+    {
+      auto ret = co_await getRoutingInfo(mgmtd, "clusterId", flat::RoutingInfoVersion(1));
+      CO_ASSERT_OK(ret);
+      CO_ASSERT_TRUE(!ret->info.has_value());
+    }
+
+    // return error when version is newer than server's
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kInvalidRoutingInfoVersion,
+                          getRoutingInfo(mgmtd, "clusterId", flat::RoutingInfoVersion(2)));
+  }());
+}
+
+TEST_F(MgmtdOperatorTest, testHeartbeat) {
+  MgmtdOperator mgmtd(env_, defaultConfig_);
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    // wrong clusterid
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kClusterIdMismatch, heartbeat(mgmtd, "wrong", from(flat::NodeId{1}), now_));
+
+    flat::HeartbeatInfo info1;
+    {
+      info1.app.nodeId = flat::NodeId(1);
+      info1.app.hostname = "hostname1";
+      info1.app.pid = 101;
+
+      flat::ServiceGroupInfo sgi0;
+      sgi0.services.emplace("MgmtdService");
+      sgi0.endpoints.emplace_back(net::Address::fromString("127.0.0.1:9876"));
+
+      info1.app.serviceGroups.push_back(sgi0);
+      info1.set(flat::MgmtdHeartbeatInfo{});
+      info1.hbVersion = flat::HeartbeatVersion{1};
+    }
+    // not primary
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary, heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info1), now_));
+
+    co_await MgmtdTestHelper(mgmtd).extendLease();
+
+    // forbid self heartbeat
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kHeartbeatFail, heartbeat(mgmtd, "clusterId", from(flat::NodeId(10000)), now_));
+
+    // timestamp deviation
+    CO_AWAIT_ASSERT_ERROR(
+        MgmtdCode::kHeartbeatFail,
+        heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info1), now_ + std::chrono::microseconds(2500)));
+
+    // unregistered
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kHeartbeatFail, heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info1), now_));
+
+    CO_AWAIT_ASSERT_OK(registerNode(mgmtd, "clusterId", flat::NodeId(1), flat::NodeType::MGMTD));
+
+    // type mismatch with registered
+    {
+      auto info2 = info1;
+      info2.set(flat::MetaHeartbeatInfo{});
+      CO_AWAIT_ASSERT_ERROR(MgmtdCode::kHeartbeatFail,
+                            heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info2), now_));
+    }
+
+    // first successful heartbeat, status changed to CONNECTED
+    {
+      now_ = UtcTime::fromMicroseconds(10500);
+      CO_AWAIT_ASSERT_OK(heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info1), now_));
+      auto node = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+      CO_ASSERT_EQ(node.base().type, flat::NodeType::MGMTD);
+      CO_ASSERT_EQ(node.base().status, flat::NodeStatus::HEARTBEAT_CONNECTED);
+      CO_ASSERT_EQ(node.base().app, info1.app) << fmt::format("node.app = {} info1.app = {}",
+                                                              serde::toJsonString(node.base().app),
+                                                              serde::toJsonString(info1.app));
+      CO_ASSERT_EQ(node.base().lastHeartbeatTs, now_);
+    }
+
+    // extend heartbeat: nothing changed
+    {
+      info1.hbVersion = flat::HeartbeatVersion{3};
+      auto node1 = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+      now_ = UtcTime::fromMicroseconds(10501);
+      CO_AWAIT_ASSERT_OK(heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info1), now_));
+      auto node2 = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+      CO_ASSERT_EQ(node2.base().lastHeartbeatTs, now_);
+      // only lastHeartbeatTs differs
+      node2.base().lastHeartbeatTs = node1.base().lastHeartbeatTs;
+      CO_ASSERT_EQ(node1.base(), node2.base());
+    }
+
+    // stale heartbeat version
+    {
+      info1.hbVersion = flat::HeartbeatVersion{2};
+      now_ = UtcTime::fromMicroseconds(10502);
+      CO_AWAIT_ASSERT_ERROR(MgmtdCode::kHeartbeatVersionStale,
+                            heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info1), now_));
+    }
+
+    // service info changed without restarting
+    {
+      auto info2 = info1;
+      info2.hbVersion = flat::HeartbeatVersion{4};
+      info2.app.serviceGroups[0].services.emplace("OtherService");
+      CO_AWAIT_ASSERT_OK(heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info2), now_));
+
+      auto node = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+      CO_ASSERT_EQ(node.base().app, info2.app);
+    }
+
+    // restart: lastHeartbeatTs changed
+    {
+      auto node1 = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+
+      auto info2 = info1;
+      info2.app.hostname = "hostname0-2";
+      info2.hbVersion = flat::HeartbeatVersion{1};
+
+      now_ = UtcTime::fromMicroseconds(10502);
+      CO_AWAIT_ASSERT_ERROR(MgmtdCode::kHeartbeatVersionStale,
+                            heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info2), now_));
+
+      info2.hbVersion = flat::HeartbeatVersion{5};
+      CO_AWAIT_ASSERT_OK(heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info2), now_));
+
+      auto node2 = co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1));
+
+      CO_ASSERT_EQ(node2->base().type, node1.base().type);
+      CO_ASSERT_EQ(node2->base().status, node1.base().status);
+      CO_ASSERT_EQ(node2->base().app, info2.app);
+      CO_ASSERT_EQ(node2->base().lastHeartbeatTs, now_);
+
+      info2.app.pid = 102;
+      info2.hbVersion = flat::HeartbeatVersion{6};
+      now_ = UtcTime::fromMicroseconds(10503);
+      CO_AWAIT_ASSERT_OK(heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info2), now_));
+
+      node2 = co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1));
+
+      CO_ASSERT_EQ(node2->base().type, node1.base().type);
+      CO_ASSERT_EQ(node2->base().status, node1.base().status);
+      CO_ASSERT_EQ(node2->base().app, info2.app);
+      CO_ASSERT_EQ(node2->base().lastHeartbeatTs, now_);
+    }
+
+    // reconnected
+    {
+      auto &innerNode = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+      innerNode.base().status = flat::NodeStatus::HEARTBEAT_FAILED;
+
+      now_ = UtcTime::fromMicroseconds(10504);
+      info1.hbVersion = flat::HeartbeatVersion{7};
+      CO_AWAIT_ASSERT_OK(heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info1), now_));
+
+      auto node = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+      CO_ASSERT_EQ(node.base().status, flat::NodeStatus::HEARTBEAT_CONNECTED);
+      CO_ASSERT_EQ(node.base().lastHeartbeatTs, now_);
+    }
+
+    // reject heartbeats from offlined nodes
+    {
+      // TODO: migrate to a formal method
+      auto &innerNode = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+      innerNode.base().status = flat::NodeStatus::DISABLED;
+
+      CO_AWAIT_ASSERT_ERROR(MgmtdCode::kHeartbeatFail,
+                            heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), info1), now_));
+    }
+
+    // allow timestamp deviation and unregistered
+    {
+      auto options = defaultConfig_.clone();
+      options.set_heartbeat_timestamp_valid_window(0_ms);
+      options.set_allow_heartbeat_from_unregistered(true);
+      auto env = env_;
+      env->setKvEngine(std::make_shared<kv::MemKVEngine>());
+      MgmtdOperator mgmtd1(std::move(env), options);
+
+      now_ = UtcTime::fromMicroseconds(10000);
+      co_await MgmtdTestHelper(mgmtd1).extendLease();
+
+      CO_AWAIT_ASSERT_OK(
+          heartbeat(mgmtd1, "clusterId", from(flat::NodeId(1), info1), UtcTime::fromMicroseconds(12500)));
+      auto node = *(co_await MgmtdTestHelper(mgmtd1).getNodeInfo(flat::NodeId(1)));
+      CO_ASSERT_EQ(node.base().type, flat::NodeType::MGMTD);
+      CO_ASSERT_EQ(node.base().status, flat::NodeStatus::HEARTBEAT_CONNECTED);
+      CO_ASSERT_EQ(node.base().app, info1.app);
+      CO_ASSERT_EQ(node.base().lastHeartbeatTs, now_);
+    }
+  }());
+}
+
+TEST_F(MgmtdOperatorTest, testRecovery) {
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    auto mgmtd1 = std::make_unique<MgmtdOperator>(env_, defaultConfig_);
+
+    co_await MgmtdTestHelper(*mgmtd1).extendLease();
+
+    flat::HeartbeatInfo info1;
+    {
+      info1.app.hostname = "hostname1";
+      info1.app.pid = 101;
+
+      flat::ServiceGroupInfo sgi0;
+      sgi0.services.emplace("MgmtdService");
+      sgi0.endpoints.emplace_back(net::Address::fromString("127.0.0.1:9876"));
+
+      info1.app.serviceGroups.push_back(sgi0);
+      info1.set(flat::MgmtdHeartbeatInfo{});
+    }
+
+    auto info2 = info1;
+    info2.set(flat::MetaHeartbeatInfo{});
+
+    CO_AWAIT_ASSERT_OK(registerNode(*mgmtd1, "clusterId", flat::NodeId(1), flat::NodeType::MGMTD));
+    CO_AWAIT_ASSERT_OK(registerNode(*mgmtd1, "clusterId", flat::NodeId(2), flat::NodeType::META));
+    CO_AWAIT_ASSERT_OK(registerNode(*mgmtd1, "clusterId", flat::NodeId(3), flat::NodeType::STORAGE));
+
+    now_ += std::chrono::microseconds(10);
+    CO_AWAIT_ASSERT_OK(heartbeat(*mgmtd1, "clusterId", from(flat::NodeId(1), info1), now_));
+    CO_AWAIT_ASSERT_OK(heartbeat(*mgmtd1, "clusterId", from(flat::NodeId(2), info2), now_));
+
+    co_await disableNode(*mgmtd1, "clusterId", flat::NodeId(2));
+
+    mgmtd1.reset();
+    mgmtd1 = std::make_unique<MgmtdOperator>(env_, defaultConfig_);
+
+    co_await MgmtdTestHelper(*mgmtd1).extendLease();
+
+    auto node = co_await MgmtdTestHelper(*mgmtd1).getNodeInfo(flat::NodeId(1));
+    CO_ASSERT_EQ(node->base().status, flat::NodeStatus::HEARTBEAT_CONNECTING);
+
+    node = co_await MgmtdTestHelper(*mgmtd1).getNodeInfo(flat::NodeId(2));
+    CO_ASSERT_EQ(node->base().status, flat::NodeStatus::DISABLED);
+
+    node = co_await MgmtdTestHelper(*mgmtd1).getNodeInfo(flat::NodeId(3));
+    CO_ASSERT_EQ(node->base().status, flat::NodeStatus::HEARTBEAT_CONNECTING);
+  }());
+}
+
+TEST_F(MgmtdOperatorTest, testDisable) {
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    MgmtdOperator mgmtd(env_, defaultConfig_);
+    // wrong clusterid
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kClusterIdMismatch, disableNode(mgmtd, "wrong", flat::NodeId(1)));
+    // offline primary mgmtd
+    CO_AWAIT_ASSERT_ERROR(StatusCode::kInvalidArg, disableNode(mgmtd, "clusterId", flat::NodeId(10000)));
+    // not primary
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary, disableNode(mgmtd, "clusterId", flat::NodeId(1)));
+
+    co_await MgmtdTestHelper(mgmtd).extendLease();
+    // not registered
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNodeNotFound, disableNode(mgmtd, "clusterId", flat::NodeId(1)));
+
+    CO_AWAIT_ASSERT_OK(registerNode(mgmtd, "clusterId", flat::NodeId(1), flat::NodeType::MGMTD));
+    CO_AWAIT_ASSERT_OK(disableNode(mgmtd, "clusterId", flat::NodeId(1)));
+
+    auto info1 = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+    CO_ASSERT_EQ(info1.base().status, flat::NodeStatus::DISABLED);
+
+    // duplicated offline has no effect
+    CO_AWAIT_ASSERT_OK(disableNode(mgmtd, "clusterId", flat::NodeId(1)));
+    auto info2 = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+    CO_ASSERT_EQ(info1.base(), info2.base());
+  }());
+}
+
+TEST_F(MgmtdOperatorTest, testEnable) {
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    MgmtdOperator mgmtd(env_, defaultConfig_);
+    // wrong clusterid
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kClusterIdMismatch, enableNode(mgmtd, "wrong", flat::NodeId(1)));
+    // offline primary mgmtd
+    CO_AWAIT_ASSERT_ERROR(StatusCode::kInvalidArg, enableNode(mgmtd, "clusterId", flat::NodeId(10000)));
+    // not primaryA
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary, enableNode(mgmtd, "clusterId", flat::NodeId(1)));
+
+    co_await MgmtdTestHelper(mgmtd).extendLease();
+    // not registered
+    CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNodeNotFound, enableNode(mgmtd, "clusterId", flat::NodeId(1)));
+
+    CO_AWAIT_ASSERT_OK(registerNode(mgmtd, "clusterId", flat::NodeId(1), flat::NodeType::MGMTD));
+
+    flat::HeartbeatInfo hb;
+    hb.set(flat::MgmtdHeartbeatInfo{});
+    CO_AWAIT_ASSERT_OK(heartbeat(mgmtd, "clusterId", from(flat::NodeId(1), hb), now_));
+
+    auto info1 = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+    CO_ASSERT_EQ(info1.base().status, flat::NodeStatus::HEARTBEAT_CONNECTED);
+
+    // online an online node has no effect
+    CO_AWAIT_ASSERT_OK(enableNode(mgmtd, "clusterId", flat::NodeId(1)));
+    auto info2 = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+    CO_ASSERT_EQ(info1.base(), info2.base());
+
+    CO_AWAIT_ASSERT_OK(disableNode(mgmtd, "clusterId", flat::NodeId(1)));
+    CO_AWAIT_ASSERT_OK(enableNode(mgmtd, "clusterId", flat::NodeId(1)));
+
+    auto info3 = *(co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)));
+    CO_ASSERT_EQ(info3.base().status, flat::NodeStatus::HEARTBEAT_CONNECTING);
+  }());
+}
+
+TEST_F(MgmtdOperatorTest, testSetNodeTags) {
+  auto f = [&]() -> CoTask<void> {
+    MgmtdOperator mgmtd(env_, defaultConfig_);
+    co_await MgmtdTestHelper(mgmtd).extendLease();
+
+    CO_AWAIT_ASSERT_OK(registerNode(mgmtd, "clusterId", flat::NodeId(1), flat::NodeType::MGMTD));
+
+    auto info1 = (co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)))->base();
+    CO_ASSERT_TRUE(info1.tags.empty());
+
+    auto setRes = co_await setNodeTags(mgmtd,
+                                       "clusterId",
+                                       flat::NodeId(1),
+                                       {flat::TagPair("A", "1"), flat::TagPair("B", "2"), flat::TagPair("C")},
+                                       flat::SetTagMode::UPSERT);
+    CO_ASSERT_OK(setRes);
+
+    auto expected = std::vector{flat::TagPair("A", "1"), flat::TagPair("B", "2"), flat::TagPair("C")};
+    info1 = (co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)))->base();
+    CO_ASSERT_EQ(setRes->info, info1);
+    CO_ASSERT_EQ(info1.tags, expected);
+
+    setRes = co_await setNodeTags(mgmtd,
+                                  "clusterId",
+                                  flat::NodeId(1),
+                                  {flat::TagPair("B", "3"), flat::TagPair("D", "4")},
+                                  flat::SetTagMode::UPSERT);
+    CO_ASSERT_OK(setRes);
+
+    expected =
+        std::vector{flat::TagPair("A", "1"), flat::TagPair("B", "3"), flat::TagPair("C"), flat::TagPair("D", "4")};
+    info1 = (co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)))->base();
+    CO_ASSERT_EQ(setRes->info, info1);
+    CO_ASSERT_EQ(info1.tags, expected) << fmt::format("actual: \n{}\nexpectd: {}",
+                                                      serde::toJsonString(info1.tags),
+                                                      serde::toJsonString(expected));
+
+    setRes = co_await setNodeTags(mgmtd,
+                                  "clusterId",
+                                  flat::NodeId(1),
+                                  {flat::TagPair("A", "2"), flat::TagPair("C", "3")},
+                                  flat::SetTagMode::REPLACE);
+    CO_ASSERT_OK(setRes);
+
+    expected = std::vector{flat::TagPair("A", "2"), flat::TagPair("C", "3")};
+    info1 = (co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)))->base();
+    CO_ASSERT_EQ(setRes->info, info1);
+    CO_ASSERT_EQ(info1.tags, expected);
+
+    setRes =
+        co_await setNodeTags(mgmtd, "clusterId", flat::NodeId(1), {flat::TagPair("A", "1")}, flat::SetTagMode::REMOVE);
+    CO_ASSERT_ERROR(setRes, MgmtdCode::kInvalidTag);
+
+    setRes = co_await setNodeTags(mgmtd, "clusterId", flat::NodeId(1), {flat::TagPair("A")}, flat::SetTagMode::REMOVE);
+
+    expected = std::vector{flat::TagPair("C", "3")};
+    info1 = (co_await MgmtdTestHelper(mgmtd).getNodeInfo(flat::NodeId(1)))->base();
+    CO_ASSERT_EQ(setRes->info, info1);
+  };
+  folly::coro::blockingWait(f());
+}
+}  // namespace hf3fs::mgmtd::testing

tests/mgmtd/TestMgmtdStore.cc

@@ -0,0 +1,224 @@
+#include <folly/experimental/coro/BlockingWait.h>
+
+#include "common/kv/mem/MemKVEngine.h"
+#include "mgmtd/store/MgmtdStore.h"
+#include "tests/GtestHelpers.h"
+
+namespace hf3fs::testing {
+namespace {
+using namespace ::hf3fs::mgmtd;
+
+class MgmtdStoreTest : public ::testing::Test {
+ protected:
+  MgmtdStoreTest() {
+    node0.nodeId = flat::NodeId(0);
+    node1.nodeId = flat::NodeId(1);
+  }
+
+  flat::PersistentNodeInfo from(flat::NodeId id, flat::PersistentNodeInfo info) {
+    info.nodeId = id;
+    return info;
+  }
+
+  kv::MemKVEngine engine;
+  UtcTime now;
+  std::function<UtcTime()> generator = [this] { return now; };
+  flat::PersistentNodeInfo node0;
+  flat::PersistentNodeInfo node1;
+  MgmtdStore store;
+  static constexpr auto leaseLength = std::chrono::microseconds(1000);
+};
+
+TEST_F(MgmtdStoreTest, testExtendLease) {
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    {
+      // nobody holds the lease
+      now = UtcTime::fromMicroseconds(10000);
+      auto txn = engine.createReadWriteTransaction();
+      auto result = co_await store.extendLease(*txn, node0, leaseLength, now);
+      CO_ASSERT_OK(result);
+      CO_ASSERT_EQ(result->primary.nodeId, 0);
+      CO_ASSERT_EQ(result->leaseStart.toMicroseconds(), 10000);
+      CO_ASSERT_EQ(result->leaseEnd.toMicroseconds(), 11000);
+      CO_AWAIT_ASSERT_OK(txn->commit());
+    }
+    {
+      // lease won't go back
+      now = UtcTime::fromMicroseconds(9900);
+      auto txn = engine.createReadWriteTransaction();
+      auto result = co_await store.extendLease(*txn, node1, leaseLength, now);
+      CO_ASSERT_OK(result);
+      CO_ASSERT_EQ(result->primary.nodeId, node0.nodeId);
+      CO_ASSERT_EQ(result->leaseStart.toMicroseconds(), 10000);
+      CO_ASSERT_EQ(result->leaseEnd.toMicroseconds(), 11000);
+      CO_AWAIT_ASSERT_OK(txn->commit());
+    }
+    {
+      // node1 won't become the lease owner before lease is end
+      now = UtcTime::fromMicroseconds(10900);
+      auto txn = engine.createReadWriteTransaction();
+      auto result = co_await store.extendLease(*txn, node1, leaseLength, now);
+      CO_ASSERT_OK(result);
+      CO_ASSERT_EQ(result->primary.nodeId, node0.nodeId);
+      CO_ASSERT_EQ(result->leaseStart.toMicroseconds(), 10000);
+      CO_ASSERT_EQ(result->leaseEnd.toMicroseconds(), 11000);
+      CO_AWAIT_ASSERT_OK(txn->commit());
+    }
+    {
+      // node1 becomes the lease owner after lease is end
+      now = UtcTime::fromMicroseconds(11100);
+      auto txn = engine.createReadWriteTransaction();
+      auto result = co_await store.extendLease(*txn, node1, leaseLength, now);
+      CO_ASSERT_OK(result);
+      CO_ASSERT_EQ(result->primary.nodeId, node1.nodeId);
+      CO_ASSERT_EQ(result->leaseStart.toMicroseconds(), 11100);
+      CO_ASSERT_EQ(result->leaseEnd.toMicroseconds(), 12100);
+      CO_AWAIT_ASSERT_OK(txn->commit());
+    }
+    {
+      // node1 extends lease
+      now = UtcTime::fromMicroseconds(11500);
+      auto txn = engine.createReadWriteTransaction();
+      auto result = co_await store.extendLease(*txn, node1, leaseLength, now);
+      CO_ASSERT_OK(result);
+      CO_ASSERT_EQ(result->primary.nodeId, node1.nodeId);
+      CO_ASSERT_EQ(result->leaseStart.toMicroseconds(), 11100);
+      CO_ASSERT_EQ(result->leaseEnd.toMicroseconds(), 12500);
+      CO_AWAIT_ASSERT_OK(txn->commit());
+    }
+    {
+      // node1 starts new lease after the previous lease was finished
+      now = UtcTime::fromMicroseconds(13000);
+      auto txn = engine.createReadWriteTransaction();
+      auto result = co_await store.extendLease(*txn, node1, leaseLength, now);
+      CO_ASSERT_OK(result);
+      CO_ASSERT_EQ(result->primary.nodeId, node1.nodeId);
+      CO_ASSERT_EQ(result->leaseStart.toMicroseconds(), 13000);
+      CO_ASSERT_EQ(result->leaseEnd.toMicroseconds(), 14000);
+      CO_AWAIT_ASSERT_OK(txn->commit());
+    }
+    {
+      // node1 cannot extend lease with a stale release version
+      auto rv = flat::ReleaseVersion::fromVersionInfo();
+      --rv.buildTimeInSeconds;
+      now = UtcTime::fromMicroseconds(13001);
+      auto txn = engine.createReadWriteTransaction();
+      auto result = co_await store.extendLease(*txn, node1, leaseLength, now, rv, true);
+      CO_ASSERT_OK(result);
+      CO_ASSERT_EQ(result->primary.nodeId, node1.nodeId);
+      CO_ASSERT_EQ(result->leaseStart.toMicroseconds(), 13000);
+      CO_ASSERT_EQ(result->leaseEnd.toMicroseconds(), 14000);
+      CO_AWAIT_ASSERT_OK(txn->commit());
+    }
+    co_return;
+  }());
+}
+
+TEST_F(MgmtdStoreTest, testEnsureLeaseValid) {
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    {
+      // let node0 get the lease
+      now = UtcTime::fromMicroseconds(10000);
+      auto txn = engine.createReadWriteTransaction();
+      auto result = co_await store.extendLease(*txn, node0, leaseLength, now);
+      CO_ASSERT_OK(result);
+      CO_AWAIT_ASSERT_OK(txn->commit());
+    }
+    {
+      // node1 got "not primary"
+      auto txn = engine.createReadonlyTransaction();
+      CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary, store.ensureLeaseValid(*txn, node1.nodeId, now));
+    }
+    {
+      // node0 ensures self holding the lease
+      now = UtcTime::fromMicroseconds(10500);
+      auto txn = engine.createReadonlyTransaction();
+      CO_AWAIT_ASSERT_OK(store.ensureLeaseValid(*txn, node0.nodeId, now));
+    }
+    {
+      // node0 got "not primary" after lease is expired
+      now = UtcTime::fromMicroseconds(11500);
+      auto txn = engine.createReadonlyTransaction();
+      CO_AWAIT_ASSERT_ERROR(MgmtdCode::kNotPrimary, store.ensureLeaseValid(*txn, node0.nodeId, now));
+    }
+    co_return;
+  }());
+}
+
+TEST_F(MgmtdStoreTest, testLoadStoreNode) {
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    {
+      // node0 got lease, then store will succeed
+      flat::PersistentNodeInfo ni = {
+          .type = flat::NodeType::META,
+          .serviceGroups = {flat::ServiceGroupInfo({"meta0", "meta1"}, {net::Address::fromString("127.0.0.1:9876")})},
+      };
+
+      now = UtcTime::fromMicroseconds(10000);
+      auto txn = engine.createReadWriteTransaction();
+      CO_AWAIT_ASSERT_OK(store.extendLease(*txn, node0, leaseLength, now));
+      CO_AWAIT_ASSERT_OK(txn->commit());
+
+      txn = engine.createReadWriteTransaction();
+      CO_AWAIT_ASSERT_OK(store.storeNodeInfo(*txn, from(flat::NodeId(2), ni)));
+      CO_AWAIT_ASSERT_OK(store.storeNodeInfo(*txn, from(flat::NodeId(3), ni)));
+      CO_AWAIT_ASSERT_OK(store.storeNodeInfo(*txn, from(flat::NodeId(4), ni)));
+      CO_AWAIT_ASSERT_OK(txn->commit());
+
+      // store now could load node info without lease ownership
+      auto txn1 = engine.createReadonlyTransaction();
+      auto result = co_await store.loadNodeInfo(*txn1, flat::NodeId(2));
+      CO_ASSERT_OK(result);
+      CO_ASSERT_TRUE(result->has_value());
+      CO_ASSERT_EQ(from(flat::NodeId(2), ni), result->value());
+
+      kv::mem::memKvDefaultLimit = 1;
+      auto result1 = co_await store.loadAllNodes(*txn1);
+      CO_ASSERT_OK(result1);
+      CO_ASSERT_EQ(result1->size(), 3);
+      CO_ASSERT_EQ(from(flat::NodeId(2), ni), result1->at(0));
+      CO_ASSERT_EQ(from(flat::NodeId(3), ni), result1->at(1));
+      CO_ASSERT_EQ(from(flat::NodeId(4), ni), result1->at(2));
+    }
+    co_return;
+  }());
+}
+
+TEST_F(MgmtdStoreTest, testLoadStoreConfig) {
+  folly::coro::blockingWait([&]() -> CoTask<void> {
+    auto txn = engine.createReadWriteTransaction();
+    auto loadAllRes = co_await store.loadAllConfigs(*txn);
+    CO_ASSERT_OK(loadAllRes);
+    CO_ASSERT_TRUE(loadAllRes->empty());
+
+    auto cfg = flat::ConfigInfo::create(flat::ConfigVersion(5), "abc");
+    CO_AWAIT_ASSERT_OK(
+        store.storeConfig(*txn, flat::NodeType::MGMTD, flat::ConfigInfo::create(flat::ConfigVersion(5), "abc")));
+    CO_AWAIT_ASSERT_OK(
+        store.storeConfig(*txn, flat::NodeType::MGMTD, flat::ConfigInfo::create(flat::ConfigVersion(6), "abcd")));
+
+    auto loadRes = co_await store.loadConfig(*txn, flat::NodeType::MGMTD, flat::ConfigVersion(4));
+    CO_ASSERT_OK(loadRes);
+    CO_ASSERT_TRUE(!loadRes->has_value());
+
+    loadRes = co_await store.loadConfig(*txn, flat::NodeType::MGMTD, flat::ConfigVersion(5));
+    CO_ASSERT_OK(loadRes);
+    CO_ASSERT_TRUE(loadRes->has_value());
+    CO_ASSERT_EQ(loadRes->value().configVersion, 5);
+    CO_ASSERT_EQ(loadRes->value().content, "abc");
+
+    loadRes = co_await store.loadConfig(*txn, flat::NodeType::MGMTD, flat::ConfigVersion(6));
+    CO_ASSERT_OK(loadRes);
+    CO_ASSERT_TRUE(loadRes->has_value());
+    CO_ASSERT_EQ(loadRes->value().configVersion, 6);
+    CO_ASSERT_EQ(loadRes->value().content, "abcd");
+
+    loadRes = co_await store.loadLatestConfig(*txn, flat::NodeType::MGMTD);
+    CO_ASSERT_OK(loadRes);
+    CO_ASSERT_TRUE(loadRes->has_value());
+    CO_ASSERT_EQ(loadRes->value().configVersion, 6);
+    CO_ASSERT_EQ(loadRes->value().content, "abcd");
+  }());
+}
+}  // namespace
+}  // namespace hf3fs::testing

tests/mgmtd/TestNodeSelector.cc

@@ -0,0 +1,61 @@
+#include "fbs/mgmtd/NodeInfo.h"
+#include "tests/GtestHelpers.h"
+
+namespace hf3fs::test {
+namespace {
+TEST(NodeSelector, testSelectByTrifficZone) {
+  auto createNode = [](const String &zone) {
+    flat::NodeInfo n;
+    if (!zone.empty()) {
+      n.tags = {flat::TagPair(flat::kTrafficZoneTagKey, zone)};
+    }
+    return n;
+  };
+  auto n0 = createNode("");
+  auto n1 = createNode("zone1");
+  auto n2 = createNode("zone2");
+
+  {
+    auto sel = flat::selectNodeByTrafficZone("");
+    ASSERT_TRUE(sel(n0));
+    ASSERT_TRUE(sel(n1));
+    ASSERT_TRUE(sel(n2));
+  }
+
+  {
+    auto sel = flat::selectNodeByTrafficZone("zone1");
+    ASSERT_TRUE(!sel(n0));
+    ASSERT_TRUE(sel(n1));
+    ASSERT_TRUE(!sel(n2));
+  }
+
+  {
+    auto sel = flat::selectNodeByTrafficZone(" zone1");
+    ASSERT_TRUE(!sel(n0));
+    ASSERT_TRUE(sel(n1));
+    ASSERT_TRUE(!sel(n2));
+  }
+
+  {
+    auto sel = flat::selectNodeByTrafficZone(" zone1 ");
+    ASSERT_TRUE(!sel(n0));
+    ASSERT_TRUE(sel(n1));
+    ASSERT_TRUE(!sel(n2));
+  }
+
+  {
+    auto sel = flat::selectNodeByTrafficZone(" zone1 ,");
+    ASSERT_TRUE(!sel(n0));
+    ASSERT_TRUE(sel(n1));
+    ASSERT_TRUE(!sel(n2));
+  }
+
+  {
+    auto sel = flat::selectNodeByTrafficZone(" zone1 , zone2");
+    ASSERT_TRUE(!sel(n0));
+    ASSERT_TRUE(sel(n1));
+    ASSERT_TRUE(sel(n2));
+  }
+}
+}  // namespace
+}  // namespace hf3fs::test