parse diskPartsKey and add unit test

src/common/utils/MetaKeyUtils.cpp

@@ -1174,18 +1174,43 @@ GraphSpaceID MetaKeyUtils::parseLocalIdSpace(folly::StringPiece rawData) {
   return *reinterpret_cast<const GraphSpaceID*>(rawData.data() + offset);
 }
 
-GraphSpaceID MetaKeyUtils::parseDiskPartsSpace(folly::StringPiece rawData) {
+/**
+ * diskPartsKey = kDiskPartsTable + len(serialized(hostAddr)) + serialized(hostAddr) + path
+ */
+
+HostAddr MetaKeyUtils::parseDiskPartsHost(const folly::StringPiece& rawData) {
   auto offset = kDiskPartsTable.size();
-  return *reinterpret_cast<const GraphSpaceID*>(rawData.data() + offset);
+  auto hostAddrLen = *reinterpret_cast<const size_t*>(rawData.begin() + offset);
+  offset += sizeof(size_t);
+  std::string hostAddrStr(rawData.data() + offset, hostAddrLen);
+  return deserializeHostAddr(hostAddrStr);
+}
+
+GraphSpaceID MetaKeyUtils::parseDiskPartsSpace(const folly::StringPiece& rawData) {
+  auto offset = kDiskPartsTable.size();
+  size_t hostAddrLen = *reinterpret_cast<const size_t*>(rawData.begin() + offset);
+  offset += sizeof(size_t) + hostAddrLen;
+  return *reinterpret_cast<const GraphSpaceID*>(rawData.begin() + offset);
+}
+
+std::string MetaKeyUtils::parseDiskPartsPath(const folly::StringPiece& rawData) {
+  auto offset = kDiskPartsTable.size();
+  size_t hostAddrLen = *reinterpret_cast<const size_t*>(rawData.begin() + offset);
+  offset += sizeof(size_t) + hostAddrLen + sizeof(GraphSpaceID);
+  std::string path(rawData.begin() + offset, rawData.size() - offset);
+  return path;
 }
 
 std::string MetaKeyUtils::diskPartsPrefix() { return kDiskPartsTable; }
 
 std::string MetaKeyUtils::diskPartsPrefix(HostAddr addr) {
   std::string key;
   std::string hostStr = serializeHostAddr(addr);
-  key.reserve(kDiskPartsTable.size() + hostStr.size());
-  key.append(kDiskPartsTable.data(), kDiskPartsTable.size()).append(hostStr.data(), hostStr.size());
+  size_t hostAddrLen = hostStr.size();
+  key.reserve(kDiskPartsTable.size() + sizeof(size_t) + hostStr.size());
+  key.append(kDiskPartsTable.data(), kDiskPartsTable.size())
+      .append(reinterpret_cast<const char*>(&hostAddrLen), sizeof(size_t))
+      .append(hostStr.data(), hostStr.size());
   return key;
 }
 
@@ -1198,7 +1223,9 @@ std::string MetaKeyUtils::diskPartsPrefix(HostAddr addr, GraphSpaceID spaceId) {
   return key;
 }
 
-std::string MetaKeyUtils::diskPartsKey(HostAddr addr, GraphSpaceID spaceId, std::string path) {
+std::string MetaKeyUtils::diskPartsKey(HostAddr addr,
+                                       GraphSpaceID spaceId,
+                                       const std::string& path) {
   std::string key;
   std::string prefix = diskPartsPrefix(addr, spaceId);
   key.reserve(prefix.size() + path.size());

src/common/utils/MetaKeyUtils.h

@@ -382,15 +382,19 @@ class MetaKeyUtils final {
 
   static std::unordered_map<std::string, std::pair<std::string, bool>> getSystemTableMaps();
 
-  static GraphSpaceID parseDiskPartsSpace(folly::StringPiece rawData);
+  static GraphSpaceID parseDiskPartsSpace(const folly::StringPiece& rawData);
+
+  static HostAddr parseDiskPartsHost(const folly::StringPiece& rawData);
+
+  static std::string parseDiskPartsPath(const folly::StringPiece& rawData);
 
   static std::string diskPartsPrefix();
 
   static std::string diskPartsPrefix(HostAddr addr);
 
   static std::string diskPartsPrefix(HostAddr addr, GraphSpaceID spaceId);
 
-  static std::string diskPartsKey(HostAddr addr, GraphSpaceID spaceId, std::string path);
+  static std::string diskPartsKey(HostAddr addr, GraphSpaceID spaceId, const std::string& path);
 
   static std::string diskPartsVal(const meta::cpp2::PartitionList& partList);
 

src/common/utils/test/MetaKeyUtilsTest.cpp

@@ -201,6 +201,17 @@ TEST(MetaKeyUtilsTest, ZoneTest) {
   ASSERT_EQ(nodes, MetaKeyUtils::parseZoneHosts(zoneValue));
 }
 
+TEST(MetaKeyUtilsTest, DiskPathsTest) {
+  HostAddr addr{"192.168.0.1", 1234};
+  GraphSpaceID spaceId = 1;
+  std::string path = "/data/storage/test_part1";
+
+  auto diskPartsKey = MetaKeyUtils::diskPartsKey(addr, spaceId, path);
+  ASSERT_EQ(addr, MetaKeyUtils::parseDiskPartsHost(diskPartsKey));
+  ASSERT_EQ(spaceId, MetaKeyUtils::parseDiskPartsSpace(diskPartsKey));
+  ASSERT_EQ(path, MetaKeyUtils::parseDiskPartsPath(diskPartsKey));
+}
+
 }  // namespace nebula
 
 int main(int argc, char** argv) {

src/kvstore/test/DiskManagerTest.cpp

@@ -145,6 +145,39 @@ TEST(DiskManagerTest, WalNoSpaceTest) {
   }
 }
 
+TEST(DiskManagerTest, GetDiskPartsTest) {
+  GraphSpaceID spaceId = 1;
+  fs::TempDir disk1("/tmp/get_disk_part_test.XXXXXX");
+  auto path1 = folly::stringPrintf("%s/nebula/%d", disk1.path(), spaceId);
+  boost::filesystem::create_directories(path1);
+  fs::TempDir disk2("/tmp/get_disk_part_test.XXXXXX");
+  auto path2 = folly::stringPrintf("%s/nebula/%d", disk2.path(), spaceId);
+  boost::filesystem::create_directories(path2);
+  GraphSpaceID spaceId2 = 2;
+  fs::TempDir disk3("/tmp/get_disk_part_test.XXXXXX");
+  auto path3 = folly::stringPrintf("%s/nebula/%d", disk3.path(), spaceId2);
+  boost::filesystem::create_directories(path3);
+
+  std::vector<std::string> dataPaths = {disk1.path(), disk2.path(), disk3.path()};
+  DiskManager diskMan(dataPaths);
+  for (PartitionID partId = 1; partId <= 10; partId++) {
+    diskMan.addPartToPath(spaceId, partId, path1);
+  }
+  for (PartitionID partId = 11; partId <= 20; partId++) {
+    diskMan.addPartToPath(spaceId, partId, path2);
+  }
+  for (PartitionID partId = 1; partId <= 10; partId++) {
+    diskMan.addPartToPath(spaceId2, partId, path3);
+  }
+
+  SpaceDiskPartsMap diskParts;
+  diskMan.getDiskParts(diskParts);
+  ASSERT_EQ(2, diskParts.size());
+  ASSERT_EQ(2, diskParts[spaceId].size());
+  ASSERT_EQ(1, diskParts[spaceId2].size());
+  ASSERT_EQ(10, diskParts[spaceId2][path3].get_part_list().size());
+}
+
 }  // namespace kvstore
 }  // namespace nebula