From f1ad49e1e82ddb73665a8dd67f6d12848725109b Mon Sep 17 00:00:00 2001
From: Doodle <13706157+critical27@users.noreply.github.com>
Date: Wed, 19 Oct 2022 14:33:39 +0800
Subject: [PATCH] stats handle the flag use_vertex_key (#4738)

---
 src/storage/admin/StatsTask.cpp | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/storage/admin/StatsTask.cpp b/src/storage/admin/StatsTask.cpp
index 38cf4564306..1a0af245fd8 100644
--- a/src/storage/admin/StatsTask.cpp
+++ b/src/storage/admin/StatsTask.cpp
@@ -10,6 +10,7 @@
 #include "common/base/MurmurHash2.h"
 #include "common/utils/NebulaKeyUtils.h"
 #include "kvstore/Common.h"
+#include "storage/StorageFlags.h"
 
 DEFINE_int32(stats_sleep_interval_ms,
              0,
@@ -111,6 +112,8 @@ nebula::cpp2::ErrorCode StatsTask::genSubTask(GraphSpaceID spaceId,
   std::unique_ptr<kvstore::KVIterator> tagIter;
   auto edgePrefix = NebulaKeyUtils::edgePrefix(part);
   std::unique_ptr<kvstore::KVIterator> edgeIter;
+  auto vertexPrefix = NebulaKeyUtils::vertexPrefix(part);
+  std::unique_ptr<kvstore::KVIterator> vertexIter;
 
   // When the storage occurs leader change, continue to read data from the
   // follower instead of reporting an error.
@@ -124,6 +127,13 @@ nebula::cpp2::ErrorCode StatsTask::genSubTask(GraphSpaceID spaceId,
     LOG(INFO) << "Stats task failed";
     return ret;
   }
+  if (FLAGS_use_vertex_key) {
+    ret = env_->kvstore_->prefix(spaceId, part, vertexPrefix, &vertexIter, true);
+    if (ret != nebula::cpp2::ErrorCode::SUCCEEDED) {
+      LOG(INFO) << "Stats task failed";
+      return ret;
+    }
+  }
   std::unordered_map<TagID, int64_t> tagsVertices;
   std::unordered_map<EdgeType, int64_t> edgetypeEdges;
   std::unordered_map<PartitionID, int64_t> positiveRelevancy;
@@ -230,6 +240,14 @@ nebula::cpp2::ErrorCode StatsTask::genSubTask(GraphSpaceID spaceId,
     edgeIter->next();
     sleepIfScannedSomeRecord(++countToSleep);
   }
+  int64_t verticesCountByVertexKey = 0;
+  if (FLAGS_use_vertex_key) {
+    while (vertexIter && vertexIter->valid()) {
+      verticesCountByVertexKey++;
+      vertexIter->next();
+      sleepIfScannedSomeRecord(++countToSleep);
+    }
+  }
   nebula::meta::cpp2::StatsItem statsItem;
 
   // convert tagId/edgeType to tagName/edgeName
@@ -246,7 +264,7 @@ nebula::cpp2::ErrorCode StatsTask::genSubTask(GraphSpaceID spaceId,
     }
   }
 
-  statsItem.space_vertices_ref() = spaceVertices;
+  statsItem.space_vertices_ref() = FLAGS_use_vertex_key ? verticesCountByVertexKey : spaceVertices;
   statsItem.space_edges_ref() = spaceEdges;
   using Correlativities = std::vector<nebula::meta::cpp2::Correlativity>;
   Correlativities positiveCorrelativity;